Fix(snowflake): support fqns in masking/projection policy constraint (#3620)

georgesittas · web-flow · commit 05fe847aeb65 · 2024-06-10T17:49:04.000+03:00
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
@@ -510,15 +510,18 @@ def _parse_with_constraint(self) -> t.Optional[exp.Expression]:
                 self._retreat(self._index - 1)
 
             if self._match_text_seq("MASKING", "POLICY"):
+                policy = self._parse_column()
                 return self.expression(
                     exp.MaskingPolicyColumnConstraint,
-                    this=self._parse_id_var(),
+                    this=policy.to_dot() if isinstance(policy, exp.Column) else policy,
                     expressions=self._match(TokenType.USING)
                     and self._parse_wrapped_csv(self._parse_id_var),
                 )
             if self._match_text_seq("PROJECTION", "POLICY"):
+                policy = self._parse_column()
                 return self.expression(
-                    exp.ProjectionPolicyColumnConstraint, this=self._parse_id_var()
+                    exp.ProjectionPolicyColumnConstraint,
+                    this=policy.to_dot() if isinstance(policy, exp.Column) else policy,
                 )
             if self._match(TokenType.TAG):
                 return self.expression(
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -1196,16 +1196,16 @@ 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)",
+                    f"CREATE TABLE t (id INT {constraint_prefix}MASKING POLICY p.q.r)",
+                    "CREATE TABLE t (id INT MASKING POLICY p.q.r)",
                 )
                 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)",
+                    f"CREATE TABLE t (id INT {constraint_prefix}PROJECTION POLICY p.q.r)",
+                    "CREATE TABLE t (id INT PROJECTION POLICY p.q.r)",
                 )
                 self.validate_identity(
                     f"CREATE TABLE t (id INT {constraint_prefix}TAG (key1='value_1', key2='value_2'))",
