Feat(postgres): add support for col int[size] column def syntax (#3548)

* Feat(postgres): add support for col int[size] column def syntax * Add support for the INT ARRAY[3] syntax * Move value generation inside expression.expressions branch
tobymao · May 24, 2024 · fe56e64 · fe56e64
1 parent 1a8a16b
commit fe56e64
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 3 deletions.
diff --git a/sqlglot/dialects/postgres.py b/sqlglot/dialects/postgres.py
@@ -116,7 +116,10 @@ def _string_agg_sql(self: Postgres.Generator, expression: exp.GroupConcat) -> st
 
 def _datatype_sql(self: Postgres.Generator, expression: exp.DataType) -> str:
     if expression.is_type("array"):
-        return f"{self.expressions(expression, flat=True)}[]" if expression.expressions else "ARRAY"
+        if expression.expressions:
+            values = self.expressions(expression, key="values", flat=True)
+            return f"{self.expressions(expression, flat=True)}[{values}]"
+        return "ARRAY"
     return self.datatype_sql(expression)
 
 

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -4375,8 +4375,26 @@ def _parse_types(
         elif expressions:
             this.set("expressions", expressions)
 
-        while self._match_pair(TokenType.L_BRACKET, TokenType.R_BRACKET):
-            this = exp.DataType(this=exp.DataType.Type.ARRAY, expressions=[this], nested=True)
+        index = self._index
+
+        # Postgres supports the INT ARRAY[3] syntax as a synonym for INT[3]
+        matched_array = self._match(TokenType.ARRAY)
+
+        while self._curr:
+            matched_l_bracket = self._match(TokenType.L_BRACKET)
+            if not matched_l_bracket and not matched_array:
+                break
+
+            matched_array = False
+            values = self._parse_csv(self._parse_conjunction) or None
+            if values and not schema:
+                self._retreat(index)
+                break
+
+            this = exp.DataType(
+                this=exp.DataType.Type.ARRAY, expressions=[this], values=values, nested=True
+            )
+            self._match(TokenType.R_BRACKET)
 
         if self.TYPE_CONVERTER and isinstance(this.this, exp.DataType.Type):
             converter = self.TYPE_CONVERTER.get(this.this)

diff --git a/tests/dialects/test_postgres.py b/tests/dialects/test_postgres.py
@@ -734,6 +734,8 @@ def test_ddl(self):
         cdef.args["kind"].assert_is(exp.DataType)
         self.assertEqual(expr.sql(dialect="postgres"), "CREATE TABLE t (x INTERVAL DAY)")
 
+        self.validate_identity("CREATE TABLE t (col INT[3][5])")
+        self.validate_identity("CREATE TABLE t (col INT[3])")
         self.validate_identity("CREATE INDEX IF NOT EXISTS ON t(c)")
         self.validate_identity("CREATE INDEX et_vid_idx ON et(vid) INCLUDE (fid)")
         self.validate_identity("CREATE INDEX idx_x ON x USING BTREE(x, y) WHERE (NOT y IS NULL)")
@@ -856,6 +858,14 @@ def test_ddl(self):
         self.validate_identity(
             "CREATE UNLOGGED TABLE foo AS WITH t(c) AS (SELECT 1) SELECT * FROM (SELECT c AS c FROM t) AS temp"
         )
+        self.validate_identity(
+            "CREATE TABLE t (col integer ARRAY[3])",
+            "CREATE TABLE t (col INT[3])",
+        )
+        self.validate_identity(
+            "CREATE TABLE t (col integer ARRAY)",
+            "CREATE TABLE t (col INT[])",
+        )
         self.validate_identity(
             "CREATE FUNCTION x(INT) RETURNS INT SET search_path TO 'public'",
             "CREATE FUNCTION x(INT) RETURNS INT SET search_path = 'public'",