fix: allow unquoted names DuckDB reserves but Snowflake doesn't

fakesnow/copy_into.py

@@ -16,6 +16,7 @@
 
 import fakesnow.transforms.stage as stage
 from fakesnow import logger
+from fakesnow.dialect import FakesnowDuckDB
 from fakesnow.params import MutableParams, pop_qmark_param
 from fakesnow.transforms.file_format import lookup_file_format
 from fakesnow.transforms.options import parse_options
@@ -110,7 +111,7 @@ def copy_into(
                 error_count = 1
                 first_error_message = "File was loaded before."
             else:
-                sql = i.sql(dialect="duckdb")
+                sql = i.sql(dialect=FakesnowDuckDB)
                 logger.log_sql(sql, params)
                 duck_conn.execute(sql, params)
                 (affected_count,) = duck_conn.fetchall()[0]
@@ -554,7 +555,7 @@ def _strip_json_extract(expr: exp.Select) -> exp.Select:
     for p in expr.find_all(exp.Parameter):
         if p == dollar1 and p.parent and (key := p.parent.expression.find(exp.JSONPathKey)):
             assert p.parent.parent, expr
-            p.parent.parent.args["this"] = exp.Identifier(this=key.this)
+            p.parent.parent.args["this"] = exp.Identifier(this=key.this, quoted=True)
 
     return expr
 

fakesnow/cursor.py

@@ -29,7 +29,7 @@
 import fakesnow.transforms as transforms
 from fakesnow import logger
 from fakesnow.copy_into import copy_into
-from fakesnow.dialect import parse_one, syntax_error
+from fakesnow.dialect import FakesnowDuckDB, parse_one, syntax_error
 from fakesnow.params import MutableParams
 from fakesnow.rowtype import describe_as_result_metadata
 from fakesnow.transforms import stage
@@ -419,7 +419,7 @@ def _execute(self, transformed: Expr, params: MutableParams | None = None) -> No
 
         cmd = expr.key_command(transformed)
 
-        sql = transformed.sql(dialect="duckdb")
+        sql = transformed.sql(dialect=FakesnowDuckDB)
 
         if not sql:
             raise NotImplementedError(transformed.sql(dialect="snowflake"))

fakesnow/dialect.py

@@ -1,10 +1,104 @@
-"""Parse Snowflake SQL with Snowflake-compatible syntax errors."""
+"""Parse Snowflake SQL with Snowflake-compatible syntax errors, and generate DuckDB SQL."""
 
+import duckdb
 import snowflake.connector.errors
 import sqlglot
 from sqlglot import Expr
+from sqlglot.dialects.duckdb import DuckDB
 from sqlglot.errors import ParseError
 
+# Keywords Snowflake reserves in every context, see https://docs.snowflake.com/en/sql-reference/reserved-keywords
+# Maintain this list manually: sqlglot's Snowflake reserved-keyword set is empty. Include only the
+# "Reserved by ANSI", "Reserved by Snowflake", and "Reserved by Snowflake and others" categories.
+SNOWFLAKE_RESERVED_KEYWORDS = {
+    "all",
+    "alter",
+    "and",
+    "any",
+    "as",
+    "between",
+    "by",
+    "check",
+    "column",
+    "connect",
+    "create",
+    "current",
+    "delete",
+    "distinct",
+    "drop",
+    "else",
+    "exists",
+    "following",
+    "for",
+    "from",
+    "grant",
+    "group",
+    "having",
+    "ilike",
+    "in",
+    "increment",
+    "insert",
+    "intersect",
+    "into",
+    "is",
+    "like",
+    "minus",
+    "not",
+    "null",
+    "of",
+    "on",
+    "or",
+    "order",
+    "qualify",
+    "regexp",
+    "revoke",
+    "rlike",
+    "row",
+    "rows",
+    "sample",
+    "select",
+    "set",
+    "some",
+    "start",
+    "table",
+    "tablesample",
+    "then",
+    "to",
+    "trigger",
+    "union",
+    "unique",
+    "update",
+    "values",
+    "whenever",
+    "where",
+    "window",
+    "with",
+}
+
+
+def _duckdb_reserved_keywords() -> set[str]:
+    """Keywords DuckDB can't take as an unquoted table, column or alias name."""
+    with duckdb.connect() as conn:
+        rows = conn.execute(
+            "SELECT keyword_name FROM duckdb_keywords() WHERE keyword_category IN ('reserved', 'type_function')"
+        ).fetchall()
+    return {name for (name,) in rows}
+
+
+class FakesnowDuckDB(DuckDB):
+    """DuckDB, quoting every identifier DuckDB reserves and Snowflake doesn't.
+
+    Snowflake allows names such as SEMI or PIVOT unquoted but DuckDB doesn't, and sqlglot's list of DuckDB
+    reserved keywords misses some, so take them from DuckDB itself. Quoting doesn't change how a name
+    resolves, because DuckDB matches quoted identifiers case-insensitively too. Words Snowflake reserves
+    too, such as QUALIFY, stay unquoted so DuckDB rejects them as Snowflake does.
+    """
+
+    class Generator(DuckDB.Generator):
+        RESERVED_KEYWORDS = DuckDB.Generator.RESERVED_KEYWORDS | (
+            _duckdb_reserved_keywords() - SNOWFLAKE_RESERVED_KEYWORDS
+        )
+
 
 def parse(command: str) -> list[Expr | None]:
     try:

fakesnow/transforms/transforms.py

@@ -10,6 +10,7 @@
 from duckdb import DuckDBPyConnection
 from sqlglot import Expr, exp
 
+from fakesnow.dialect import FakesnowDuckDB
 from fakesnow.params import MutableParams, pop_qmark_param
 from fakesnow.variables import Variables
 
@@ -1440,10 +1441,10 @@ def create_table_as(expression: Expr, duck_conn: DuckDBPyConnection) -> Expr:
 
         if any(isinstance(expr, exp.Star) for expr in select_query.expressions):
             # convert SELECT * to SELECT <col1>, <col2>, ...
-            duck_conn.execute(f"DESCRIBE {select_query}")
+            duck_conn.execute(f"DESCRIBE {select_query.sql(dialect=FakesnowDuckDB)}")
             select_query.set(
                 "expressions",
-                [exp.Column(this=exp.Identifier(this=col[0], quoted=False)) for col in duck_conn.fetchall()],
+                [exp.Column(this=exp.Identifier(this=col[0], quoted=True)) for col in duck_conn.fetchall()],
             )
 
         if len(select_query.expressions) != len(create_col_defs):

tests/test_copy_into.py

@@ -131,7 +131,7 @@ class Case(NamedTuple):
             FILE_FORMAT = (TYPE = 'PARQUET')
             """,
             expected_inserts=[
-                "INSERT INTO TABLE1 SELECT CAST((B) AS BIGINT), CAST((A) AS BIGINT) FROM READ_PARQUET('s3://{bucket}/foo.parquet')"
+                'INSERT INTO TABLE1 SELECT CAST(("B") AS BIGINT), CAST(("A") AS BIGINT) FROM READ_PARQUET(\'s3://{bucket}/foo.parquet\')'
             ],
             data=pd.DataFrame({"A": [1, 2], "B": [11, 12]}).to_parquet(),
             expected_rows_loaded=2,
@@ -267,6 +267,21 @@ def test_copy_uses_named_csv_with_inline_override(dcur: snowflake.connector.curs
         assert dcur.fetchall() == [{"A": 1, "B": 2}, {"A": 3, "B": 4}]
 
 
+def test_copy_into_table_name_reserved_by_duckdb(dcur: snowflake.connector.cursor.DictCursor) -> None:
+    # DuckDB reserves PIVOT, Snowflake doesn't
+    dcur.execute("CREATE TABLE pivot (a INT, b INT)")
+    with tempfile.NamedTemporaryFile(mode="w+", suffix=".csv") as temp_file:
+        temp_file.write("1,2\n")
+        temp_file.flush()
+
+        dcur.execute("CREATE STAGE reserved_word_stage")
+        dcur.execute(f"PUT 'file://{temp_file.name}' @reserved_word_stage")
+        dcur.execute("COPY INTO pivot FROM @reserved_word_stage")
+
+        dcur.execute("SELECT * FROM pivot")
+        assert dcur.fetchall() == [{"A": 1, "B": 2}]
+
+
 def test_copy_internal_stage_path(dcur: snowflake.connector.cursor.DictCursor) -> None:
     create_table(dcur)
     with tempfile.TemporaryDirectory() as tmp_dir:
@@ -895,6 +910,21 @@ def test_errors_parquet(dcur: snowflake.connector.cursor.DictCursor, s3_client:
     )
 
 
+@pytest.mark.parametrize("name", ["null", "qualify"])
+def test_copy_parquet_reserved_field_name(
+    cur: snowflake.connector.cursor.SnowflakeCursor, s3_client: S3Client, name: str
+) -> None:
+    bucket = upload_file(s3_client, pd.DataFrame({name: [7]}).to_parquet(), key="reserved.parquet")
+    cur.execute("CREATE TABLE target (x INT)")
+    cur.execute(f"CREATE STAGE source_stage URL='s3://{bucket}/'")
+    cur.execute(
+        f"""COPY INTO target FROM (SELECT $1:"{name}"::INT FROM @source_stage)
+        FILES=('reserved.parquet') FILE_FORMAT=(TYPE=PARQUET)"""
+    )
+    cur.execute("SELECT x FROM target")
+    assert cur.fetchall() == [(7,)]
+
+
 def test_copy_parquet_match_by_column_name(dcur: snowflake.connector.cursor.DictCursor, s3_client: S3Client) -> None:
     """Test COPY INTO with MATCH_BY_COLUMN_NAME for parquet files.
 
@@ -1420,8 +1450,7 @@ def test__strip_json_extract():
         )
         .transform(_strip_json_extract)
         .sql()
-        # TODO: quoted????
-        == "SELECT CAST(A AS INT), CAST(B AS INT) FROM @stage1"
+        == 'SELECT CAST("A" AS INT), CAST("B" AS INT) FROM @stage1'
     )
 
     # after all our transforms sql becomes
@@ -1434,7 +1463,7 @@ def test__strip_json_extract():
         )
         .transform(_strip_json_extract)
         .sql()
-        == "SELECT CAST((A) AS BIGINT), CAST((B) AS BIGINT) FROM @stage1"
+        == 'SELECT CAST(("A") AS BIGINT), CAST(("B") AS BIGINT) FROM @stage1'
     )