Repository: tobymao/sqlglot
Pull request: Fix parsing for snowflake AUTOINCREMENT/IDENTITY constraints
Fixes #1129 

Problem statement:
snowflake: cannot create or alter tables with AUTOINCREMENT or IDENTITY keyword
I can use the "IDENTITY" and "AUTOINCREMENT" keywords in snowflake:

```sql
create or replace table colors as
    select name
    from (values ('blue'),('red'),('green')) colors (name);
create or replace table identity_column_example like colors;
alter table identity_column_example add column id int identity(1,1);
```

```sql
CREATE OR REPLACE TEMPORARY TABLE "X" ("IDENTITY_COLUMN" NUMBER AUTOINCREMENT(0, 1))
```

In sqlglot 11.0.0, python 3.10.4 I get ParseError:

```python
import sqlglot

sqlglot.transpile('CREATE OR REPLACE TEMPORARY TABLE "X" ("IDENTITY_COLUMN" NUMBER AUTOINCREMENT(0, 1))', read='snowflake', write='snowflake')
sqlglot.transpile('alter table identity_column_example add column id int identity(1,1);')
```

<details>

<summary>stack trace for first error in example</summary>

```python-traceback
ParseError                                Traceback (most recent call last)
Cell In[4], line 1
----> 1 sqlglot.transpile('CREATE OR REPLACE TEMPORARY TABLE "X" ("IDENTITY_COLUMN" NUMBER AUTOINCREMENT(0, 1))', read='snowflake', write='snowflake')

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/__init__.py:177, in transpile(sql, read, write, identity, error_level, **opts)
    158 """
    159 Parses the given SQL string in accordance with the source dialect and returns a list of SQL strings transformed
    160 to conform to the target dialect. Each string in the returned list represents a single transformed SQL statement.
   (...)
    172     The list of transpiled SQL statements.
    173 """
    174 write = write or read if identity else write
    175 return [
    176     Dialect.get_or_raise(write)().generate(expression, **opts)
--> 177     for expression in parse(sql, read, error_level=error_level)
    178 ]

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/__init__.py:65, in parse(sql, read, **opts)
     53 """
     54 Parses the given SQL string into a collection of syntax trees, one per parsed SQL statement.
     55
   (...)
     62     The resulting syntax tree collection.
     63 """
     64 dialect = Dialect.get_or_raise(read)()
---> 65 return dialect.parse(sql, **opts)

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/dialects/dialect.py:163, in Dialect.parse(self, sql, **opts)
    162 def parse(self, sql: str, **opts) -> t.List[t.Optional[exp.Expression]]:
--> 163     return self.parser(**opts).parse(self.tokenizer.tokenize(sql), sql)

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:707, in Parser.parse(self, raw_tokens, sql)
    693 def parse(
    694     self, raw_tokens: t.List[Token], sql: t.Optional[str] = None
    695 ) -> t.List[t.Optional[exp.Expression]]:
    696     """
    697     Parses a list of tokens and returns a list of syntax trees, one tree
    698     per parsed SQL statement.
   (...)
    705         The list of syntax trees.
    706     """
--> 707     return self._parse(
    708         parse_method=self.__class__._parse_statement, raw_tokens=raw_tokens, sql=sql
    709     )

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:770, in Parser._parse(self, parse_method, raw_tokens, sql)
    767 self._tokens = tokens
    768 self._advance()
--> 770 expressions.append(parse_method(self))
    772 if self._index < len(self._tokens):
    773     self.raise_error("Invalid expression / Unexpected token")

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:900, in Parser._parse_statement(self)
    897     return None
    899 if self._match_set(self.STATEMENT_PARSERS):
--> 900     return self.STATEMENT_PARSERS[self._prev.token_type](self)
    902 if self._match_set(Tokenizer.COMMANDS):
    903     return self._parse_command()

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:431, in Parser.<lambda>(self)
    365 COLUMN_OPERATORS = {
    366     TokenType.DOT: None,
    367     TokenType.DCOLON: lambda self, this, to: self.expression(
   (...)
    396     ),
    397 }
    399 EXPRESSION_PARSERS = {
    400     exp.Column: lambda self: self._parse_column(),
    401     exp.DataType: lambda self: self._parse_types(),
   (...)
    423     "JOIN_TYPE": lambda self: self._parse_join_side_and_kind(),
    424 }
    426 STATEMENT_PARSERS = {
    427     TokenType.ALTER: lambda self: self._parse_alter(),
    428     TokenType.BEGIN: lambda self: self._parse_transaction(),
    429     TokenType.CACHE: lambda self: self._parse_cache(),
    430     TokenType.COMMIT: lambda self: self._parse_commit_or_rollback(),
--> 431     TokenType.CREATE: lambda self: self._parse_create(),
    432     TokenType.DELETE: lambda self: self._parse_delete(),
    433     TokenType.DESC: lambda self: self._parse_describe(),
    434     TokenType.DESCRIBE: lambda self: self._parse_describe(),
    435     TokenType.DROP: lambda self: self._parse_drop(),
    436     TokenType.END: lambda self: self._parse_commit_or_rollback(),
    437     TokenType.INSERT: lambda self: self._parse_insert(),
    438     TokenType.LOAD_DATA: lambda self: self._parse_load_data(),
    439     TokenType.MERGE: lambda self: self._parse_merge(),
    440     TokenType.ROLLBACK: lambda self: self._parse_commit_or_rollback(),
    441     TokenType.UNCACHE: lambda self: self._parse_uncache(),
    442     TokenType.UPDATE: lambda self: self._parse_update(),
    443     TokenType.USE: lambda self: self.expression(
    444         exp.Use,
    445         kind=self._match_texts(("ROLE", "WAREHOUSE", "DATABASE", "SCHEMA"))
    446         and exp.Var(this=self._prev.text),
    447         this=self._parse_table(schema=False),
    448     ),
    449 }
    451 UNARY_PARSERS = {
    452     TokenType.PLUS: lambda self: self._parse_unary(),  # Unary + is handled as a no-op
    453     TokenType.NOT: lambda self: self.expression(exp.Not, this=self._parse_equality()),
    454     TokenType.TILDA: lambda self: self.expression(exp.BitwiseNot, this=self._parse_unary()),
    455     TokenType.DASH: lambda self: self.expression(exp.Neg, this=self._parse_unary()),
    456 }
    458 PRIMARY_PARSERS = {
    459     TokenType.STRING: lambda self, token: self.expression(
    460         exp.Literal, this=token.text, is_string=True
   (...)
    477     TokenType.SESSION_PARAMETER: lambda self, _: self._parse_session_parameter(),
    478 }

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:997, in Parser._parse_create(self)
    994 if self._match(TokenType.COMMA):  # comma-separated properties before schema definition
    995     properties = self._parse_properties(before=True)
--> 997 this = self._parse_schema(this=table_parts)
    999 if not properties:  # properties after schema definition
   1000     properties = self._parse_properties()

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:2720, in Parser._parse_schema(self, this)
   2714     return this
   2716 args = self._parse_csv(
   2717     lambda: self._parse_constraint()
   2718     or self._parse_column_def(self._parse_field(any_token=True))
   2719 )
-> 2720 self._match_r_paren()
   2721 return self.expression(exp.Schema, this=this, expressions=args)

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:3635, in Parser._match_r_paren(self, expression)
   3633 def _match_r_paren(self, expression=None):
   3634     if not self._match(TokenType.R_PAREN):
-> 3635         self.raise_error("Expecting )")
   3636     if expression and self._prev_comments:
   3637         expression.comments = self._prev_comments

File ~/opt/anaconda3/envs/ponder-product-testing/lib/python3.10/site-packages/sqlglot/parser.py:816, in Parser.raise_error(self, message, token)
    804 error = ParseError.new(
    805     f"{message}. Line {token.line}, Col: {token.col}.\n"
    806     f"  {start_context}\033[4m{highlight}\033[0m{end_context}",
   (...)
    812     end_context=end_context,
    813 )
    815 if self.error_level == ErrorLevel.IMMEDIATE:
--> 816     raise error
    818 self.errors.append(error)

ParseError: Expecting ). Line 1, Col: 65.
  CREATE OR REPLACE TEMPORARY TABLE "X" ("IDENTITY_COLUMN" NUMBER AUTOINCREMENT(0, 1))
```

</details>

Hints from issue discussion:
I can take a look at this and #1132 tomorrow if nobody beats me to it.

Tip: if you want to use `transpile` with `read` and `write` pointing to the same dialect, you can omit the `write` argument because of `identity` being true by default for this method. Even shorter: you can just do `transpile(..., "snowflake")`.

Commit message: Fix parsing for snowflake AUTOINCREMENT/IDENTITY constraints

Patch to review:
diff --git a/sqlglot/dialects/snowflake.py b/sqlglot/dialects/snowflake.py
index 8ca1d362ea..7b74fb5eda 100644
--- a/sqlglot/dialects/snowflake.py
+++ b/sqlglot/dialects/snowflake.py
@@ -178,6 +178,24 @@ class Parser(parser.Parser):
             ),
         }
 
+        # https://docs.snowflake.com/en/sql-reference/sql/create-table.html
+        def _parse_autoincrement(self) -> exp.Expression:
+            start = None
+            increment = None
+
+            if self._match(TokenType.L_PAREN, advance=False):
+                args = self._parse_wrapped_csv(self._parse_bitwise)
+                start = seq_get(args, 0)
+                increment = seq_get(args, 1)
+            elif self._match_text_seq("START"):
+                start = self._parse_bitwise()
+                self._match_text_seq("INCREMENT")
+                increment = self._parse_bitwise()
+
+            if start and increment:
+                return exp.GeneratedAsIdentityColumnConstraint(start=start, increment=increment)
+            return exp.AutoIncrementColumnConstraint()
+
     class Tokenizer(tokens.Tokenizer):
         QUOTES = ["'", "$$"]
         STRING_ESCAPES = ["\\", "'"]
@@ -189,6 +207,7 @@ class Tokenizer(tokens.Tokenizer):
 
         KEYWORDS = {
             **tokens.Tokenizer.KEYWORDS,
+            "AUTOINCREMENT": TokenType.AUTO_INCREMENT,
             "EXCLUDE": TokenType.EXCEPT,
             "MATCH_RECOGNIZE": TokenType.MATCH_RECOGNIZE,
             "RENAME": TokenType.REPLACE,
@@ -295,3 +314,12 @@ def describe_sql(self, expression: exp.Describe) -> str:
             kind = f" {kind_value}" if kind_value else ""
             this = f" {self.sql(expression, 'this')}"
             return f"DESCRIBE{kind}{this}"
+
+        def generatedasidentitycolumnconstraint_sql(
+            self, expression: exp.GeneratedAsIdentityColumnConstraint
+        ) -> str:
+            start = expression.args.get("start")
+            start = f" START {start}" if start else ""
+            increment = expression.args.get("increment")
+            increment = f" INCREMENT {increment}" if increment else ""
+            return f"AUTOINCREMENT{start}{increment}"
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
index 38764fa2e0..ba09c8b2e3 100644
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -939,7 +939,7 @@ class EncodeColumnConstraint(ColumnConstraintKind):
 
 class GeneratedAsIdentityColumnConstraint(ColumnConstraintKind):
     # this: True -> ALWAYS, this: False -> BY DEFAULT
-    arg_types = {"this": True, "start": False, "increment": False}
+    arg_types = {"this": False, "start": False, "increment": False}
 
 
 class NotNullColumnConstraint(ColumnConstraintKind):
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
index 5b68688be8..0d72fe313f 100644
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -453,6 +453,9 @@ def defaultcolumnconstraint_sql(self, expression: exp.DefaultColumnConstraint) -
     def generatedasidentitycolumnconstraint_sql(
         self, expression: exp.GeneratedAsIdentityColumnConstraint
     ) -> str:
+        this = ""
+        if expression.this is not None:
+            this = " ALWAYS " if expression.this else " BY DEFAULT "
         start = expression.args.get("start")
         start = f"START WITH {start}" if start else ""
         increment = expression.args.get("increment")
@@ -461,9 +464,7 @@ def generatedasidentitycolumnconstraint_sql(
         if start or increment:
             sequence_opts = f"{start} {increment}"
             sequence_opts = f" ({sequence_opts.strip()})"
-        return (
-            f"GENERATED {'ALWAYS' if expression.this else 'BY DEFAULT'} AS IDENTITY{sequence_opts}"
-        )
+        return f"GENERATED{this}AS IDENTITY{sequence_opts}"
 
     def notnullcolumnconstraint_sql(self, expression: exp.NotNullColumnConstraint) -> str:
         return f"{'' if expression.args.get('allow_null') else 'NOT '}NULL"
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
index dfda4eeddd..7bea286228 100644
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2738,6 +2738,9 @@ def _parse_column_def(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.
 
         return self.expression(exp.ColumnDef, this=this, kind=kind, constraints=constraints)
 
+    def _parse_autoincrement(self) -> exp.Expression:
+        return exp.AutoIncrementColumnConstraint()
+
     def _parse_column_constraint(self) -> t.Optional[exp.Expression]:
         this = self._parse_references()
 
@@ -2749,8 +2752,8 @@ def _parse_column_constraint(self) -> t.Optional[exp.Expression]:
 
         kind: exp.Expression
 
-        if self._match(TokenType.AUTO_INCREMENT):
-            kind = exp.AutoIncrementColumnConstraint()
+        if self._match_set((TokenType.AUTO_INCREMENT, TokenType.IDENTITY)):
+            kind = self._parse_autoincrement()
         elif self._match(TokenType.CHECK):
             constraint = self._parse_wrapped(self._parse_conjunction)
             kind = self.expression(exp.CheckColumnConstraint, this=constraint)
diff --git a/tests/dialects/test_snowflake.py b/tests/dialects/test_snowflake.py
index 05e7643039..d1fa3cfdc5 100644
--- a/tests/dialects/test_snowflake.py
+++ b/tests/dialects/test_snowflake.py
@@ -7,7 +7,32 @@ class TestSnowflake(Validator):
 
     def test_snowflake(self):
         self.validate_identity("SELECT REGEXP_LIKE(a, b, c)")
+        self.validate_identity("CREATE TABLE foo (bar FLOAT AUTOINCREMENT START 0 INCREMENT 1)")
 
+        self.validate_all(
+            "CREATE OR REPLACE TEMPORARY TABLE x (y NUMBER IDENTITY(0, 1))",
+            write={
+                "snowflake": "CREATE OR REPLACE TEMPORARY TABLE x (y DECIMAL AUTOINCREMENT START 0 INCREMENT 1)",
+            },
+        )
+        self.validate_all(
+            "CREATE TEMPORARY TABLE x (y NUMBER AUTOINCREMENT(0, 1))",
+            write={
+                "snowflake": "CREATE OR REPLACE TEMPORARY TABLE x (y DECIMAL AUTOINCREMENT START 0 INCREMENT 1)",
+            },
+        )
+        self.validate_all(
+            "CREATE TABLE x (y NUMBER IDENTITY START 0 INCREMENT 1)",
+            write={
+                "snowflake": "CREATE TABLE x (y DECIMAL AUTOINCREMENT START 0 INCREMENT 1)",
+            },
+        )
+        self.validate_all(
+            "ALTER TABLE foo ADD COLUMN id INT identity(1, 1)",
+            write={
+                "snowflake": "ALTER TABLE foo ADD COLUMN id INT AUTOINCREMENT START 1 INCREMENT 1",
+            },
+        )
         self.validate_all(
             "SELECT DAYOFWEEK('2016-01-02T23:39:20.123-07:00'::TIMESTAMP)",
             write={


Code review comments:
diff_hunk: @@ -2738,6 +2738,9 @@ def _parse_column_def(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.
 
         return self.expression(exp.ColumnDef, this=this, kind=kind, constraints=constraints)
 
+    def _parse_autoincrement(self) -> exp.Expression:
path: sqlglot/parser.py
text: @user1:
this is not a parse, it's not matching anything, why is the snowflake one in it's own dialect?

@author:
I factored this out so that it can be overridden easily, e.g. in Snowflake

@author:
No problem with moving the snowflake method in the base parser, just thought I'd override it there because the generator method needed to be overridden for snowflake too.

@user1:
how come?

@author:
discussed in slack