feat(snowflake)!: add support for GET statements (#5019)

* feat(snowflake): add support for GET statements

* types in helper method

* address review comments

* refactor

Author: eruditmorina

sqlglot/dialects/snowflake.py

@@ -512,6 +512,7 @@ class Parser(parser.Parser):
 
         STATEMENT_PARSERS = {
             **parser.Parser.STATEMENT_PARSERS,
+            TokenType.GET: lambda self: self._parse_get(),
             TokenType.PUT: lambda self: self._parse_put(),
             TokenType.SHOW: lambda self: self._parse_show(),
         }
@@ -857,6 +858,21 @@ def _parse_put(self) -> exp.Put | exp.Command:
                 properties=self._parse_properties(),
             )
 
+        def _parse_get(self) -> exp.Get | exp.Command:
+            start = self._prev
+            target = self._parse_location_path()
+
+            # Parse as command if unquoted file path
+            if self._curr.token_type == TokenType.URI_START:
+                return self._parse_as_command(start)
+
+            return self.expression(
+                exp.Get,
+                this=self._parse_string(),
+                target=target,
+                properties=self._parse_properties(),
+            )
+
         def _parse_location_property(self) -> exp.LocationProperty:
             self._match(TokenType.EQ)
             return self.expression(exp.LocationProperty, this=self._parse_location_path())
@@ -937,6 +953,7 @@ class Tokenizer(tokens.Tokenizer):
             "CHARACTER VARYING": TokenType.VARCHAR,
             "EXCLUDE": TokenType.EXCEPT,
             "FILE FORMAT": TokenType.FILE_FORMAT,
+            "GET": TokenType.GET,
             "ILIKE ANY": TokenType.ILIKE_ANY,
             "LIKE ANY": TokenType.LIKE_ANY,
             "MATCH_CONDITION": TokenType.MATCH_CONDITION,

sqlglot/expressions.py

@@ -3307,6 +3307,11 @@ class Put(Expression):
     arg_types = {"this": True, "target": True, "properties": False}
 
 
+# https://docs.snowflake.com/en/sql-reference/sql/get
+class Get(Expression):
+    arg_types = {"this": True, "target": True, "properties": False}
+
+
 class Table(Expression):
     arg_types = {
         "this": False,

sqlglot/generator.py

@@ -146,6 +146,7 @@ class Generator(metaclass=_Generator):
         exp.Except: lambda self, e: self.set_operations(e),
         exp.ExternalProperty: lambda *_: "EXTERNAL",
         exp.Floor: lambda self, e: self.ceil_floor(e),
+        exp.Get: lambda self, e: self.get_put_sql(e),
         exp.GlobalProperty: lambda *_: "GLOBAL",
         exp.HeapProperty: lambda *_: "HEAP",
         exp.IcebergProperty: lambda *_: "ICEBERG",
@@ -175,6 +176,7 @@ class Generator(metaclass=_Generator):
         exp.PivotAny: lambda self, e: f"ANY{self.sql(e, 'this')}",
         exp.ProjectionPolicyColumnConstraint: lambda self,
         e: f"PROJECTION POLICY {self.sql(e, 'this')}",
+        exp.Put: lambda self, e: self.get_put_sql(e),
         exp.RemoteWithConnectionModelProperty: lambda self,
         e: f"REMOTE WITH CONNECTION {self.sql(e, 'this')}",
         exp.ReturnsProperty: lambda self, e: (
@@ -4903,9 +4905,16 @@ def show_sql(self, expression: exp.Show) -> str:
         self.unsupported("Unsupported SHOW statement")
         return ""
 
-    def put_sql(self, expression: exp.Put) -> str:
+    def get_put_sql(self, expression: exp.Put | exp.Get) -> str:
+        # Snowflake GET/PUT statements:
+        #   PUT <file> <internalStage> <properties>
+        #   GET <internalStage> <file> <properties>
         props = expression.args.get("properties")
         props_sql = self.properties(props, prefix=" ", sep=" ", wrapped=False) if props else ""
         this = self.sql(expression, "this")
         target = self.sql(expression, "target")
-        return f"PUT {this} {target}{props_sql}"
+
+        if isinstance(expression, exp.Put):
+            return f"PUT {this} {target}{props_sql}"
+        else:
+            return f"GET {target} {this}{props_sql}"

sqlglot/parser.py

@@ -512,6 +512,7 @@ class Parser(metaclass=_Parser):
         TokenType.FINAL,
         TokenType.FORMAT,
         TokenType.FULL,
+        TokenType.GET,
         TokenType.IDENTIFIER,
         TokenType.IS,
         TokenType.ISNULL,

sqlglot/tokens.py

@@ -290,6 +290,7 @@ class TokenType(AutoName):
     FROM = auto()
     FULL = auto()
     FUNCTION = auto()
+    GET = auto()
     GLOB = auto()
     GLOBAL = auto()
     GRANT = auto()

tests/dialects/test_snowflake.py

@@ -110,6 +110,10 @@ def test_snowflake(self):
             "SELECT 1 put",
             "SELECT 1 AS put",
         )
+        self.validate_identity(
+            "SELECT 1 get",
+            "SELECT 1 AS get",
+        )
         self.validate_identity(
             "WITH t (SELECT 1 AS c) SELECT c FROM t",
             "WITH t AS (SELECT 1 AS c) SELECT c FROM t",
@@ -2465,6 +2469,33 @@ def test_put_to_stage(self):
             check_command_warning=True,
         )
 
+    def test_get_from_stage(self):
+        self.validate_identity('GET @"my_DB"."schEMA1"."MYstage" \'file:///dir/tmp.csv\'')
+        self.validate_identity("GET @s1/test 'file:///dir/tmp.csv'")
+
+        # GET with file path and stage ref containing spaces (wrapped in single quotes)
+        ast = parse_one("GET '@s1/my folder' 'file://my file.txt'", read="snowflake")
+        self.assertIsInstance(ast, exp.Get)
+        self.assertEqual(ast.args["target"], exp.Var(this="'@s1/my folder'"))
+        self.assertEqual(ast.this, exp.Literal(this="file://my file.txt", is_string=True))
+        self.assertEqual(ast.sql("snowflake"), "GET '@s1/my folder' 'file://my file.txt'")
+
+        # expression with additional properties
+        ast = parse_one("GET @stage1/folder 'file:///tmp/my.txt' PARALLEL = 1", read="snowflake")
+        self.assertIsInstance(ast, exp.Get)
+        self.assertEqual(ast.args["target"], exp.Var(this="@stage1/folder"))
+        self.assertEqual(ast.this, exp.Literal(this="file:///tmp/my.txt", is_string=True))
+        properties = ast.args.get("properties")
+        props_dict = {prop.this.this: prop.args["value"].this for prop in properties.expressions}
+        self.assertEqual(props_dict, {"PARALLEL": "1"})
+
+        # the unquoted URI variant is not fully supported yet
+        self.validate_identity("GET @%table file:///dir/tmp.csv", check_command_warning=True)
+        self.validate_identity(
+            "GET @s1/test file:///dir/tmp.csv PARALLEL=1",
+            check_command_warning=True,
+        )
+
     def test_querying_semi_structured_data(self):
         self.validate_identity("SELECT $1")
         self.validate_identity("SELECT $1.elem")