Repository: tobymao/sqlglot

Issue:
Support staged files in Snowflake
Hi,

Looking at Snowflake documentation of 'Querying data in staged files" (https://docs.snowflake.com/en/user-guide/querying-stage), it seems like the below query is a valid one:

`
SELECT parse_json($1):a.b FROM @mystage2/data1.json.gz;
`

However I get a parse error, seems like because of the from suffix ("/data1.json.gz")

`
FROM @mystage2\data1.json.gz 
sqlglot.errors.ParseError: Invalid expression / Unexpected token
`

Any suggestions?

Thanks!


Discussion:

Patch:
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index 03b116e4..c02858e1 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -318,6 +318,7 @@ class ClickHouse(Dialect):
         QUERY_HINTS = False
         STRUCT_DELIMITER = ("(", ")")
         NVL2_SUPPORTED = False
+        TABLESAMPLE_REQUIRES_PARENS = False
 
         STRING_TYPE_MAPPING = {
             exp.DataType.Type.CHAR: "String",
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index a4278707..afca44ed 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -51,6 +51,32 @@ TIME_DIFF_FACTOR = {
 DIFF_MONTH_SWITCH = ("YEAR", "QUARTER", "MONTH")
 
 
+def _create_sql(self, expression: exp.Create) -> str:
+    expression = expression.copy()
+
+    # remove UNIQUE column constraints
+    for constraint in expression.find_all(exp.UniqueColumnConstraint):
+        if constraint.parent:
+            constraint.parent.pop()
+
+    properties = expression.args.get("properties")
+    temporary = any(
+        isinstance(prop, exp.TemporaryProperty)
+        for prop in (properties.expressions if properties else [])
+    )
+
+    # CTAS with temp tables map to CREATE TEMPORARY VIEW
+    kind = expression.args["kind"]
+    if kind.upper() == "TABLE" and temporary:
+        if expression.expression:
+            return f"CREATE TEMPORARY VIEW {self.sql(expression, 'this')} AS {self.sql(expression, 'expression')}"
+        else:
+            # CREATE TEMPORARY TABLE may require storage provider
+            expression = self.temporary_storage_provider(expression)
+
+    return create_with_partitions_sql(self, expression)
+
+
 def _add_date_sql(self: Hive.Generator, expression: exp.DateAdd | exp.DateSub) -> str:
     unit = expression.text("unit").upper()
     func, multiplier = DATE_DELTA_INTERVAL.get(unit, ("DATE_ADD", 1))
@@ -429,7 +455,7 @@ class Hive(Dialect):
             if e.args.get("allow_null")
             else "NOT NULL",
             exp.VarMap: var_map_sql,
-            exp.Create: create_with_partitions_sql,
+            exp.Create: _create_sql,
             exp.Quantile: rename_func("PERCENTILE"),
             exp.ApproxQuantile: rename_func("PERCENTILE_APPROX"),
             exp.RegexpExtract: regexp_extract_sql,
@@ -480,6 +506,10 @@ class Hive(Dialect):
             exp.VolatileProperty: exp.Properties.Location.UNSUPPORTED,
         }
 
+        def temporary_storage_provider(self, expression: exp.Create) -> exp.Create:
+            # Hive has no temporary storage provider (there are hive settings though)
+            return expression
+
         def parameter_sql(self, expression: exp.Parameter) -> str:
             this = self.sql(expression, "this")
             parent = expression.parent
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 3edd84af..0ba6a113 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -318,6 +318,28 @@ class Snowflake(Dialect):
             "TERSE PRIMARY KEYS": _show_parser("PRIMARY KEYS"),
         }
 
+        STAGED_FILE_SINGLE_TOKENS = {
+            TokenType.DOT,
+            TokenType.MOD,
+            TokenType.SLASH,
+        }
+
+        def _parse_table_parts(self, schema: bool = False) -> exp.Table:
+            # https://docs.snowflake.com/en/user-guide/querying-stage
+            if self._match_text_seq("@"):
+                table_name = "@"
+                while True:
+                    self._advance()
+                    table_name += self._prev.text
+                    if not self._match_set(self.STAGED_FILE_SINGLE_TOKENS, advance=False):
+                        break
+                    while self._match_set(self.STAGED_FILE_SINGLE_TOKENS):
+                        table_name += self._prev.text
+
+                return self.expression(exp.Table, this=exp.Identifier(this=table_name))
+
+            return super()._parse_table_parts(schema=schema)
+
         def _parse_id_var(
             self,
             any_token: bool = True,
diff --git a/sqlglot/dialects/spark.py b/sqlglot/dialects/spark.py
index 63924d43..2eaa2ae9 100644
--- a/sqlglot/dialects/spark.py
+++ b/sqlglot/dialects/spark.py
@@ -96,17 +96,3 @@ class Spark(Spark2):
                 return self.func("DATEDIFF", unit, start, end)
 
             return self.func("DATEDIFF", end, start)
-
-        def create_sql(self, expression: exp.Create) -> str:
-            kind = self.sql(expression, "kind").upper()
-            properties = expression.args.get("properties")
-            temporary = any(
-                isinstance(prop, exp.TemporaryProperty)
-                for prop in (properties.expressions if properties else [])
-            )
-            if kind == "TABLE" and temporary:
-                provider = exp.FileFormatProperty(this=exp.Literal.string("parquet"))
-                expression = expression.copy()
-                expression.args["properties"].append("expressions", provider)
-
-            return super().create_sql(expression)
diff --git a/sqlglot/dialects/spark2.py b/sqlglot/dialects/spark2.py
index 3dc9838a..b0d64fdc 100644
--- a/sqlglot/dialects/spark2.py
+++ b/sqlglot/dialects/spark2.py
@@ -5,7 +5,6 @@ import typing as t
 from sqlglot import exp, transforms
 from sqlglot.dialects.dialect import (
     binary_from_function,
-    create_with_partitions_sql,
     format_time_lambda,
     is_parse_json,
     move_insert_cte_sql,
@@ -17,22 +16,6 @@ from sqlglot.dialects.hive import Hive
 from sqlglot.helper import seq_get
 
 
-def _create_sql(self: Spark2.Generator, e: exp.Create) -> str:
-    kind = e.args["kind"]
-    properties = e.args.get("properties")
-
-    if (
-        kind.upper() == "TABLE"
-        and e.expression
-        and any(
-            isinstance(prop, exp.TemporaryProperty)
-            for prop in (properties.expressions if properties else [])
-        )
-    ):
-        return f"CREATE TEMPORARY VIEW {self.sql(e, 'this')} AS {self.sql(e, 'expression')}"
-    return create_with_partitions_sql(self, e)
-
-
 def _map_sql(self: Spark2.Generator, expression: exp.Map) -> str:
     keys = expression.args.get("keys")
     values = expression.args.get("values")
@@ -192,7 +175,6 @@ class Spark2(Hive):
             exp.AtTimeZone: lambda self, e: f"FROM_UTC_TIMESTAMP({self.sql(e, 'this')}, {self.sql(e, 'zone')})",
             exp.BitwiseLeftShift: rename_func("SHIFTLEFT"),
             exp.BitwiseRightShift: rename_func("SHIFTRIGHT"),
-            exp.Create: _create_sql,
             exp.DateFromParts: rename_func("MAKE_DATE"),
             exp.DateTrunc: lambda self, e: self.func("TRUNC", e.this, e.args.get("unit")),
             exp.DayOfMonth: rename_func("DAYOFMONTH"),
@@ -236,6 +218,12 @@ class Spark2(Hive):
         WRAP_DERIVED_VALUES = False
         CREATE_FUNCTION_RETURN_AS = False
 
+        def temporary_storage_provider(self, expression: exp.Create) -> exp.Create:
+            # spark2, spark, Databricks require a storage provider for temporary tables
+            provider = exp.FileFormatProperty(this=exp.Literal.string("parquet"))
+            expression.args["properties"].append("expressions", provider)
+            return expression
+
         def cast_sql(self, expression: exp.Cast, safe_prefix: t.Optional[str] = None) -> str:
             if is_parse_json(expression.this):
                 schema = f"'{self.sql(expression, 'to')}'"
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index de3cc8a7..577bf4cf 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -210,6 +210,9 @@ class Generator:
     # Whether or not CREATE TABLE .. COPY .. is supported. False means we'll generate CLONE instead of COPY
     SUPPORTS_TABLE_COPY = True
 
+    # Whether or not parentheses are required around the table sample's expression
+    TABLESAMPLE_REQUIRES_PARENS = True
+
     TYPE_MAPPING = {
         exp.DataType.Type.NCHAR: "CHAR",
         exp.DataType.Type.NVARCHAR: "VARCHAR",
@@ -1353,6 +1356,7 @@ class Generator:
         else:
             this = self.sql(expression, "this")
             alias = ""
+
         method = self.sql(expression, "method")
         method = f"{method.upper()} " if method and self.TABLESAMPLE_WITH_METHOD else ""
         numerator = self.sql(expression, "bucket_numerator")
@@ -1364,13 +1368,20 @@ class Generator:
         percent = f"{percent} PERCENT" if percent else ""
         rows = self.sql(expression, "rows")
         rows = f"{rows} ROWS" if rows else ""
+
         size = self.sql(expression, "size")
         if size and self.TABLESAMPLE_SIZE_IS_PERCENT:
             size = f"{size} PERCENT"
+
         seed = self.sql(expression, "seed")
         seed = f" {seed_prefix} ({seed})" if seed else ""
         kind = expression.args.get("kind", "TABLESAMPLE")
-        return f"{this} {kind} {method}({bucket}{percent}{rows}{size}){seed}{alias}"
+
+        expr = f"{bucket}{percent}{rows}{size}"
+        if self.TABLESAMPLE_REQUIRES_PARENS:
+            expr = f"({expr})"
+
+        return f"{this} {kind} {method}{expr}{seed}{alias}"
 
     def pivot_sql(self, expression: exp.Pivot) -> str:
         expressions = self.expressions(expression, flat=True)
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index 5f8a8448..d536d17d 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2734,14 +2734,18 @@ class Parser(metaclass=_Parser):
         )
         method = self._parse_var(tokens=(TokenType.ROW,))
 
-        self._match(TokenType.L_PAREN)
+        matched_l_paren = self._match(TokenType.L_PAREN)
 
         if self.TABLESAMPLE_CSV:
             num = None
             expressions = self._parse_csv(self._parse_primary)
         else:
             expressions = None
-            num = self._parse_primary()
+            num = (
+                self._parse_factor()
+                if self._match(TokenType.NUMBER, advance=False)
+                else self._parse_primary()
+            )
 
         if self._match_text_seq("BUCKET"):
             bucket_numerator = self._parse_number()
@@ -2756,7 +2760,8 @@ class Parser(metaclass=_Parser):
         elif num:
             size = num
 
-        self._match(TokenType.R_PAREN)
+        if matched_l_paren:
+            self._match_r_paren()
 
         if self._match(TokenType.L_PAREN):
             method = self._parse_var()


Test patch:
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index e32b1abf..2a2fd811 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -24,7 +24,9 @@ class TestClickhouse(Validator):
         self.assertEqual(expr.sql(dialect="clickhouse"), "COUNT(x)")
         self.assertIsNone(expr._meta)
 
-        self.validate_identity("SELECT sum(foo * bar) FROM bla SAMPLE (10000000)")
+        self.validate_identity("SELECT * FROM (SELECT a FROM b SAMPLE 0.01)")
+        self.validate_identity("SELECT * FROM (SELECT a FROM b SAMPLE 1 / 10 OFFSET 1 / 2)")
+        self.validate_identity("SELECT sum(foo * bar) FROM bla SAMPLE 10000000")
         self.validate_identity("CAST(x AS Nested(ID UInt32, Serial UInt32, EventTime DATETIME))")
         self.validate_identity("CAST(x AS Enum('hello' = 1, 'world' = 2))")
         self.validate_identity("CAST(x AS Enum('hello', 'world'))")
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 70fc2926..a85dc841 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -8,34 +8,11 @@ class TestSnowflake(Validator):
     dialect = "snowflake"
 
     def test_snowflake(self):
-        self.validate_identity(
-            'DESCRIBE TABLE "SNOWFLAKE_SAMPLE_DATA"."TPCDS_SF100TCL"."WEB_SITE" type=stage'
-        )
-
-        self.validate_all(
-            "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
-            read={
-                "oracle": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
-            },
-            write={
-                "oracle": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
-                "snowflake": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
-            },
-        )
-        self.validate_all(
-            "SELECT INSERT(a, 0, 0, 'b')",
-            read={
-                "mysql": "SELECT INSERT(a, 0, 0, 'b')",
-                "snowflake": "SELECT INSERT(a, 0, 0, 'b')",
-                "tsql": "SELECT STUFF(a, 0, 0, 'b')",
-            },
-            write={
-                "mysql": "SELECT INSERT(a, 0, 0, 'b')",
-                "snowflake": "SELECT INSERT(a, 0, 0, 'b')",
-                "tsql": "SELECT STUFF(a, 0, 0, 'b')",
-            },
-        )
-
+        self.validate_identity("SELECT * FROM @~")
+        self.validate_identity("SELECT * FROM @~/some/path/to/file.csv")
+        self.validate_identity("SELECT * FROM @mystage")
+        self.validate_identity("SELECT * FROM @namespace.mystage/path/to/file.json.gz")
+        self.validate_identity("SELECT * FROM @namespace.%table_name/path/to/file.json.gz")
         self.validate_identity("LISTAGG(data['some_field'], ',')")
         self.validate_identity("WEEKOFYEAR(tstamp)")
         self.validate_identity("SELECT SUM(amount) FROM mytable GROUP BY ALL")
@@ -64,6 +41,9 @@ class TestSnowflake(Validator):
         self.validate_identity("COMMENT IF EXISTS ON TABLE foo IS 'bar'")
         self.validate_identity("SELECT CONVERT_TIMEZONE('UTC', 'America/Los_Angeles', col)")
         self.validate_identity("REGEXP_REPLACE('target', 'pattern', '\n')")
+        self.validate_identity(
+            'DESCRIBE TABLE "SNOWFLAKE_SAMPLE_DATA"."TPCDS_SF100TCL"."WEB_SITE" type=stage'
+        )
         self.validate_identity(
             'COPY INTO NEW_TABLE ("foo", "bar") FROM (SELECT $1, $2, $3, $4 FROM @%old_table)'
         )
@@ -82,11 +62,38 @@ class TestSnowflake(Validator):
             "SELECT {'test': 'best'}::VARIANT",
             "SELECT CAST(OBJECT_CONSTRUCT('test', 'best') AS VARIANT)",
         )
+        self.validate_identity(
+            "SELECT parse_json($1):a.b FROM @mystage2/data1.json.gz",
+            "SELECT PARSE_JSON($1)['a'].b FROM @mystage2/data1.json.gz",
+        )
 
         self.validate_all("CAST(x AS BYTEINT)", write={"snowflake": "CAST(x AS INT)"})
         self.validate_all("CAST(x AS CHAR VARYING)", write={"snowflake": "CAST(x AS VARCHAR)"})
         self.validate_all("CAST(x AS CHARACTER VARYING)", write={"snowflake": "CAST(x AS VARCHAR)"})
         self.validate_all("CAST(x AS NCHAR VARYING)", write={"snowflake": "CAST(x AS VARCHAR)"})
+        self.validate_all(
+            "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
+            read={
+                "oracle": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
+            },
+            write={
+                "oracle": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
+                "snowflake": "SELECT * FROM x START WITH a = b CONNECT BY c = PRIOR d",
+            },
+        )
+        self.validate_all(
+            "SELECT INSERT(a, 0, 0, 'b')",
+            read={
+                "mysql": "SELECT INSERT(a, 0, 0, 'b')",
+                "snowflake": "SELECT INSERT(a, 0, 0, 'b')",
+                "tsql": "SELECT STUFF(a, 0, 0, 'b')",
+            },
+            write={
+                "mysql": "SELECT INSERT(a, 0, 0, 'b')",
+                "snowflake": "SELECT INSERT(a, 0, 0, 'b')",
+                "tsql": "SELECT STUFF(a, 0, 0, 'b')",
+            },
+        )
         self.validate_all(
             "ARRAY_GENERATE_RANGE(0, 3)",
             write={
diff --git a/tests/dialects/test_tsql.py b/tests/dialects/test_tsql.py
index 960f65ef..23a6037c 100644
--- a/tests/dialects/test_tsql.py
+++ b/tests/dialects/test_tsql.py
@@ -6,6 +6,45 @@ class TestTSQL(Validator):
     dialect = "tsql"
 
     def test_tsql(self):
+        self.validate_all(
+            "CREATE TABLE #mytemptable (a INTEGER)",
+            read={
+                "duckdb": "CREATE TEMPORARY TABLE mytemptable (a INT)",
+            },
+            write={
+                "tsql": "CREATE TABLE #mytemptable (a INTEGER)",
+                "snowflake": "CREATE TEMPORARY TABLE mytemptable (a INT)",
+                "duckdb": "CREATE TEMPORARY TABLE mytemptable (a INT)",
+                "oracle": "CREATE TEMPORARY TABLE mytemptable (a NUMBER)",
+                "hive": "CREATE TEMPORARY TABLE mytemptable (a INT)",
+                "spark2": "CREATE TEMPORARY TABLE mytemptable (a INT) USING PARQUET",
+                "spark": "CREATE TEMPORARY TABLE mytemptable (a INT) USING PARQUET",
+                "databricks": "CREATE TEMPORARY TABLE mytemptable (a INT) USING PARQUET",
+            },
+        )
+        self.validate_all(
+            "CREATE TABLE #mytemp (a INTEGER, b CHAR(2), c TIME(4), d FLOAT(24))",
+            write={
+                "spark": "CREATE TEMPORARY TABLE mytemp (a INT, b CHAR(2), c TIMESTAMP, d FLOAT) USING PARQUET",
+                "tsql": "CREATE TABLE #mytemp (a INTEGER, b CHAR(2), c TIME(4), d FLOAT(24))",
+            },
+        )
+        self.validate_all(
+            """CREATE TABLE [dbo].[mytable](
+                [email] [varchar](255) NOT NULL,
+                CONSTRAINT [UN_t_mytable] UNIQUE NONCLUSTERED 
+                (
+                    [email] ASC
+                )
+                )""",
+            write={
+                "hive": "CREATE TABLE `dbo`.`mytable` (`email` VARCHAR(255) NOT NULL)",
+                "spark2": "CREATE TABLE `dbo`.`mytable` (`email` VARCHAR(255) NOT NULL)",
+                "spark": "CREATE TABLE `dbo`.`mytable` (`email` VARCHAR(255) NOT NULL)",
+                "databricks": "CREATE TABLE `dbo`.`mytable` (`email` VARCHAR(255) NOT NULL)",
+            },
+        )
+
         self.validate_all(
             "CREATE TABLE x ( A INTEGER NOT NULL, B INTEGER NULL )",
             write={
@@ -568,18 +607,6 @@ class TestTSQL(Validator):
                 "tsql": "CREATE TABLE #mytemp (a INTEGER, b CHAR(2), c TIME(4), d FLOAT(24))",
             },
         )
-        self.validate_all(
-            "CREATE TABLE #mytemptable (a INTEGER)",
-            read={
-                "duckdb": "CREATE TEMPORARY TABLE mytemptable (a INT)",
-            },
-            write={
-                "tsql": "CREATE TABLE #mytemptable (a INTEGER)",
-                "snowflake": "CREATE TEMPORARY TABLE mytemptable (a INT)",
-                "duckdb": "CREATE TEMPORARY TABLE mytemptable (a INT)",
-                "oracle": "CREATE TEMPORARY TABLE mytemptable (a NUMBER)",
-            },
-        )
 
     def test_insert_cte(self):
         self.validate_all(
diff --git a/tests/test_transpile.py b/tests/test_transpile.py
index a5b19778..72960b02 100644
--- a/tests/test_transpile.py
+++ b/tests/test_transpile.py
@@ -604,7 +604,7 @@ LEFT OUTER JOIN b""",
         self.validate(
             "CREATE TEMPORARY TABLE test AS SELECT 1",
             "CREATE TEMPORARY VIEW test AS SELECT 1",
-            write="spark",
+            write="spark2",
         )
 
     @mock.patch("sqlglot.helper.logger")