Repository: tobymao/sqlglot

Issue:
Parser support for deferrable foreign key constraints in column def
The parser does not expect an optional "deferrable" clause for foreign key constraints inside column definitions.

Example:

```python
sql = "CREATE TABLE foo (baz_id INTEGER REFERENCES baz (id) DEFERRABLE);"
sqlglot.parse_one(sql)
```

```shell
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/foo/product/sqlglot/sqlglot/__init__.py", line 84, in parse_one
    result = dialect.parse(sql, **opts)
  File "/Users/foo/product/sqlglot/sqlglot/dialects/dialect.py", line 151, in parse
    return self.parser(**opts).parse(self.tokenizer.tokenize(sql), sql)
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 672, in parse
    return self._parse(
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 734, in _parse
    expressions.append(parse_method(self))
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 874, in _parse_statement
    return self.STATEMENT_PARSERS[self._prev.token_type](self)
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 433, in <lambda>
    TokenType.CREATE: lambda self: self._parse_create(),
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 955, in _parse_create
    this = self._parse_table(schema=True)
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 1603, in _parse_table
    return self._parse_schema(this=this)
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 2366, in _parse_schema
    self._match_r_paren()
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 3173, in _match_r_paren
    self.raise_error("Expecting )")
  File "/Users/foo/product/sqlglot/sqlglot/parser.py", line 780, in raise_error
    raise error
sqlglot.errors.ParseError: Expecting ). Line 1, Col: 54.
  CREATE TABLE foo (baz_id INTEGER REFERENCES baz (id) DEFERRABLE);
```

Variations: `DEFERRABLE`, `DEFERRABLE INITIALLY DEFERRED`, `DEFERRABLE IMMEDIATELY DEFERRED`

Refs:
[sqlite syntax](https://www.sqlite.org/syntax/foreign-key-clause.html)
[postgres syntax](https://www.postgresql.org/docs/current/sql-createtable.html)



Discussion:

Resolving patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -887,6 +887,10 @@ class ForeignKey(Expression):
     }
 
 
+class PrimaryKey(Expression):
+    arg_types = {"expressions": True, "options": False}
+
+
 class Unique(Expression):
     arg_types = {"expressions": True}
 
@@ -1298,7 +1302,7 @@ class Return(Expression):
 
 
 class Reference(Expression):
-    arg_types = {"this": True, "expressions": True}
+    arg_types = {"this": True, "expressions": False, "options": False}
 
 
 class Tuple(Expression):
@@ -2341,14 +2345,7 @@ class AlterTable(Expression):
 
 
 class AddConstraint(Expression):
-    arg_types = {
-        "this": False,
-        "expression": False,  # Only used in the CHECK constraint
-        "primary": False,  # None: CHECK, False: FOREIGN KEY, True: PRIMARY KEY
-        "columns": False,
-        "references": False,
-        "options": False,
-    }
+    arg_types = {"this": False, "expression": False, "enforced": False}
 
 
 # Binary expressions like (ADD a b)
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -1199,6 +1199,12 @@ def foreignkey_sql(self, expression: exp.ForeignKey) -> str:
         update = f" ON UPDATE {update}" if update else ""
         return f"FOREIGN KEY ({expressions}){reference}{delete}{update}"
 
+    def primarykey_sql(self, expression: exp.ForeignKey) -> str:
+        expressions = self.expressions(expression, flat=True)
+        options = self.expressions(expression, "options", flat=True, sep=" ")
+        options = f" {options}" if options else ""
+        return f"PRIMARY KEY ({expressions}){options}"
+
     def unique_sql(self, expression: exp.Unique) -> str:
         columns = self.expressions(expression, key="expressions")
         return f"UNIQUE ({columns})"
@@ -1241,7 +1247,10 @@ def return_sql(self, expression: exp.Return) -> str:
     def reference_sql(self, expression: exp.Reference) -> str:
         this = self.sql(expression, "this")
         expressions = self.expressions(expression, flat=True)
-        return f"REFERENCES {this}({expressions})"
+        expressions = f"({expressions})" if expressions else ""
+        options = self.expressions(expression, "options", flat=True, sep=" ")
+        options = f" {options}" if options else ""
+        return f"REFERENCES {this}{expressions}{options}"
 
     def anonymous_sql(self, expression: exp.Anonymous) -> str:
         args = self.format_args(*expression.expressions)
@@ -1381,23 +1390,14 @@ def altertable_sql(self, expression: exp.AlterTable) -> str:
 
     def addconstraint_sql(self, expression: exp.AddConstraint) -> str:
         this = self.sql(expression, "this")
-        check = self.sql(expression, "expression")
-        primary = expression.args.get("primary")
-        options = self.expressions(expression, "options", sep=" ")
-        options = f" {options}" if options else ""
-
-        add_constraint = f"ADD CONSTRAINT {this}" if this else "ADD CONSTRAINT"
-
-        if primary is None:
-            return f"{add_constraint} CHECK ({check}){options}"
+        expression_ = self.sql(expression, "expression")
+        add_constraint = f"ADD CONSTRAINT {this}" if this else "ADD"
 
-        columns = self.expressions(expression, "columns")
-        if primary:
-            return f"{add_constraint} PRIMARY KEY {self.wrap(columns)}{options}"
+        enforced = expression.args.get("enforced")
+        if enforced is not None:
+            return f"{add_constraint} CHECK ({expression_}){' ENFORCED' if enforced else ''}"
 
-        references = expression.args.get("references")
-        references = f" REFERENCES {self.sql(references)}" if references else ""
-        return f"{add_constraint} FOREIGN KEY {self.wrap(columns)}{options}{references}"
+        return f"{add_constraint} {expression_}"
 
     def distinct_sql(self, expression: exp.Distinct) -> str:
         this = self.expressions(expression, flat=True)
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -2477,15 +2477,43 @@ def _parse_unnamed_constraint(self) -> t.Optional[exp.Expression]:
     def _parse_unique(self) -> exp.Expression:
         return self.expression(exp.Unique, expressions=self._parse_wrapped_id_vars())
 
+    def _parse_key_constraint_options(self) -> t.List[str]:
+        options = []
+        while True:
+            if not self._curr:
+                break
+
+            if self._match_text_seq("NOT", "ENFORCED"):
+                options.append("NOT ENFORCED")
+            elif self._match_text_seq("DEFERRABLE"):
+                options.append("DEFERRABLE")
+            elif self._match_text_seq("INITIALLY", "DEFERRED"):
+                options.append("INITIALLY DEFERRED")
+            elif self._match_text_seq("NORELY"):
+                options.append("NORELY")
+            elif self._match_text_seq("MATCH", "FULL"):
+                options.append("MATCH FULL")
+            elif self._match_text_seq("ON", "UPDATE", "NO ACTION"):
+                options.append("ON UPDATE NO ACTION")
+            elif self._match_text_seq("ON", "DELETE", "NO ACTION"):
+                options.append("ON DELETE NO ACTION")
+            else:
+                break
+
+        return options
+
     def _parse_references(self) -> t.Optional[exp.Expression]:
         if not self._match(TokenType.REFERENCES):
             return None
 
-        return self.expression(
-            exp.Reference,
-            this=self._parse_id_var(),
-            expressions=self._parse_wrapped_id_vars(),
-        )
+        expressions = None
+        this = self._parse_id_var()
+
+        if self._match(TokenType.L_PAREN, advance=False):
+            expressions = self._parse_wrapped_id_vars()
+
+        options = self._parse_key_constraint_options()
+        return self.expression(exp.Reference, this=this, expressions=expressions, options=options)
 
     def _parse_foreign_key(self) -> exp.Expression:
         expressions = self._parse_wrapped_id_vars()
@@ -2510,12 +2538,14 @@ def _parse_foreign_key(self) -> exp.Expression:
             options[kind] = action
 
         return self.expression(
-            exp.ForeignKey,
-            expressions=expressions,
-            reference=reference,
-            **options,  # type: ignore
+            exp.ForeignKey, expressions=expressions, reference=reference, **options  # type: ignore
         )
 
+    def _parse_primary_key(self) -> exp.Expression:
+        expressions = self._parse_wrapped_id_vars()
+        options = self._parse_key_constraint_options()
+        return self.expression(exp.PrimaryKey, expressions=expressions, options=options)
+
     def _parse_bracket(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
         if not self._match(TokenType.L_BRACKET):
             return this
@@ -3031,62 +3061,25 @@ def _parse_drop_column(self) -> t.Optional[exp.Expression]:
 
     def _parse_add_constraint(self) -> t.Optional[exp.Expression]:
         this = None
-        options = None
-
         kind = self._prev.token_type
+
         if kind == TokenType.CONSTRAINT:
             this = self._parse_id_var()
 
             if self._match(TokenType.CHECK):
                 expression = self._parse_wrapped(self._parse_conjunction)
-
-                if self._match_text_seq("ENFORCED"):
-                    options = ["ENFORCED"]
+                enforced = self._match_text_seq("ENFORCED")
 
                 return self.expression(
-                    exp.AddConstraint, this=this, expression=expression, options=options
+                    exp.AddConstraint, this=this, expression=expression, enforced=enforced
                 )
 
-        if self._match_set((TokenType.PRIMARY_KEY, TokenType.FOREIGN_KEY)):
-            kind = self._prev.token_type
-
-        references = None
-        primary = kind == TokenType.PRIMARY_KEY
-        columns = self._parse_wrapped_csv(self._parse_column)
-
-        if self._match(TokenType.REFERENCES):
-            references = self._parse_table(schema=True)
-
-        options = []
-        while True:
-            if not self._curr:
-                break
+        if kind == TokenType.FOREIGN_KEY or self._match(TokenType.FOREIGN_KEY):
+            expression = self._parse_foreign_key()
+        elif kind == TokenType.PRIMARY_KEY or self._match(TokenType.PRIMARY_KEY):
+            expression = self._parse_primary_key()
 
-            if self._match_text_seq("NOT ENFORCED"):
-                options.append("NOT ENFORCED")
-            elif self._match_text_seq("DEFERRABLE"):
-                options.append("DEFERRABLE")
-            elif self._match_text_seq("INITIALLY DEFERRED"):
-                options.append("INITIALLY DEFERRED")
-            elif self._match_text_seq("NORELY"):
-                options.append("NORELY")
-            elif self._match_text_seq("MATCH FULL"):
-                options.append("MATCH FULL")
-            elif self._match_text_seq("ON UPDATE NO ACTION"):
-                options.append("ON UPDATE NO ACTION")
-            elif self._match_text_seq("ON DELETE NO ACTION"):
-                options.append("ON DELETE NO ACTION")
-            else:
-                self.raise_error(f"Invalid ADD CONSTRAINT option {self._curr.text}.")
-
-        return self.expression(
-            exp.AddConstraint,
-            this=this,
-            primary=primary,
-            columns=columns,
-            references=references,
-            options=options,
-        )
+        return self.expression(exp.AddConstraint, this=this, expression=expression)
 
     def _parse_alter(self) -> t.Optional[exp.Expression]:
         if not self._match(TokenType.TABLE):


Test patch:
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -679,3 +679,9 @@ ALTER TABLE pets ADD CONSTRAINT pets_persons_fk FOREIGN KEY (owner_first_name, o
 ALTER TABLE pets ADD CONSTRAINT pets_name_not_cute_chk CHECK (LENGTH(name) < 20)
 ALTER TABLE people10m ADD CONSTRAINT dateWithinRange CHECK (birthDate > '1900-01-01')
 ALTER TABLE people10m ADD CONSTRAINT validIds CHECK (id > 1 AND id < 99999999) ENFORCED
+ALTER TABLE baa ADD CONSTRAINT boo PRIMARY KEY (x, y) NOT ENFORCED DEFERRABLE INITIALLY DEFERRED NORELY
+ALTER TABLE baa ADD CONSTRAINT boo PRIMARY KEY (x, y) NOT ENFORCED DEFERRABLE INITIALLY DEFERRED NORELY
+ALTER TABLE baa ADD CONSTRAINT boo FOREIGN KEY (x, y) REFERENCES persons ON UPDATE NO ACTION ON DELETE NO ACTION MATCH FULL
+ALTER TABLE a ADD PRIMARY KEY (x, y) NOT ENFORCED
+ALTER TABLE a ADD FOREIGN KEY (x, y) REFERENCES bla
+CREATE TABLE foo (baz_id INT REFERENCES baz(id) DEFERRABLE)