Repository: tobymao/sqlglot
Pull request: fix(snowflake): COPY Subquery postfix
Fixes #3434

This PR addresses the following issues:
- COPY would only allow files in the `FROM` clause, extend it to allow subqueries as well

- When parsing staged files such as `@foo` we'd also consume the following parentheses coming from (optional) staged file options such as `SELECT * FROM @foo (FILE_FORMAT => ...)`. However, if there are no options specified we would end up consuming _any_ following `R_PAREN` as well e.g. `SELECT * FROM (SELECT * FROM @foo)`

- For transforming semi-structured data Snowflake specifies the subquery columns as `[<alias>.]$<file_col_num>[.<element>]`. However, the tokenizer always appends the `TokenType.DOT` to following a number e.g. `$1.elem -> $1.e lem`. To solve this issue, the tokenizer will only parse an integer (and not a real) if the number is prefixed by a token of type `TokenType.PARAMETER`

- The base parser parses `exp.Parameter` as `'[' <this> ':' <expression> ']'` but only generates `this`. Snowflake parsers `$<num>` as `exp.Parameter` which can optionally be followed by a path & element name such as `$1:a.c`. For this reason, `exp.Parameter` now generates `<this> [':' <expression>]` 

Docs
-----
- [Snowflake COPY](https://docs.snowflake.com/en/sql-reference/sql/copy-into-table)
- [Snowflake COPY Transform Subquery specs](https://docs.snowflake.com/en/sql-reference/sql/copy-into-table#transformation-parameters)
- [Snowflake Querying semi-structured data](https://docs.snowflake.com/en/user-guide/querying-semistructured#traversing-semi-structured-data)

Problem statement:
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!

Hints from issue discussion:

Commit message: fix(snowflake): COPY Subquery postfix

Patch to review:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 538cb9cb4a..9a28488570 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -601,8 +601,8 @@ def _parse_table_parts(
                 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,23 @@ def _parse_location_property(self) -> exp.LocationProperty:
             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_primary()
+                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):
+            while (
+                self._is_connected()
+                and not self._match(TokenType.COMMA, advance=False)
+                and not self._match(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/generator.py b/sqlglot/generator.py
index ab14381eb3..ebcab6c070 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -2310,7 +2310,10 @@ def star_sql(self, expression: exp.Star) -> str:
 
     def parameter_sql(self, expression: exp.Parameter) -> str:
         this = self.sql(expression, "this")
-        return f"{self.PARAMETER_TOKEN}{this}"
+        expr = self.sql(expression, "expression")
+        expr = f":{expr}" if expr else ""
+
+        return f"{self.PARAMETER_TOKEN}{this}{expr}"
 
     def sessionparameter_sql(self, expression: exp.SessionParameter) -> str:
         this = self.sql(expression, "this")
diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
index 989477fc44..96fc433c28 100644
--- a/sqlglot/tokens.py
+++ b/sqlglot/tokens.py
@@ -1187,6 +1187,8 @@ def _scan_number(self) -> None:
             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 cdc61d7084..e79d0e900f 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 {
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index f05bd66cf1..f3285d8d10 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -894,6 +894,10 @@ def test_staged_files(self):
         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 @@ def test_try_cast(self):
         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 @@ def test_copy(self):
         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""",
+            """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.a")
+        self.validate_identity("SELECT $1:a.b")
+        self.validate_identity("SELECT t.$23:a.b")
+        self.validate_identity("SELECT tbl.$17:a[0].b[0].c")


Code review comments:
diff_hunk: @@ -682,14 +682,23 @@ def _parse_location_property(self) -> exp.LocationProperty:
             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_primary()
path: sqlglot/dialects/snowflake.py
text: @user1:
maybe it should be self._parse_select(table=True)

diff_hunk: @@ -601,8 +601,8 @@ def _parse_table_parts(
                 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)):
path: sqlglot/dialects/snowflake.py
text: @user1:
i don't think you need the parenthesis

self._curr and wrappend and not self._match right?

@author:
Right, shouldn't make a difference. I usually wrap such paired bools to put a mental (and future) note when re-reading these chained conditions. Do we want to preserve the most simplified version of branches in general?

@user1:
yea let's simplify