feat(postgres): Support generated columns (#4472)

* feat(postgres): Support STORED computed columns * Update sqlglot/dialects/postgres.py Co-authored-by: Jo <[email protected]> --------- Co-authored-by: Jo <[email protected]>
tobymao · Dec 4, 2024 · a9dca8d · a9dca8d
1 parent 3945acc
commit a9dca8d
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 0 deletions.
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -476,6 +476,20 @@ def _parse_jsonb_exists(self) -> exp.JSONBExists:
                 and self.dialect.to_json_path(self._parse_bitwise()),
             )
 
+        def _parse_generated_as_identity(
+            self,
+        ) -> (
+            exp.GeneratedAsIdentityColumnConstraint
+            | exp.ComputedColumnConstraint
+            | exp.GeneratedAsRowColumnConstraint
+        ):
+            this = super()._parse_generated_as_identity()
+
+            if self._match_text_seq("STORED"):
+                this = self.expression(exp.ComputedColumnConstraint, this=this.expression)
+
+            return this
+
     class Generator(generator.Generator):
         SINGLE_STRING_INTERVAL = True
         RENAME_TABLE_WITH_DB = False
@@ -691,3 +705,6 @@ def array_sql(self, expression: exp.Array) -> str:
                 if isinstance(seq_get(exprs, 0), exp.Select)
                 else f"{self.normalize_func('ARRAY')}[{self.expressions(expression, flat=True)}]"
             )
+
+        def computedcolumnconstraint_sql(self, expression: exp.ComputedColumnConstraint) -> str:
+            return f"GENERATED ALWAYS AS ({self.sql(expression, 'this')}) STORED"
diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -1047,6 +1047,7 @@ def test_ddl(self):
         self.validate_identity("CREATE TABLE tbl (col INT UNIQUE NULLS NOT DISTINCT DEFAULT 9.99)")
         self.validate_identity("CREATE TABLE tbl (col UUID UNIQUE DEFAULT GEN_RANDOM_UUID())")
         self.validate_identity("CREATE TABLE tbl (col UUID, UNIQUE NULLS NOT DISTINCT (col))")
+        self.validate_identity("CREATE TABLE tbl (col_a INT GENERATED ALWAYS AS (1 + 2) STORED)")
 
         self.validate_identity("CREATE INDEX CONCURRENTLY ix_table_id ON tbl USING btree(id)")
         self.validate_identity(