Fix!: parse column instead of identifier for SET assignment LHS (#3417)

tobymao · May 7, 2024 · 258ad3b · 258ad3b
1 parent e1b6483
commit 258ad3b
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 10 deletions.
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -5974,7 +5974,7 @@ def _parse_set_item_assignment(
         if kind in ("GLOBAL", "SESSION") and self._match_text_seq("TRANSACTION"):
             return self._parse_set_transaction(global_=kind == "GLOBAL")
 
-        left = self._parse_primary() or self._parse_id_var()
+        left = self._parse_primary() or self._parse_column()
         assignment_delimiter = self._match_texts(("=", "TO"))
 
         if not left or (self.SET_REQUIRES_ASSIGNMENT_DELIMITER and not assignment_delimiter):

diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -10,14 +10,6 @@ class TestSnowflake(Validator):
     dialect = "snowflake"
 
     def test_snowflake(self):
-        self.validate_identity(
-            "MERGE INTO my_db AS ids USING (SELECT new_id FROM my_model WHERE NOT col IS NULL) AS new_ids ON ids.type = new_ids.type AND ids.source = new_ids.source WHEN NOT MATCHED THEN INSERT VALUES (new_ids.new_id)"
-        )
-        self.validate_identity("ALTER TABLE table1 CLUSTER BY (name DESC)")
-        self.validate_identity(
-            "INSERT OVERWRITE TABLE t SELECT 1", "INSERT OVERWRITE INTO t SELECT 1"
-        )
-        self.validate_identity("SELECT rename, replace")
         expr = parse_one("SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
         expr.selects[0].assert_is(exp.AggFunc)
         self.assertEqual(expr.sql(dialect="snowflake"), "SELECT APPROX_TOP_K(C4, 3, 5) FROM t")
@@ -43,6 +35,9 @@ def test_snowflake(self):
   )""",
         )
 
+        self.validate_identity("ALTER TABLE table1 CLUSTER BY (name DESC)")
+        self.validate_identity("SELECT rename, replace")
+        self.validate_identity("ALTER TABLE table1 SET TAG foo.bar = 'baz'")
         self.validate_identity("SELECT TIMEADD(HOUR, 2, CAST('09:05:03' AS TIME))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS MAP(VARCHAR, INT))")
         self.validate_identity("SELECT CAST(OBJECT_CONSTRUCT('a', 1) AS OBJECT(a CHAR NOT NULL))")
@@ -95,6 +90,12 @@ def test_snowflake(self):
         self.validate_identity("SELECT CONVERT_TIMEZONE('UTC', 'America/Los_Angeles', col)")
         self.validate_identity("ALTER TABLE a SWAP WITH b")
         self.validate_identity("SELECT MATCH_CONDITION")
+        self.validate_identity(
+            "MERGE INTO my_db AS ids USING (SELECT new_id FROM my_model WHERE NOT col IS NULL) AS new_ids ON ids.type = new_ids.type AND ids.source = new_ids.source WHEN NOT MATCHED THEN INSERT VALUES (new_ids.new_id)"
+        )
+        self.validate_identity(
+            "INSERT OVERWRITE TABLE t SELECT 1", "INSERT OVERWRITE INTO t SELECT 1"
+        )
         self.validate_identity(
             'DESCRIBE TABLE "SNOWFLAKE_SAMPLE_DATA"."TPCDS_SF100TCL"."WEB_SITE" type=stage'
         )

diff --git a/tests/test_parser.py b/tests/test_parser.py
@@ -503,7 +503,7 @@ def test_set_expression(self):
 
         self.assertIsInstance(set_item, exp.SetItem)
         self.assertIsInstance(set_item.this, exp.EQ)
-        self.assertIsInstance(set_item.this.this, exp.Identifier)
+        self.assertIsInstance(set_item.this.this, exp.Column)
         self.assertIsInstance(set_item.this.expression, exp.Literal)
 
         self.assertEqual(set_item.args.get("kind"), "SESSION")