feat: support OBJECT_CONSTRUCT_KEEP_NULL

and strip NULL from OBJECT_CONSTRUCT arguments

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

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]

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'")

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')})"
     )