Repository: tobymao/sqlglot

Issue:
Support alter table set tag (snowflake)
```
>>> sqlglot.parse_one("alter table table1 set tag foo='bar'", read="snowflake")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/tekumara/code3/sqlglot/sqlglot/__init__.py", line 147, in parse_one
    result = dialect.parse(sql, **opts)
  File "/Users/tekumara/code3/sqlglot/sqlglot/dialects/dialect.py", line 163, in parse
    return self.parser(**opts).parse(self.tokenizer.tokenize(sql), sql)
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 757, in parse
    return self._parse(
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 820, in _parse
    expressions.append(parse_method(self))
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 950, in _parse_statement
    return self.STATEMENT_PARSERS[self._prev.token_type](self)
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 436, in <lambda>
    TokenType.ALTER: lambda self: self._parse_alter(),
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 3702, in _parse_alter
    return self.expression(exp.AlterTable, this=this, exists=exists, actions=actions)
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 890, in expression
    self.validate_expression(instance)
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 908, in validate_expression
    self.raise_error(error_message)
  File "/Users/tekumara/code3/sqlglot/sqlglot/parser.py", line 866, in raise_error
    raise error
sqlglot.errors.ParseError: Required keyword: 'actions' missing for <class 'sqlglot.expressions.AlterTable'>. Line 1, Col: 20.
  alter table table1 set tag foo='bar'
```
See https://docs.snowflake.com/en/sql-reference/sql/alter-table



Discussion:

Patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 63e9ad9d..bd6cd2db 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import typing as t
+
 from sqlglot import exp, generator, parser, tokens
 from sqlglot.dialects.dialect import (
     Dialect,
@@ -195,6 +197,17 @@ class Snowflake(Dialect):
             ),
         }
 
+        ALTER_PARSERS = {
+            **parser.Parser.ALTER_PARSERS,  # type: ignore
+            "UNSET": lambda self: self._parse_alter_table_set_tag(unset=True),
+            "SET": lambda self: self._parse_alter_table_set_tag(),
+        }
+
+        def _parse_alter_table_set_tag(self, unset: bool = False) -> exp.Expression:
+            self._match_text_seq("TAG")
+            parser = t.cast(t.Callable, self._parse_id_var if unset else self._parse_conjunction)
+            return self.expression(exp.SetTag, expressions=self._parse_csv(parser), unset=unset)
+
     class Tokenizer(tokens.Tokenizer):
         QUOTES = ["'", "$$"]
         STRING_ESCAPES = ["\\", "'"]
@@ -297,6 +310,10 @@ class Snowflake(Dialect):
                 return self.no_identify(lambda: super(self.__class__, self).values_sql(expression))
             return super().values_sql(expression)
 
+        def settag_sql(self, expression: exp.SetTag) -> str:
+            action = "UNSET" if expression.args.get("unset") else "SET"
+            return f"{action} TAG {self.expressions(expression)}"
+
         def select_sql(self, expression: exp.Select) -> str:
             """Due to a bug in Snowflake we want to make sure that all columns in a VALUES table alias are unquoted and also
             that all columns in a SELECT are unquoted. We also want to make sure that after we find matches where we need
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 9744a1f4..8db74dd2 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -930,6 +930,10 @@ class RenameTable(Expression):
     pass
 
 
+class SetTag(Expression):
+    arg_types = {"expressions": True, "unset": False}
+
+
 class ColumnConstraint(Expression):
     arg_types = {"this": False, "kind": True}
 
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index f3e81a8e..3fcfba2a 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -640,6 +640,14 @@ class Parser(metaclass=_Parser):
         "UPPERCASE": lambda self: self.expression(exp.UppercaseColumnConstraint),
     }
 
+    ALTER_PARSERS = {
+        "ADD": lambda self: self._parse_alter_table_add(),
+        "ALTER": lambda self: self._parse_alter_table_alter(),
+        "DELETE": lambda self: self.expression(exp.Delete, where=self._parse_where()),
+        "DROP": lambda self: self._parse_alter_table_drop(),
+        "RENAME": lambda self: self._parse_alter_table_rename(),
+    }
+
     SCHEMA_UNNAMED_CONSTRAINTS = {"CHECK", "FOREIGN KEY", "LIKE", "PRIMARY KEY", "UNIQUE"}
 
     NO_PAREN_FUNCTION_PARSERS = {
@@ -3648,6 +3656,47 @@ class Parser(metaclass=_Parser):
 
         return self.expression(exp.AddConstraint, this=this, expression=expression)
 
+    def _parse_alter_table_add(self) -> t.List[t.Optional[exp.Expression]]:
+        index = self._index - 1
+
+        if self._match_set(self.ADD_CONSTRAINT_TOKENS):
+            return self._parse_csv(self._parse_add_constraint)
+
+        self._retreat(index)
+        return self._parse_csv(self._parse_add_column)
+
+    def _parse_alter_table_alter(self) -> exp.Expression:
+        self._match(TokenType.COLUMN)
+        column = self._parse_field(any_token=True)
+
+        if self._match_pair(TokenType.DROP, TokenType.DEFAULT):
+            return self.expression(exp.AlterColumn, this=column, drop=True)
+        if self._match_pair(TokenType.SET, TokenType.DEFAULT):
+            return self.expression(exp.AlterColumn, this=column, default=self._parse_conjunction())
+
+        self._match_text_seq("SET", "DATA")
+        return self.expression(
+            exp.AlterColumn,
+            this=column,
+            dtype=self._match_text_seq("TYPE") and self._parse_types(),
+            collate=self._match(TokenType.COLLATE) and self._parse_term(),
+            using=self._match(TokenType.USING) and self._parse_conjunction(),
+        )
+
+    def _parse_alter_table_drop(self) -> t.List[t.Optional[exp.Expression]]:
+        index = self._index - 1
+
+        partition_exists = self._parse_exists()
+        if self._match(TokenType.PARTITION, advance=False):
+            return self._parse_csv(lambda: self._parse_drop_partition(exists=partition_exists))
+
+        self._retreat(index)
+        return self._parse_csv(self._parse_drop_column)
+
+    def _parse_alter_table_rename(self) -> exp.Expression:
+        self._match_text_seq("TO")
+        return self.expression(exp.RenameTable, this=self._parse_table(schema=True))
+
     def _parse_alter(self) -> t.Optional[exp.Expression]:
         if not self._match(TokenType.TABLE):
             return self._parse_as_command(self._prev)
@@ -3655,50 +3704,12 @@ class Parser(metaclass=_Parser):
         exists = self._parse_exists()
         this = self._parse_table(schema=True)
 
-        actions: t.Optional[exp.Expression | t.List[t.Optional[exp.Expression]]] = None
-
-        index = self._index
-        if self._match(TokenType.DELETE):
-            actions = [self.expression(exp.Delete, where=self._parse_where())]
-        elif self._match_text_seq("ADD"):
-            if self._match_set(self.ADD_CONSTRAINT_TOKENS):
-                actions = self._parse_csv(self._parse_add_constraint)
-            else:
-                self._retreat(index)
-                actions = self._parse_csv(self._parse_add_column)
-        elif self._match_text_seq("DROP"):
-            partition_exists = self._parse_exists()
+        if not self._curr:
+            return None
 
-            if self._match(TokenType.PARTITION, advance=False):
-                actions = self._parse_csv(
-                    lambda: self._parse_drop_partition(exists=partition_exists)
-                )
-            else:
-                self._retreat(index)
-                actions = self._parse_csv(self._parse_drop_column)
-        elif self._match_text_seq("RENAME", "TO"):
-            actions = self.expression(exp.RenameTable, this=self._parse_table(schema=True))
-        elif self._match_text_seq("ALTER"):
-            self._match(TokenType.COLUMN)
-            column = self._parse_field(any_token=True)
-
-            if self._match_pair(TokenType.DROP, TokenType.DEFAULT):
-                actions = self.expression(exp.AlterColumn, this=column, drop=True)
-            elif self._match_pair(TokenType.SET, TokenType.DEFAULT):
-                actions = self.expression(
-                    exp.AlterColumn, this=column, default=self._parse_conjunction()
-                )
-            else:
-                self._match_text_seq("SET", "DATA")
-                actions = self.expression(
-                    exp.AlterColumn,
-                    this=column,
-                    dtype=self._match_text_seq("TYPE") and self._parse_types(),
-                    collate=self._match(TokenType.COLLATE) and self._parse_term(),
-                    using=self._match(TokenType.USING) and self._parse_conjunction(),
-                )
+        parser = self.ALTER_PARSERS.get(self._curr.text.upper())
+        actions = ensure_list(self._advance() or parser(self)) if parser else []  # type: ignore
 
-        actions = ensure_list(actions)
         return self.expression(exp.AlterTable, this=this, exists=exists, actions=actions)
 
     def _parse_show(self) -> t.Optional[exp.Expression]:


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 5e0f5796..b955f22a 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -10,6 +10,8 @@ class TestSnowflake(Validator):
         self.validate_identity("SELECT REGEXP_LIKE(a, b, c)")
         self.validate_identity("PUT file:///dir/tmp.csv @%table")
         self.validate_identity("CREATE TABLE foo (bar FLOAT AUTOINCREMENT START 0 INCREMENT 1)")
+        self.validate_identity("ALTER TABLE IF EXISTS foo SET TAG a = 'a', b = 'b', c = 'c'")
+        self.validate_identity("ALTER TABLE foo UNSET TAG a, b, c")
         self.validate_identity(
             'COPY INTO NEW_TABLE ("foo", "bar") FROM (SELECT $1, $2, $3, $4 FROM @%old_table)'
         )