From f96191f33d8738694b03990b582aeb89698d814b Mon Sep 17 00:00:00 2001
From: Kyle Lahnakoski <kyle@lahnakoski.com>
Date: Wed, 30 Oct 2024 21:29:25 -0400
Subject: [PATCH] include nulls

---
 mo_sql_parsing/formatting.py   |  6 +++++-
 tests/test_format_and_parse.py | 21 ++++++++++++++++++---
 2 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/mo_sql_parsing/formatting.py b/mo_sql_parsing/formatting.py
index 7ae6e46..3f2aab1 100644
--- a/mo_sql_parsing/formatting.py
+++ b/mo_sql_parsing/formatting.py
@@ -519,7 +519,11 @@ def pivot(self, json, prec):
 
     def unpivot(self, json, prec):
         pivot = json["unpivot"]
-        return self._pivot("UNPIVOT", pivot, self.dispatch(pivot["value"]))
+        if "nulls" in pivot:
+            nulls = " INCLUDE NULLS" if pivot["nulls"] else " EXCLUDE NULLS"
+        else:
+            nulls = ""
+        return self._pivot(f"UNPIVOT{nulls}", pivot, self.dispatch(pivot["value"]))
 
     def _pivot(self, op, pivot, value):
         for_ = self.dispatch(pivot["for"])
diff --git a/tests/test_format_and_parse.py b/tests/test_format_and_parse.py
index 3b5b256..b36ad3b 100644
--- a/tests/test_format_and_parse.py
+++ b/tests/test_format_and_parse.py
@@ -11,7 +11,7 @@
 import re
 from unittest import TestCase
 
-from mo_logs import logger, Except
+from mo_logs import logger
 from mo_testing.fuzzytestcase import add_error_reporting
 
 from mo_sql_parsing import format, parse
@@ -43,7 +43,6 @@ def verify_formatting(self, expected_sql, expected_json):
             new_json = parse(new_sql)
             self.assertEqual(new_json, expected_json)
         except Exception as cause:
-            cause = Except.wrap(cause)
             logger.error(
                 EXCEPTION_MESSAGE,
                 expected_sql=expected_sql,
@@ -1838,4 +1837,20 @@ def test_issue_255_unpivot(self):
                 "name": "u",
             },
         }
-        self.verify_formatting(sql, json)
\ No newline at end of file
+        self.verify_formatting(sql, json)
+
+    def test_issue_255_unpivot2(self):
+        sql = """SELECT col_a, col_b FROM tab UNPIVOT EXCLUDE NULLS (col_a FOR col_b IN (col_1, col_2)) AS u"""
+        json = {
+            "select": [{"value": "col_a"}, {"value": "col_b"}],
+            "from": "tab",
+            "unpivot": {
+                "nulls": False,
+                "value": "col_a",
+                "for": "col_b",
+                "in": ["col_1", "col_2"],
+                "name": "u",
+            },
+        }
+        self.verify_formatting(sql, json)
+