mraba/underscore_column_id: more test cases and description.md update

snowflakedb · Oct 22, 2024 · f58df7c · f58df7c
1 parent 6f4c15d
commit f58df7c
Show file tree

Hide file tree

Showing 4 changed files with 48 additions and 9 deletions.
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
@@ -11,6 +11,7 @@ Source code is also available at:
 
 - (Unreleased)
 
+  - Fixed quoting of `_` as column name
   - Add support for dynamic tables and required options
   - Add support for hybrid tables
   - Fixed SAWarning when registering functions with existing name in default namespace

diff --git a/src/snowflake/sqlalchemy/base.py b/src/snowflake/sqlalchemy/base.py
@@ -107,6 +107,7 @@
 AUTOCOMMIT_REGEXP = re.compile(
     r"\s*(?:UPDATE|INSERT|DELETE|MERGE|COPY)", re.I | re.UNICODE
 )
+# used for quoting identifiers ie. table names, column names, etc.
 ILLEGAL_INITIAL_CHARACTERS = frozenset({d for d in string.digits}.union({"_", "$"}))
 
 """

diff --git a/tests/test_compiler.py b/tests/test_compiler.py
@@ -33,6 +33,21 @@ def test_now_func(self):
             dialect="snowflake",
         )
 
+    def test_dot_as_valid_identifier(self):
+        _table = table(
+            "table_1745924",
+            column("ca", Integer),
+            column("cb", String),
+            column(".", String),
+        )
+
+        stmt = insert(_table).values({"ca": 1, "cb": "test", ".": "test_"})
+        self.assert_compile(
+            stmt,
+            'INSERT INTO table_1745924 (ca, cb, ".") VALUES (%(ca)s, %(cb)s, %(_)s)',
+            dialect="snowflake",
+        )
+
     def test_underscore_as_valid_identifier(self):
         _table = table(
             "table_1745924",

diff --git a/tests/test_quote_identifiers.py b/tests/test_quote_identifiers.py
@@ -1,6 +1,6 @@
 #
 # Copyright (c) 2012-2023 Snowflake Computing Inc. All rights reserved.
-
+import pytest
 from sqlalchemy import (
     Column,
     Integer,
@@ -16,15 +16,35 @@
 
 from .parameters import CONNECTION_PARAMETERS
 
+# https://docs.snowflake.com/en/sql-reference/identifiers-syntax#double-quoted-identifiers
+DOUBLE_QUOTE_IDENTIFIERS = {i for i in ".'!@#$%^&*"}
+
 
-def test_insert_with_identifier():
+@pytest.mark.parametrize(
+    "identifier",
+    (
+        pytest.param(".", id="dot"),
+        pytest.param("'", id="single_quote"),
+        pytest.param("_", id="underscore"),
+        pytest.param("!", id="exclamation"),
+        pytest.param("@", id="at"),
+        pytest.param("#", id="hash"),
+        pytest.param("$", id="dollar"),
+        pytest.param("%", id="percent"),
+        pytest.param("^", id="caret"),
+        pytest.param("&", id="ampersand"),
+        pytest.param("*", id="asterisk"),
+    ),
+)
+def test_insert_with_identifier_as_column_name(identifier: str):
+    expected_identifier = f"test: {identifier}"
     metadata = MetaData()
     table = Table(
         "table_1745924",
         metadata,
         Column("ca", Integer),
         Column("cb", String),
-        Column("_", String),
+        Column(identifier, String),
     )
 
     engine = create_engine(URL(**CONNECTION_PARAMETERS))
@@ -33,14 +53,16 @@ def test_insert_with_identifier():
         metadata.create_all(engine)
 
         with engine.connect() as connection:
-            connection.execute(insert(table).values(ca=1, cb="test", _="test_"))
             connection.execute(
-                insert(table).values({"ca": 2, "cb": "test", "_": "test_"})
+                insert(table).values(
+                    {
+                        "ca": 1,
+                        "cb": "test",
+                        identifier: f"test: {identifier}",
+                    }
+                )
             )
             result = connection.execute(select(table)).fetchall()
-            assert result == [
-                (1, "test", "test_"),
-                (2, "test", "test_"),
-            ]
+            assert result == [(1, "test", expected_identifier)]
     finally:
         metadata.drop_all(engine)