Repository: tobymao/sqlglot

Issue:
Parsing issue for Snowflake COPY INTO queries with subqueries
Hi folks, 

looks like we're facing an issue when parsing Snowflake `COPY INTO` queries with `SELECT` subqueries:
```
query = "COPY INTO test (c1) FROM (SELECT $1:c1::string FROM @mystage)"
parse_one(query, read="snowflake")
```

... which yields the following error:
```
sqlglot.errors.ParseError: Expected table name but got <Token token_type: TokenType.L_PAREN, text: (, line: 1, col: 32, start: 31, end: 31, comments: []>. Line 1, Col: 32.
COPY INTO t_d56103e1 (c1) FROM (SELECT $1:c1::string FROM @"mystage")
```

(The query has been validated to run successfully against Snowflake.)

Please let me know if you need any additional information, and whether I can help out with a PR/fix for this. 🙌  Thanks!

Discussion:

Patch:
diff --git a/sqlglot/dialects/bigquery.py b/sqlglot/dialects/bigquery.py
index 5450a6cb..0d2cfc1a 100644
--- a/sqlglot/dialects/bigquery.py
+++ b/sqlglot/dialects/bigquery.py
@@ -633,6 +633,7 @@ class BigQuery(Dialect):
                     transforms.eliminate_semi_and_anti_joins,
                 ]
             ),
+            exp.SHA: rename_func("SHA1"),
             exp.SHA2: lambda self, e: self.func(
                 "SHA256" if e.text("length") == "256" else "SHA512", e.this
             ),
diff --git a/sqlglot/dialects/clickhouse.py b/sqlglot/dialects/clickhouse.py
index b0ca65ab..362ccde3 100644
--- a/sqlglot/dialects/clickhouse.py
+++ b/sqlglot/dialects/clickhouse.py
@@ -146,6 +146,9 @@ class ClickHouse(Dialect):
             "TUPLE": exp.Struct.from_arg_list,
             "UNIQ": exp.ApproxDistinct.from_arg_list,
             "XOR": lambda args: exp.Xor(expressions=args),
+            "MD5": exp.MD5Digest.from_arg_list,
+            "SHA256": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(256)),
+            "SHA512": lambda args: exp.SHA2(this=seq_get(args, 0), length=exp.Literal.number(512)),
         }
 
         AGG_FUNCTIONS = {
@@ -717,6 +720,12 @@ class ClickHouse(Dialect):
             ),
             exp.VarMap: lambda self, e: _lower_func(var_map_sql(self, e)),
             exp.Xor: lambda self, e: self.func("xor", e.this, e.expression, *e.expressions),
+            exp.MD5Digest: rename_func("MD5"),
+            exp.MD5: lambda self, e: self.func("LOWER", self.func("HEX", self.func("MD5", e.this))),
+            exp.SHA: rename_func("SHA1"),
+            exp.SHA2: lambda self, e: self.func(
+                "SHA256" if e.text("length") == "256" else "SHA512", e.this
+            ),
         }
 
         PROPERTIES_LOCATION = {
diff --git a/sqlglot/dialects/hive.py b/sqlglot/dialects/hive.py
index d86691e4..22b70291 100644
--- a/sqlglot/dialects/hive.py
+++ b/sqlglot/dialects/hive.py
@@ -422,6 +422,15 @@ class Hive(Dialect):
                 super()._parse_order(skip_order_token=self._match(TokenType.SORT_BY)),
             )
 
+        def _parse_parameter(self) -> exp.Parameter:
+            self._match(TokenType.L_BRACE)
+            this = self._parse_identifier() or self._parse_primary_or_var()
+            expression = self._match(TokenType.COLON) and (
+                self._parse_identifier() or self._parse_primary_or_var()
+            )
+            self._match(TokenType.R_BRACE)
+            return self.expression(exp.Parameter, this=this, expression=expression)
+
     class Generator(generator.Generator):
         LIMIT_FETCH = "LIMIT"
         TABLESAMPLE_WITH_METHOD = False
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 538cb9cb..1ae854ef 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -601,8 +601,8 @@ class Snowflake(Dialect):
                 file_format = None
                 pattern = None
 
-                self._match(TokenType.L_PAREN)
-                while self._curr and not self._match(TokenType.R_PAREN):
+                wrapped = self._match(TokenType.L_PAREN)
+                while self._curr and wrapped and not self._match(TokenType.R_PAREN):
                     if self._match_text_seq("FILE_FORMAT", "=>"):
                         file_format = self._parse_string() or super()._parse_table_parts(
                             is_db_reference=is_db_reference
@@ -682,14 +682,22 @@ class Snowflake(Dialect):
             return self.expression(exp.LocationProperty, this=self._parse_location_path())
 
         def _parse_file_location(self) -> t.Optional[exp.Expression]:
-            return self._parse_table_parts()
+            # Parse either a subquery or a staged file
+            return (
+                self._parse_select(table=True)
+                if self._match(TokenType.L_PAREN, advance=False)
+                else self._parse_table_parts()
+            )
 
         def _parse_location_path(self) -> exp.Var:
             parts = [self._advance_any(ignore_reserved=True)]
 
             # We avoid consuming a comma token because external tables like @foo and @bar
-            # can be joined in a query with a comma separator.
-            while self._is_connected() and not self._match(TokenType.COMMA, advance=False):
+            # can be joined in a query with a comma separator, as well as closing paren
+            # in case of subqueries
+            while self._is_connected() and not self._match_set(
+                (TokenType.COMMA, TokenType.R_PAREN), advance=False
+            ):
                 parts.append(self._advance_any(ignore_reserved=True))
 
             return exp.var("".join(part.text for part in parts if part))
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 78defc6d..7bf7bcc8 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -1948,6 +1948,7 @@ class Drop(Expression):
         "cascade": False,
         "constraints": False,
         "purge": False,
+        "cluster": False,
     }
 
 
@@ -4189,6 +4190,7 @@ class AlterTable(Expression):
         "exists": False,
         "only": False,
         "options": False,
+        "cluster": False,
     }
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index ab14381e..77a0188b 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1211,12 +1211,14 @@ class Generator(metaclass=_Generator):
         expressions = f" ({expressions})" if expressions else ""
         kind = expression.args["kind"]
         exists_sql = " IF EXISTS " if expression.args.get("exists") else " "
+        on_cluster = self.sql(expression, "cluster")
+        on_cluster = f" {on_cluster}" if on_cluster else ""
         temporary = " TEMPORARY" if expression.args.get("temporary") else ""
         materialized = " MATERIALIZED" if expression.args.get("materialized") else ""
         cascade = " CASCADE" if expression.args.get("cascade") else ""
         constraints = " CONSTRAINTS" if expression.args.get("constraints") else ""
         purge = " PURGE" if expression.args.get("purge") else ""
-        return f"DROP{temporary}{materialized} {kind}{exists_sql}{this}{expressions}{cascade}{constraints}{purge}"
+        return f"DROP{temporary}{materialized} {kind}{exists_sql}{this}{on_cluster}{expressions}{cascade}{constraints}{purge}"
 
     def except_sql(self, expression: exp.Except) -> str:
         return self.set_operations(expression)
@@ -3023,10 +3025,12 @@ class Generator(metaclass=_Generator):
             actions = self.expressions(expression, key="actions", flat=True)
 
         exists = " IF EXISTS" if expression.args.get("exists") else ""
+        on_cluster = self.sql(expression, "cluster")
+        on_cluster = f" {on_cluster}" if on_cluster else ""
         only = " ONLY" if expression.args.get("only") else ""
         options = self.expressions(expression, key="options")
         options = f", {options}" if options else ""
-        return f"ALTER TABLE{exists}{only} {self.sql(expression, 'this')} {actions}{options}"
+        return f"ALTER TABLE{exists}{only} {self.sql(expression, 'this')}{on_cluster} {actions}{options}"
 
     def add_column_sql(self, expression: exp.AlterTable) -> str:
         if self.ALTER_TABLE_INCLUDE_COLUMN_KEYWORD:
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index cd089296..e7ee89a6 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -1510,6 +1510,8 @@ class Parser(metaclass=_Parser):
             schema=True, is_db_reference=self._prev.token_type == TokenType.SCHEMA
         )
 
+        cluster = self._parse_on_property() if self._match(TokenType.ON) else None
+
         if self._match(TokenType.L_PAREN, advance=False):
             expressions = self._parse_wrapped_csv(self._parse_types)
         else:
@@ -1527,6 +1529,7 @@ class Parser(metaclass=_Parser):
             cascade=self._match_text_seq("CASCADE"),
             constraints=self._match_text_seq("CONSTRAINTS"),
             purge=self._match_text_seq("PURGE"),
+            cluster=cluster,
         )
 
     def _parse_exists(self, not_: bool = False) -> t.Optional[bool]:
@@ -4345,9 +4348,7 @@ class Parser(metaclass=_Parser):
             if not this and self._match(TokenType.R_PAREN, advance=False):
                 this = self.expression(exp.Tuple)
             elif isinstance(this, exp.UNWRAPPED_QUERIES):
-                this = self._parse_set_operations(
-                    self._parse_subquery(this=this, parse_alias=False)
-                )
+                this = self._parse_subquery(this=this, parse_alias=False)
             elif isinstance(this, exp.Subquery):
                 this = self._parse_subquery(
                     this=self._parse_set_operations(this), parse_alias=False
@@ -5644,13 +5645,8 @@ class Parser(metaclass=_Parser):
         return self._parse_placeholder()
 
     def _parse_parameter(self) -> exp.Parameter:
-        self._match(TokenType.L_BRACE)
         this = self._parse_identifier() or self._parse_primary_or_var()
-        expression = self._match(TokenType.COLON) and (
-            self._parse_identifier() or self._parse_primary_or_var()
-        )
-        self._match(TokenType.R_BRACE)
-        return self.expression(exp.Parameter, this=this, expression=expression)
+        return self.expression(exp.Parameter, this=this)
 
     def _parse_placeholder(self) -> t.Optional[exp.Expression]:
         if self._match_set(self.PLACEHOLDER_PARSERS):
@@ -5912,6 +5908,7 @@ class Parser(metaclass=_Parser):
         exists = self._parse_exists()
         only = self._match_text_seq("ONLY")
         this = self._parse_table(schema=True)
+        cluster = self._parse_on_property() if self._match(TokenType.ON) else None
 
         if self._next:
             self._advance()
@@ -5929,6 +5926,7 @@ class Parser(metaclass=_Parser):
                     actions=actions,
                     only=only,
                     options=options,
+                    cluster=cluster,
                 )
 
         return self._parse_as_command(start)
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 989477fc..96fc433c 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -1187,6 +1187,8 @@ class Tokenizer(metaclass=_Tokenizer):
             if self._peek.isdigit():
                 self._advance()
             elif self._peek == "." and not decimal:
+                if self.tokens and self.tokens[-1].token_type == TokenType.PARAMETER:
+                    return self._add(TokenType.NUMBER)
                 decimal = True
                 self._advance()
             elif self._peek in ("-", "+") and scientific == 1:
diff --git a/sqlglotrs/src/tokenizer.rs b/sqlglotrs/src/tokenizer.rs
index cdc61d70..e79d0e90 100644
--- a/sqlglotrs/src/tokenizer.rs
+++ b/sqlglotrs/src/tokenizer.rs
@@ -483,6 +483,9 @@ impl<'a> TokenizerState<'a> {
             if self.peek_char.is_digit(10) {
                 self.advance(1)?;
             } else if self.peek_char == '.' && !decimal {
+                if self.tokens.last().map(|t| t.token_type) == Some(self.token_types.parameter) {
+                    return self.add(self.token_types.number, None);
+                }
                 decimal = true;
                 self.advance(1)?;
             } else if (self.peek_char == '-' || self.peek_char == '+') && scientific == 1 {


Test patch:
diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
index 89fd0e87..afdc540f 100644
--- a/tests/dialects/test_bigquery.py
+++ b/tests/dialects/test_bigquery.py
@@ -626,11 +626,15 @@ class TestBigQuery(Validator):
         )
         self.validate_all(
             "MD5(x)",
+            read={
+                "clickhouse": "MD5(x)",
+            },
             write={
                 "": "MD5_DIGEST(x)",
                 "bigquery": "MD5(x)",
                 "hive": "UNHEX(MD5(x))",
                 "spark": "UNHEX(MD5(x))",
+                "clickhouse": "MD5(x)",
             },
         )
         self.validate_all(
@@ -643,18 +647,45 @@ class TestBigQuery(Validator):
                 "": "SELECT MD5(some_string)",
                 "bigquery": "SELECT TO_HEX(MD5(some_string))",
                 "duckdb": "SELECT MD5(some_string)",
+                "clickhouse": "SELECT LOWER(HEX(MD5(some_string)))",
+            },
+        )
+        self.validate_all(
+            "SHA1(x)",
+            read={
+                "clickhouse": "SHA1(x)",
+            },
+            write={
+                "clickhouse": "SHA1(x)",
+                "bigquery": "SHA1(x)",
+                "": "SHA(x)",
+            },
+        )
+        self.validate_all(
+            "SHA1(x)",
+            write={
+                "bigquery": "SHA1(x)",
+                "": "SHA(x)",
             },
         )
         self.validate_all(
             "SHA256(x)",
+            read={
+                "clickhouse": "SHA256(x)",
+            },
             write={
                 "bigquery": "SHA256(x)",
                 "spark2": "SHA2(x, 256)",
+                "clickhouse": "SHA256(x)",
             },
         )
         self.validate_all(
             "SHA512(x)",
+            read={
+                "clickhouse": "SHA512(x)",
+            },
             write={
+                "clickhouse": "SHA512(x)",
                 "bigquery": "SHA512(x)",
                 "spark2": "SHA2(x, 512)",
             },
diff --git a/tests/dialects/test_clickhouse.py b/tests/dialects/test_clickhouse.py
index d206018a..1bce7f5c 100644
--- a/tests/dialects/test_clickhouse.py
+++ b/tests/dialects/test_clickhouse.py
@@ -438,6 +438,7 @@ class TestClickhouse(Validator):
             "ALTER TABLE visits REPLACE PARTITION tuple(toYYYYMM(toDate('2019-01-25'))) FROM visits_tmp"
         )
         self.validate_identity("ALTER TABLE visits REPLACE PARTITION ID '201901' FROM visits_tmp")
+        self.validate_identity("ALTER TABLE visits ON CLUSTER test_cluster DROP COLUMN col1")
 
     def test_cte(self):
         self.validate_identity("WITH 'x' AS foo SELECT foo")
@@ -870,3 +871,8 @@ LIFETIME(MIN 0 MAX 0)""",
         )
 
         parse_one("foobar(x)").assert_is(exp.Anonymous)
+
+    def test_drop_on_cluster(self):
+        for creatable in ("DATABASE", "TABLE", "VIEW", "DICTIONARY", "FUNCTION"):
+            with self.subTest(f"Test DROP {creatable} ON CLUSTER"):
+                self.validate_identity(f"DROP {creatable} test ON CLUSTER test_cluster")
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index f05bd66c..e8f827a9 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -894,6 +894,10 @@ WHERE
         self.validate_identity("SELECT * FROM @namespace.%table_name/path/to/file.json.gz")
         self.validate_identity("SELECT * FROM '@external/location' (FILE_FORMAT => 'path.to.csv')")
         self.validate_identity("PUT file:///dir/tmp.csv @%table", check_command_warning=True)
+        self.validate_identity("SELECT * FROM (SELECT a FROM @foo)")
+        self.validate_identity(
+            "SELECT * FROM (SELECT * FROM '@external/location' (FILE_FORMAT => 'path.to.csv'))"
+        )
         self.validate_identity(
             "SELECT * FROM @foo/bar (FILE_FORMAT => ds_sandbox.test.my_csv_format, PATTERN => 'test') AS bla"
         )
@@ -1856,10 +1860,7 @@ STORAGE_ALLOWED_LOCATIONS=('s3://mybucket1/path1/', 's3://mybucket2/path2/')""",
         self.assertEqual(expression.sql(dialect="snowflake"), "SELECT TRY_CAST(FOO() AS TEXT)")
 
     def test_copy(self):
-        self.validate_identity(
-            """COPY INTO mytable (col1, col2) FROM 's3://mybucket/data/files' FILES = ('file1', 'file2') PATTERN = 'pattern' file_format = (FORMAT_NAME = my_csv_format NULL_IF = ('str1', 'str2')) PARSE_HEADER = TRUE""",
-            """COPY INTO mytable (col1, col2) FROM 's3://mybucket/data/files' FILES = ('file1', 'file2') PATTERN = 'pattern' FILE_FORMAT = (FORMAT_NAME = my_csv_format NULL_IF = ('str1', 'str2')) PARSE_HEADER = TRUE""",
-        )
+        self.validate_identity("COPY INTO test (c1) FROM (SELECT $1.c1 FROM @mystage)")
         self.validate_identity(
             """COPY INTO temp FROM @random_stage/path/ FILE_FORMAT = (TYPE = CSV FIELD_DELIMITER = '|' NULL_IF = () FIELD_OPTIONALLY_ENCLOSED_BY = '"' TIMESTAMP_FORMAT = 'TZHTZM YYYY-MM-DD HH24:MI:SS.FF9' DATE_FORMAT = 'TZHTZM YYYY-MM-DD HH24:MI:SS.FF9' BINARY_FORMAT = BASE64) VALIDATION_MODE = 'RETURN_3_ROWS'"""
         )
@@ -1869,3 +1870,14 @@ STORAGE_ALLOWED_LOCATIONS=('s3://mybucket1/path1/', 's3://mybucket2/path2/')""",
         self.validate_identity(
             """COPY INTO mytable FROM 'azure://myaccount.blob.core.windows.net/mycontainer/data/files' CREDENTIALS = (AZURE_SAS_TOKEN = 'token') ENCRYPTION = (TYPE = 'AZURE_CSE' MASTER_KEY = 'kPx...') FILE_FORMAT = (FORMAT_NAME = my_csv_format)"""
         )
+        self.validate_identity(
+            """COPY INTO mytable (col1, col2) FROM 's3://mybucket/data/files' FILES = ('file1', 'file2') PATTERN = 'pattern' FILE_FORMAT = (FORMAT_NAME = my_csv_format NULL_IF = ('str1', 'str2')) PARSE_HEADER = TRUE"""
+        )
+
+    def test_querying_semi_structured_data(self):
+        self.validate_identity("SELECT $1")
+        self.validate_identity("SELECT $1.elem")
+
+        self.validate_identity("SELECT $1:a.b", "SELECT GET_PATH($1, 'a.b')")
+        self.validate_identity("SELECT t.$23:a.b", "SELECT GET_PATH(t.$23, 'a.b')")
+        self.validate_identity("SELECT t.$17:a[0].b[0].c", "SELECT GET_PATH(t.$17, 'a[0].b[0].c')")