Feat(postgres): add support for ALTER TABLE ONLY ... (#2179)

tobymao · Sep 7, 2023 · cd301cc · cd301cc
1 parent 93b7ba2
commit cd301cc
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/sqlglot/expressions.py b/sqlglot/expressions.py
@@ -3710,7 +3710,7 @@ class Rollback(Expression):
 
 
 class AlterTable(Expression):
-    arg_types = {"this": True, "actions": True, "exists": False}
+    arg_types = {"this": True, "actions": True, "exists": False, "only": False}
 
 
 class AddConstraint(Expression):

diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -2289,7 +2289,8 @@ def altertable_sql(self, expression: exp.AlterTable) -> str:
             actions = self.expressions(expression, key="actions")
 
         exists = " IF EXISTS" if expression.args.get("exists") else ""
-        return f"ALTER TABLE{exists} {self.sql(expression, 'this')} {actions}"
+        only = " ONLY" if expression.args.get("only") else ""
+        return f"ALTER TABLE{exists}{only} {self.sql(expression, 'this')} {actions}"
 
     def droppartition_sql(self, expression: exp.DropPartition) -> str:
         expressions = self.expressions(expression)

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -4785,6 +4785,7 @@ def _parse_alter(self) -> exp.AlterTable | exp.Command:
             return self._parse_as_command(start)
 
         exists = self._parse_exists()
+        only = self._match_text_seq("ONLY")
         this = self._parse_table(schema=True)
 
         if self._next:
@@ -4800,7 +4801,9 @@ def _parse_alter(self) -> exp.AlterTable | exp.Command:
                     this=this,
                     exists=exists,
                     actions=actions,
+                    only=only,
                 )
+
         return self._parse_as_command(start)
 
     def _parse_merge(self) -> exp.Merge:

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -126,12 +126,17 @@ def test_array_offset(self, logger):
         )
 
     def test_postgres(self):
-        self.validate_identity("x @@ y")
-
         expr = parse_one("SELECT * FROM r CROSS JOIN LATERAL UNNEST(ARRAY[1]) AS s(location)")
         unnest = expr.args["joins"][0].this.this
         unnest.assert_is(exp.Unnest)
 
+        alter_table_only = """ALTER TABLE ONLY "Album" ADD CONSTRAINT "FK_AlbumArtistId" FOREIGN KEY ("ArtistId") REFERENCES "Artist" ("ArtistId") ON DELETE NO ACTION ON UPDATE NO ACTION"""
+        expr = parse_one(alter_table_only)
+
+        self.assertIsInstance(expr, exp.AlterTable)
+        self.assertEqual(expr.sql(dialect="postgres"), alter_table_only)
+
+        self.validate_identity("x @@ y")
         self.validate_identity("CAST(x AS MONEY)")
         self.validate_identity("CAST(x AS INT4RANGE)")
         self.validate_identity("CAST(x AS INT4MULTIRANGE)")