Repository: tobymao/sqlglot

Issue:
Snowflake Alter Table ... Add Column ... WITH TAG (...) doesn't parse correctly
**Before you file an issue**
- Make sure you specify the "read" dialect eg. `parse_one(sql, read="spark")`: **snowflake**
- Make sure you specify the "write" dialect eg. `ast.sql(dialect="duckdb")`: **snowflake**
- Check if the issue still exists on main

**Fully reproducible code snippet**
Please include a fully reproducible code snippet or the input sql, dialect, and expected output.

Snowflake syntax allows an optional WITH token to add a column with a tag, however, the parser doesn't like the WITH token in this position. If the WITH token is removed the parser doesn't know how to parse the TAG syntax for added columns.

This example doesn't understand the **WITH** token as part of the Add Column query:
```python
from sqlglot import parse_one

query = """
    ALTER TABLE db_name.schmaName.tblName add column COLUMN_1 varchar not null WITH TAG (key1 = 'value_1');
"""
parsed = parse_one(query, read='snowflake')
print(repr(parsed))

...
sqlglot.errors.ParseError: Expecting (. Line 2, Col: 84.
            ALTER TABLE db_name.schmaName.tblName add column COLUMN_1 varchar not null TAG (key1 = 'value_1');
```

This example doesn't raise an exception but still doesn't know how to parse the query to add a column with a tag.
```python
from sqlglot import parse_one

query = """
    ALTER TABLE db_name.schmaName.tblName add column COLUMN_1 varchar not null TAG (key1 = 'value_1');
"""
parsed = parse_one(query, read='snowflake')
print(repr(parsed))

...
'ALTER TABLE db_name.schmaName.tblName add column COLUMN_1 varchar not null TAG (key1 = 'value_1')' contains unsupported syntax. Falling back to parsing as a 'Command'.
Command(this=ALTER, expression=TABLE db_name.schmaName.tblName add column COLUMN_1 varchar not null TAG (key1 = 'value_1'))
...

```

**Official Documentation**
Please include links to official SQL documentation related to your issue.
https://docs.snowflake.com/en/sql-reference/sql/alter-table#table-column-actions-tablecolumnaction


Discussion:

Patch:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 0de4a80e..f72b29ef 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -473,6 +473,14 @@ class Snowflake(Dialect):
             "TERSE USERS": _show_parser("USERS"),
         }
 
+        CONSTRAINT_PARSERS = {
+            **parser.Parser.CONSTRAINT_PARSERS,
+            "WITH": lambda self: self._parse_with_constraint(),
+            "MASKING": lambda self: self._parse_with_constraint(),
+            "PROJECTION": lambda self: self._parse_with_constraint(),
+            "TAG": lambda self: self._parse_with_constraint(),
+        }
+
         STAGED_FILE_SINGLE_TOKENS = {
             TokenType.DOT,
             TokenType.MOD,
@@ -497,6 +505,29 @@ class Snowflake(Dialect):
             ),
         }
 
+        def _parse_with_constraint(self) -> t.Optional[exp.Expression]:
+            if self._prev.token_type != TokenType.WITH:
+                self._retreat(self._index - 1)
+
+            if self._match_text_seq("MASKING", "POLICY"):
+                return self.expression(
+                    exp.MaskingPolicyColumnConstraint,
+                    this=self._parse_id_var(),
+                    expressions=self._match(TokenType.USING)
+                    and self._parse_wrapped_csv(self._parse_id_var),
+                )
+            if self._match_text_seq("PROJECTION", "POLICY"):
+                return self.expression(
+                    exp.ProjectionPolicyColumnConstraint, this=self._parse_id_var()
+                )
+            if self._match(TokenType.TAG):
+                return self.expression(
+                    exp.TagColumnConstraint,
+                    expressions=self._parse_wrapped_csv(self._parse_property),
+                )
+
+            return None
+
         def _parse_create(self) -> exp.Create | exp.Command:
             expression = super()._parse_create()
             if isinstance(expression, exp.Create) and expression.kind in self.NON_TABLE_CREATABLES:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 6da69fa9..7f5b4c56 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -1835,6 +1835,11 @@ class NotForReplicationColumnConstraint(ColumnConstraintKind):
     arg_types = {}
 
 
+# https://docs.snowflake.com/en/sql-reference/sql/create-table
+class MaskingPolicyColumnConstraint(ColumnConstraintKind):
+    arg_types = {"this": True, "expressions": False}
+
+
 class NotNullColumnConstraint(ColumnConstraintKind):
     arg_types = {"allow_null": False}
 
@@ -1844,6 +1849,11 @@ class OnUpdateColumnConstraint(ColumnConstraintKind):
     pass
 
 
+# https://docs.snowflake.com/en/sql-reference/sql/create-table
+class TagColumnConstraint(ColumnConstraintKind):
+    arg_types = {"expressions": True}
+
+
 # https://docs.snowflake.com/en/sql-reference/sql/create-external-table#optional-parameters
 class TransformColumnConstraint(ColumnConstraintKind):
     pass
@@ -1869,6 +1879,11 @@ class PathColumnConstraint(ColumnConstraintKind):
     pass
 
 
+# https://docs.snowflake.com/en/sql-reference/sql/create-table
+class ProjectionPolicyColumnConstraint(ColumnConstraintKind):
+    pass
+
+
 # computed column expression
 # https://learn.microsoft.com/en-us/sql/t-sql/statements/create-table-transact-sql?view=sql-server-ver16
 class ComputedColumnConstraint(ColumnConstraintKind):
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 128d1958..84b336c9 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -123,6 +123,8 @@ class Generator(metaclass=_Generator):
         exp.OnUpdateColumnConstraint: lambda self, e: f"ON UPDATE {self.sql(e, 'this')}",
         exp.OutputModelProperty: lambda self, e: f"OUTPUT{self.sql(e, 'this')}",
         exp.PathColumnConstraint: lambda self, e: f"PATH {self.sql(e, 'this')}",
+        exp.ProjectionPolicyColumnConstraint: lambda self,
+        e: f"PROJECTION POLICY {self.sql(e, 'this')}",
         exp.RemoteWithConnectionModelProperty: lambda self,
         e: f"REMOTE WITH CONNECTION {self.sql(e, 'this')}",
         exp.ReturnsProperty: lambda self, e: (
@@ -139,6 +141,7 @@ class Generator(metaclass=_Generator):
         exp.StabilityProperty: lambda _, e: e.name,
         exp.StrictProperty: lambda *_: "STRICT",
         exp.TemporaryProperty: lambda *_: "TEMPORARY",
+        exp.TagColumnConstraint: lambda self, e: f"TAG ({self.expressions(e, flat=True)})",
         exp.TitleColumnConstraint: lambda self, e: f"TITLE {self.sql(e, 'this')}",
         exp.Timestamp: lambda self, e: self.func("TIMESTAMP", e.this, e.expression),
         exp.ToMap: lambda self, e: f"MAP {self.sql(e, 'this')}",
@@ -3929,3 +3932,11 @@ class Generator(metaclass=_Generator):
             on_sql = self.func("ON", filter_col, retention_period)
 
         return f"DATA_DELETION={on_sql}"
+
+    def maskingpolicycolumnconstraint_sql(
+        self, expression: exp.MaskingPolicyColumnConstraint
+    ) -> str:
+        this = self.sql(expression, "this")
+        expressions = self.expressions(expression, flat=True)
+        expressions = f" USING ({expressions})" if expressions else ""
+        return f"MASKING POLICY {this}{expressions}"


Test patch:
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index f8c2ea16..84430570 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -1186,6 +1186,25 @@ WHERE
         )
 
     def test_ddl(self):
+        for constraint_prefix in ("WITH ", ""):
+            with self.subTest(f"Constraint prefix: {constraint_prefix}"):
+                self.validate_identity(
+                    f"CREATE TABLE t (id INT {constraint_prefix}MASKING POLICY p)",
+                    "CREATE TABLE t (id INT MASKING POLICY p)",
+                )
+                self.validate_identity(
+                    f"CREATE TABLE t (id INT {constraint_prefix}MASKING POLICY p USING (c1, c2, c3))",
+                    "CREATE TABLE t (id INT MASKING POLICY p USING (c1, c2, c3))",
+                )
+                self.validate_identity(
+                    f"CREATE TABLE t (id INT {constraint_prefix}PROJECTION POLICY p)",
+                    "CREATE TABLE t (id INT PROJECTION POLICY p)",
+                )
+                self.validate_identity(
+                    f"CREATE TABLE t (id INT {constraint_prefix}TAG (key1='value_1', key2='value_2'))",
+                    "CREATE TABLE t (id INT TAG (key1='value_1', key2='value_2'))",
+                )
+
         self.validate_identity(
             """create external table et2(
   col1 date as (parse_json(metadata$external_table_partition):COL1::date),
@@ -1210,6 +1229,9 @@ WHERE
         self.validate_identity(
             "CREATE OR REPLACE TAG IF NOT EXISTS cost_center COMMENT='cost_center tag'"
         ).this.assert_is(exp.Identifier)
+        self.validate_identity(
+            "ALTER TABLE db_name.schmaName.tblName ADD COLUMN COLUMN_1 VARCHAR NOT NULL TAG (key1='value_1')"
+        )
         self.validate_identity(
             "DROP FUNCTION my_udf (OBJECT(city VARCHAR, zipcode DECIMAL(38, 0), val ARRAY(BOOLEAN)))"
         )