Fix: allow type column ops for bigquery

tobymao · May 30, 2023 · 6045b74 · 6045b74
1 parent 2792eaa
commit 6045b74
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/sqlglot/generator.py b/sqlglot/generator.py
@@ -802,7 +802,8 @@ def datatype_sql(self, expression: exp.DataType) -> str:
                 nested = f"{self.STRUCT_DELIMITER[0]}{interior}{self.STRUCT_DELIMITER[1]}"
                 if expression.args.get("values") is not None:
                     delimiters = ("[", "]") if type_value == exp.DataType.Type.ARRAY else ("(", ")")
-                    values = f"{delimiters[0]}{self.expressions(expression, key='values')}{delimiters[1]}"
+                    values = self.expressions(expression, key="values", flat=True)
+                    values = f"{delimiters[0]}{values}{delimiters[1]}"
             else:
                 nested = f"({interior})"
 

diff --git a/sqlglot/parser.py b/sqlglot/parser.py
@@ -2842,7 +2842,7 @@ def _parse_type(self) -> t.Optional[exp.Expression]:
             if not data_type.expressions:
                 self._retreat(index)
                 return self._parse_column()
-            return data_type
+            return self._parse_column_ops(data_type)
 
         return this
 
@@ -2985,6 +2985,9 @@ def _parse_column(self) -> t.Optional[exp.Expression]:
             this = self.expression(exp.Column, this=this)
         elif not this:
             return self._parse_bracket(this)
+        return self._parse_column_ops(this)
+
+    def _parse_column_ops(self, this: exp.Expression) -> exp.Expression:
         this = self._parse_bracket(this)
 
         while self._match_set(self.COLUMN_OPERATORS):
@@ -3029,7 +3032,6 @@ def _parse_column(self) -> t.Optional[exp.Expression]:
             else:
                 this = self.expression(exp.Dot, this=this, expression=field)
             this = self._parse_bracket(this)
-
         return this
 
     def _parse_primary(self) -> t.Optional[exp.Expression]:
@@ -3462,7 +3464,15 @@ def _parse_primary_key(self) -> exp.Expression:
         options = self._parse_key_constraint_options()
         return self.expression(exp.PrimaryKey, expressions=expressions, options=options)
 
+    @t.overload
+    def _parse_bracket(self, this: exp.Expression) -> exp.Expression:
+        ...
+
+    @t.overload
     def _parse_bracket(self, this: t.Optional[exp.Expression]) -> t.Optional[exp.Expression]:
+        ...
+
+    def _parse_bracket(self, this):
         if not self._match_set((TokenType.L_BRACKET, TokenType.L_BRACE)):
             return this
 
@@ -4516,9 +4526,17 @@ def _match_text_seq(self, *texts, advance=True):
 
         return True
 
+    @t.overload
+    def _replace_columns_with_dots(self, this: exp.Expression) -> exp.Expression:
+        ...
+
+    @t.overload
     def _replace_columns_with_dots(
         self, this: t.Optional[exp.Expression]
     ) -> t.Optional[exp.Expression]:
+        ...
+
+    def _replace_columns_with_dots(self, this):
         if isinstance(this, exp.Dot):
             exp.replace_children(this, self._replace_columns_with_dots)
         elif isinstance(this, exp.Column):

diff --git a/tests/dialects/test_bigquery.py b/tests/dialects/test_bigquery.py
@@ -21,6 +21,7 @@ def test_bigquery(self):
         self.validate_identity("SELECT DISTINCT AS STRUCT 1 AS a, 2 AS b")
         self.validate_identity("SELECT AS VALUE STRUCT(1 AS a, 2 AS b)")
         self.validate_identity("SELECT STRUCT<ARRAY<STRING>>(['2023-01-17'])")
+        self.validate_identity("SELECT STRUCT<STRING>((SELECT a FROM b.c LIMIT 1)).*")
         self.validate_identity("SELECT * FROM q UNPIVOT(values FOR quarter IN (b, c))")
         self.validate_identity("""CREATE TABLE x (a STRUCT<values ARRAY<INT64>>)""")
         self.validate_identity("""CREATE TABLE x (a STRUCT<b STRING OPTIONS (description='b')>)""")