Fix(snowflake): clean up PUT implementation (#4830)

Author: georgesittas

sqlglot/dialects/snowflake.py

@@ -810,10 +810,13 @@ def _parse_show_snowflake(self, this: str) -> exp.Show:
         def _parse_put(self) -> exp.Put | exp.Command:
             if self._curr.token_type != TokenType.STRING:
                 return self._parse_as_command(self._prev)
-            source = self._parse_string()
-            target = self._parse_location_path()
-            props = self._parse_properties()
-            return self.expression(exp.Put, this=source, target=target, properties=props)
+
+            return self.expression(
+                exp.Put,
+                this=self._parse_string(),
+                target=self._parse_location_path(),
+                properties=self._parse_properties(),
+            )
 
         def _parse_location_property(self) -> exp.LocationProperty:
             self._match(TokenType.EQ)
@@ -870,6 +873,7 @@ class Tokenizer(tokens.Tokenizer):
 
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
+            "FILE://": TokenType.URI_START,
             "BYTEINT": TokenType.INT,
             "CHAR VARYING": TokenType.VARCHAR,
             "CHARACTER VARYING": TokenType.VARCHAR,

sqlglot/tokens.py

@@ -80,6 +80,8 @@ class TokenType(AutoName):
     XOR = auto()
     DSTAR = auto()
 
+    URI_START = auto()
+
     BLOCK_START = auto()
     BLOCK_END = auto()
 

tests/dialects/test_snowflake.py

@@ -106,6 +106,10 @@ def test_snowflake(self):
         self.validate_identity(
             """SELECT TO_TIMESTAMP('2025-01-16T14:45:30.123+0500', 'yyyy-mm-DD"T"hh24:mi:ss.ff3TZHTZM')"""
         )
+        self.validate_identity(
+            "SELECT 1 put",
+            "SELECT 1 AS put",
+        )
         self.validate_identity(
             "WITH t (SELECT 1 AS c) SELECT c FROM t",
             "WITH t AS (SELECT 1 AS c) SELECT c FROM t",
@@ -2390,10 +2394,12 @@ def test_put_to_stage(self):
 
         # validate identity for different args and properties
         self.validate_identity("PUT 'file:///dir/tmp.csv' @s1/test")
-        # TODO: the test below is still failing!
-        self.validate_identity("PUT file:///dir/tmp.csv @%table")
+
+        # the unquoted URI variant is not fully supported yet
+        self.validate_identity("PUT file:///dir/tmp.csv @%table", check_command_warning=True)
         self.validate_identity(
-            "PUT 'file:///dir/tmp.csv' @s1/test PARALLEL=1 AUTO_COMPRESS=FALSE source_compression=gzip OVERWRITE=TRUE"
+            "PUT file:///dir/tmp.csv @s1/test PARALLEL=1 AUTO_COMPRESS=FALSE source_compression=gzip OVERWRITE=TRUE",
+            check_command_warning=True,
         )
 
     def test_querying_semi_structured_data(self):

tests/fixtures/identity.sql

@@ -889,4 +889,3 @@ CAST(x AS INT128)
 CAST(x AS UINT128)
 CAST(x AS UINT256)
 SELECT export
-SELECT 1 put

tests/test_tokens.py

@@ -1,7 +1,6 @@
 import unittest
 
 from sqlglot.dialects import BigQuery
-from sqlglot.dialects.snowflake import Snowflake
 from sqlglot.errors import TokenError
 from sqlglot.tokens import Tokenizer, TokenType
 
@@ -177,25 +176,3 @@ def test_jinja(self):
                 (TokenType.STRING, ") }}"),
             ],
         )
-
-    def test_tokenize_put_with_raw_url(self):
-        tokens = Snowflake.Tokenizer().tokenize("PUT file:///tmp.csv @%table")
-        tokens = [(token.token_type, token.text) for token in tokens]
-
-        self.assertEqual(
-            tokens,
-            [
-                (TokenType.VAR, "PUT"),
-                (TokenType.VAR, "file"),
-                (TokenType.COLON, ":"),
-                (TokenType.SLASH, "/"),
-                (TokenType.SLASH, "/"),
-                (TokenType.SLASH, "/"),
-                (TokenType.VAR, "tmp"),
-                (TokenType.DOT, "."),
-                (TokenType.VAR, "csv"),
-                (TokenType.PARAMETER, "@"),
-                (TokenType.MOD, "%"),
-                (TokenType.TABLE, "table"),
-            ],
-        )