feat(snowflake): parse CREATE SEQUENCE (#3072)

* . * . * . * .
tobymao · Mar 5, 2024 · 46c9c2c · 46c9c2c
1 parent 6c67a2b
commit 46c9c2c
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 1 deletion.
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
@@ -670,6 +670,7 @@ class Tokenizer(tokens.Tokenizer):
             "RENAME": TokenType.REPLACE,
             "RM": TokenType.COMMAND,
             "SAMPLE": TokenType.TABLE_SAMPLE,
+            "SEQUENCE": TokenType.SEQUENCE,
             "SQL_DOUBLE": TokenType.DOUBLE,
             "SQL_VARCHAR": TokenType.VARCHAR,
             "STORAGE INTEGRATION": TokenType.STORAGE_INTEGRATION,

diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -1214,6 +1214,8 @@ class Create(DDL):
         "begin": False,
         "end": False,
         "clone": False,
+        "start": False,
+        "increment": False,
     }
 
     @property
@@ -3620,6 +3622,14 @@ def output_name(self) -> str:
         return self.name
 
 
+class Increment(Expression):
+    arg_types = {"this": True}
+
+
+class Start(Expression):
+    arg_types = {"this": True}
+
+
 class Parameter(Condition):
     arg_types = {"this": True, "expression": False}
 

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -972,7 +972,13 @@ def create_sql(self, expression: exp.Create) -> str:
         clone = self.sql(expression, "clone")
         clone = f" {clone}" if clone else ""
 
-        expression_sql = f"CREATE{modifiers} {kind}{exists_sql} {this}{properties_sql}{expression_sql}{postexpression_props_sql}{index_sql}{no_schema_binding}{clone}"
+        start = self.sql(expression, "start")
+        start = f" {start}" if start else ""
+
+        increment = self.sql(expression, "increment")
+        increment = f" {increment}" if increment else ""
+
+        expression_sql = f"CREATE{modifiers} {kind}{exists_sql} {this}{properties_sql}{expression_sql}{postexpression_props_sql}{index_sql}{no_schema_binding}{clone}{start}{increment}"
         return self.prepend_ctes(expression, expression_sql)
 
     def clone_sql(self, expression: exp.Clone) -> str:
@@ -1560,6 +1566,14 @@ def historicaldata_sql(self, expression: exp.HistoricalData) -> str:
         expr = self.sql(expression, "expression")
         return f"{this} ({kind} => {expr})"
 
+    def start_sql(self, expression: exp.Start) -> str:
+        this = self.sql(expression, "this")
+        return f"START WITH {this}"
+
+    def increment_sql(self, expression: exp.Increment) -> str:
+        this = self.sql(expression, "this")
+        return f"INCREMENT BY {this}"
+
     def table_parts(self, expression: exp.Table) -> str:
         return ".".join(
             self.sql(part)

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -291,6 +291,7 @@ class Parser(metaclass=_Parser):
         TokenType.VIEW,
         TokenType.MODEL,
         TokenType.DICTIONARY,
+        TokenType.SEQUENCE,
         TokenType.STORAGE_INTEGRATION,
     }
 
@@ -1438,6 +1439,8 @@ def _parse_create(self) -> exp.Create | exp.Command:
         begin = None
         end = None
         clone = None
+        seq_start = None
+        seq_increment = None
 
         def extend_props(temp_props: t.Optional[exp.Properties]) -> None:
             nonlocal properties
@@ -1515,6 +1518,12 @@ def extend_props(temp_props: t.Optional[exp.Properties]) -> None:
             elif create_token.token_type == TokenType.VIEW:
                 if self._match_text_seq("WITH", "NO", "SCHEMA", "BINDING"):
                     no_schema_binding = True
+            elif create_token.token_type == TokenType.SEQUENCE:
+                if self._match_texts("START") or self._match(TokenType.START_WITH):
+                    seq_start = self.expression(exp.Start, this=self._parse_number())
+                if self._match_texts("INCREMENT"):
+                    self._match_texts("BY")
+                    seq_increment = self.expression(exp.Increment, this=self._parse_number())
 
             shallow = self._match_text_seq("SHALLOW")
 
@@ -1542,6 +1551,8 @@ def extend_props(temp_props: t.Optional[exp.Properties]) -> None:
             begin=begin,
             end=end,
             clone=clone,
+            start=seq_start,
+            increment=seq_increment,
         )
 
     def _parse_property_before(self) -> t.Optional[exp.Expression]:

diff --git a/sqlglot/tokens.py b/sqlglot/tokens.py
@@ -340,6 +340,7 @@ class TokenType(AutoName):
     SELECT = auto()
     SEMI = auto()
     SEPARATOR = auto()
+    SEQUENCE = auto()
     SERDE_PROPERTIES = auto()
     SET = auto()
     SETTINGS = auto()

diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
@@ -1099,6 +1099,15 @@ def test_ddl(self):
             write={"snowflake": "CREATE TABLE a (b INT)"},
         )
 
+        self.validate_identity("CREATE SEQUENCE seq1 START WITH 1 INCREMENT BY 2")
+        self.assertIsInstance(
+            parse_one("CREATE SEQUENCE seq1 START 1 INCREMENT 2", read="snowflake"), exp.Create
+        )
+        self.assertIsInstance(
+            parse_one("CREATE SEQUENCE seq1 WITH START WITH 1 INCREMENT BY 2", read="snowflake"),
+            exp.Create,
+        )
+
     def test_user_defined_functions(self):
         self.validate_all(
             "CREATE FUNCTION a(x DATE, y BIGINT) RETURNS ARRAY LANGUAGE JAVASCRIPT AS $$ SELECT 1 $$",