From 47168fb1e2684a9f0c885d216fac076226bfea52 Mon Sep 17 00:00:00 2001
From: Oliver Mannion <125105+tekumara@users.noreply.github.com>
Date: Fri, 5 Jan 2024 11:45:11 +1100
Subject: [PATCH] feat: support OBJECT_CONSTRUCT_KEEP_NULL

and strip NULL from OBJECT_CONSTRUCT arguments
---
 fakesnow/transforms.py   | 21 +++++++++------------
 pyproject.toml           |  2 +-
 tests/test_fakes.py      | 13 ++++++++++++-
 tests/test_transforms.py |  7 +++++--
 4 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/fakesnow/transforms.py b/fakesnow/transforms.py
index 7c462fe..da1d10b 100644
--- a/fakesnow/transforms.py
+++ b/fakesnow/transforms.py
@@ -410,22 +410,19 @@ def sample(expression: exp.Expression) -> exp.Expression:
 
 
 def object_construct(expression: exp.Expression) -> exp.Expression:
-    """Convert object_construct to return a json string
+    """Convert OBJECT_CONSTRUCT to TO_JSON.
 
-    Because internally snowflake stores OBJECT types as a json string.
+    Internally snowflake stores OBJECT types as a json string, so the Duckdb JSON type most closely matches.
 
-    Example:
-        >>> import sqlglot
-        >>> sqlglot.parse_one("SELECT OBJECT_CONSTRUCT('a',1,'b','BBBB', 'c',null)", read="snowflake").transform(object_construct).sql(dialect="duckdb")
-        "SELECT TO_JSON({'a': 1, 'b': 'BBBB', 'c': NULL})"
-    Args:
-        expression (exp.Expression): the expression that will be transformed.
-
-    Returns:
-        exp.Expression: The transformed expression.
-    """  # noqa: E501
+    See https://docs.snowflake.com/en/sql-reference/functions/object_construct
+    """
 
     if isinstance(expression, exp.Struct):
+        # remove expressions containing NULL
+        for enull in expression.find_all(exp.Null):
+            if enull.parent:
+                enull.parent.pop()
+
         return exp.Anonymous(this="TO_JSON", expressions=[expression])
 
     return expression
diff --git a/pyproject.toml b/pyproject.toml
index 2717d00..b396e41 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -11,7 +11,7 @@ dependencies = [
     "duckdb~=0.9.2",
     "pyarrow",
     "snowflake-connector-python",
-    "sqlglot @ git+https://github.com/tobymao/sqlglot.git@c246285",
+    "sqlglot~=20.7.1",
 ]
 
 [project.urls]
diff --git a/tests/test_fakes.py b/tests/test_fakes.py
index 34098df..766a677 100644
--- a/tests/test_fakes.py
+++ b/tests/test_fakes.py
@@ -756,7 +756,7 @@ def test_schema_drop(cur: snowflake.connector.cursor.SnowflakeCursor):
 def test_semi_structured_types(cur: snowflake.connector.cursor.SnowflakeCursor):
     def indent(rows: Sequence[tuple]) -> list[tuple]:
         # indent duckdb json strings to match snowflake json strings
-        return [(json.dumps(json.loads(r[0]), indent=2), *r[1:]) for r in rows]
+        return [(*[json.dumps(json.loads(c), indent=2) for c in r],) for r in rows]
 
     cur.execute("create or replace table semis (emails array, name object, notes variant)")
     cur.execute(
@@ -780,6 +780,17 @@ def indent(rows: Sequence[tuple]) -> list[tuple]:
     cur.execute("select notes[0] from semis")
     assert cur.fetchall() == [('"foo"',), (None,)]
 
+    cur.execute(
+        """
+            SELECT OBJECT_CONSTRUCT('key_1', 'one', 'key_2', NULL) AS WITHOUT_KEEP_NULL,
+                   OBJECT_CONSTRUCT_KEEP_NULL('key_1', 'one', 'key_2', NULL) AS KEEP_NULL_1,
+                   OBJECT_CONSTRUCT_KEEP_NULL('key_1', 'one', NULL, 'two') AS KEEP_NULL_2
+        """
+    )
+    assert indent(cur.fetchall()) == [  # type: ignore
+        ('{\n  "key_1": "one"\n}', '{\n  "key_1": "one",\n  "key_2": null\n}', '{\n  "key_1": "one"\n}')
+    ]
+
 
 def test_sqlstate(cur: snowflake.connector.cursor.SnowflakeCursor):
     cur.execute("select 'hello world'")
diff --git a/tests/test_transforms.py b/tests/test_transforms.py
index 13d52e2..7279cd1 100644
--- a/tests/test_transforms.py
+++ b/tests/test_transforms.py
@@ -196,10 +196,13 @@ def test_json_extract_cast_as_varchar() -> None:
 
 def test_object_construct() -> None:
     assert (
-        sqlglot.parse_one("SELECT OBJECT_CONSTRUCT('a',1,'b','BBBB', 'c',null)", read="snowflake")
+        sqlglot.parse_one(
+            "SELECT OBJECT_CONSTRUCT('a',1,'b','BBBB','c',null,'d',PARSE_JSON('NULL'), null, 'foo')",
+            read="snowflake",
+        )
         .transform(object_construct)
         .sql(dialect="duckdb")
-        == "SELECT TO_JSON({'a': 1, 'b': 'BBBB', 'c': NULL})"
+        == "SELECT TO_JSON({'a': 1, 'b': 'BBBB', 'd': JSON('NULL')})"
     )