Repository: tobymao/sqlglot

Issue:
Support for Check Constraint Transpiling
Both, T-SQL and Spark/Databricks support check constraints.
I had hope to be able to transpile something like this:

```t-sql
ALTER TABLE [schema].[tablename]
    ADD CONSTRAINT [CHK_Name] CHECK ([IdDwh] IS NOT NULL AND [IdDwh]<>(0));
```

Seems not to be supported?


Discussion:

Resolving patch:
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
--- a/sqlglot/expressions.py
+++ b/sqlglot/expressions.py
@@ -2292,10 +2292,17 @@ class Rollback(Expression):
 
 
 class AlterTable(Expression):
+    arg_types = {"this": True, "actions": True, "exists": False}
+
+
+class AddConstraint(Expression):
     arg_types = {
-        "this": True,
-        "actions": True,
-        "exists": False,
+        "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,
     }
 
 
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
--- a/sqlglot/generator.py
+++ b/sqlglot/generator.py
@@ -106,6 +106,8 @@ class Generator:
         exp.TableFormatProperty,
     }
 
+    WITH_SINGLE_ALTER_TABLE_ACTION = (exp.AlterColumn, exp.RenameTable, exp.AddConstraint)
+
     WITH_SEPARATED_COMMENTS = (exp.Select, exp.From, exp.Where, exp.Binary)
     SENTINEL_LINE_BREAK = "__SQLGLOT__LB__"
 
@@ -1364,7 +1366,7 @@ def altertable_sql(self, expression: exp.AlterTable) -> str:
             actions = self.expressions(expression, "actions", prefix="ADD COLUMNS ")
         elif isinstance(actions[0], exp.Drop):
             actions = self.expressions(expression, "actions")
-        elif isinstance(actions[0], (exp.AlterColumn, exp.RenameTable)):
+        elif isinstance(actions[0], self.WITH_SINGLE_ALTER_TABLE_ACTION):
             actions = self.sql(actions[0])
         else:
             self.unsupported(f"Unsupported ALTER TABLE action {actions[0].__class__.__name__}")
@@ -1372,6 +1374,26 @@ def altertable_sql(self, expression: exp.AlterTable) -> str:
         exists = " IF EXISTS" if expression.args.get("exists") else ""
         return f"ALTER TABLE{exists} {self.sql(expression, 'this')} {actions}"
 
+    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}"
+
+        columns = self.expressions(expression, "columns")
+        if primary:
+            return f"{add_constraint} PRIMARY KEY {self.wrap(columns)}{options}"
+
+        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}"
+
     def distinct_sql(self, expression: exp.Distinct) -> str:
         this = self.expressions(expression, flat=True)
         this = f" {this}" if this else ""
diff --git a/sqlglot/parser.py b/sqlglot/parser.py
--- a/sqlglot/parser.py
+++ b/sqlglot/parser.py
@@ -601,6 +601,8 @@ class Parser(metaclass=_Parser):
 
     WINDOW_ALIAS_TOKENS = ID_VAR_TOKENS - {TokenType.ROWS}
 
+    ADD_CONSTRAINT_TOKENS = {TokenType.CONSTRAINT, TokenType.PRIMARY_KEY, TokenType.FOREIGN_KEY}
+
     STRICT_CAST = True
 
     __slots__ = (
@@ -3020,6 +3022,65 @@ def _parse_add_column(self) -> t.Optional[exp.Expression]:
     def _parse_drop_column(self) -> t.Optional[exp.Expression]:
         return self._match(TokenType.DROP) and self._parse_drop(default_kind="COLUMN")
 
+    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"]
+
+                return self.expression(
+                    exp.AddConstraint, this=this, expression=expression, options=options
+                )
+
+        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 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,
+        )
+
     def _parse_alter(self) -> t.Optional[exp.Expression]:
         if not self._match(TokenType.TABLE):
             return None
@@ -3028,8 +3089,14 @@ def _parse_alter(self) -> t.Optional[exp.Expression]:
         this = self._parse_table(schema=True)
 
         actions: t.Optional[exp.Expression | t.List[t.Optional[exp.Expression]]] = None
-        if self._match_text_seq("ADD", advance=False):
-            actions = self._parse_csv(self._parse_add_column)
+
+        index = self._index
+        if self._match_text_seq("ADD"):
+            if self._match_set(self.ADD_CONSTRAINT_TOKENS):
+                actions = self._parse_csv(self._parse_add_constraint)
+            else:
+                self._retreat(index)
+                actions = self._parse_csv(self._parse_add_column)
         elif self._match_text_seq("DROP", advance=False):
             actions = self._parse_csv(self._parse_drop_column)
         elif self._match_text_seq("RENAME", "TO"):


Test patch:
diff --git a/tests/fixtures/identity.sql b/tests/fixtures/identity.sql
--- a/tests/fixtures/identity.sql
+++ b/tests/fixtures/identity.sql
@@ -673,3 +673,9 @@ CREATE TABLE products (x INT GENERATED ALWAYS AS IDENTITY)
 CREATE TABLE IF NOT EXISTS customer (pk BIGINT NOT NULL GENERATED ALWAYS AS IDENTITY (INCREMENT BY 1))
 CREATE TABLE customer (pk BIGINT NOT NULL GENERATED ALWAYS AS IDENTITY (START WITH 10 INCREMENT BY 1))
 CREATE TABLE customer (pk BIGINT NOT NULL GENERATED ALWAYS AS IDENTITY (START WITH 10))
+ALTER TABLE "schema"."tablename" ADD CONSTRAINT "CHK_Name" CHECK (NOT "IdDwh" IS NULL AND "IdDwh" <> (0))
+ALTER TABLE persons ADD CONSTRAINT persons_pk PRIMARY KEY (first_name, last_name)
+ALTER TABLE pets ADD CONSTRAINT pets_persons_fk FOREIGN KEY (owner_first_name, owner_last_name) REFERENCES persons
+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