feat: STRUCT and ARRAY support

sqlalchemy_bigquery/_struct.py

@@ -34,8 +34,13 @@
     import sqlalchemy.sql.coercions
     import sqlalchemy.sql.roles
 
-# We have to delay getting the type compiler, because of circular imports. :(
-type_compiler = None
+
+def _get_subtype_col_spec(type_):
+    global _get_subtype_col_spec
+
+    type_compiler = base.dialect.type_compiler(base.dialect())
+    _get_subtype_col_spec = type_compiler.process
+    return _get_subtype_col_spec(type_)
 
 
 class STRUCT(sqlalchemy.sql.sqltypes.Indexable, sqlalchemy.types.UserDefinedType):
@@ -55,30 +60,29 @@ def __init__(
         # Note that because:
         # https://docs.python.org/3/whatsnew/3.6.html#pep-468-preserving-keyword-argument-order
         # We know that `kwfields` preserves order.
-        self.__fields = tuple(
+        self._STRUCT_fields = tuple(
             (
                 name,
                 type_ if isinstance(type_, sqlalchemy.types.TypeEngine) else type_(),
             )
             for (name, type_) in (fields + tuple(kwfields.items()))
         )
 
-        self.__byname = {name.lower(): type_ for (name, type_) in self.__fields}
+        self._STRUCT_byname = {
+            name.lower(): type_ for (name, type_) in self._STRUCT_fields
+        }
 
     def __repr__(self):
-        fields = ", ".join(f"{name}={repr(type_)}" for name, type_ in self.__fields)
+        fields = ", ".join(
+            f"{name}={repr(type_)}" for name, type_ in self._STRUCT_fields
+        )
         return f"STRUCT({fields})"
 
     def get_col_spec(self, **kw):
-        global type_compiler
-
-        try:
-            process = type_compiler.process
-        except AttributeError:
-            type_compiler = base.dialect.type_compiler(base.dialect())
-            process = type_compiler.process
-
-        fields = ", ".join(f"{name} {process(type_)}" for name, type_ in self.__fields)
+        fields = ", ".join(
+            f"{name} {_get_subtype_col_spec(type_)}"
+            for name, type_ in self._STRUCT_fields
+        )
         return f"STRUCT<{fields}>"
 
     def bind_processor(self, dialect):
@@ -89,22 +93,26 @@ def _setup_getitem(self, name):
             if not isinstance(name, str):
                 raise TypeError(
                     f"STRUCT fields can only be accessed with strings field names,"
-                    f" not {name}."
+                    f" not {repr(name)}."
                 )
-            subtype = self.expr.type._STRUCT__byname.get(name.lower())
+            subtype = self.expr.type._STRUCT_byname.get(name.lower())
             if subtype is None:
                 raise KeyError(name)
             operator = struct_getitem_op
             index = _field_index(self, name, operator)
             return operator, index, subtype
 
         def __getattr__(self, name):
-            if name.lower() in self.expr.type._STRUCT__byname:
+            if name.lower() in self.expr.type._STRUCT_byname:
                 return self[name]
 
     comparator_factory = Comparator
 
 
+# In the implementations of _field_index below, we're stealing from
+# the JSON type implementation, but the code to steal changed in
+# 1.4. :/
+
 if sqlalchemy_1_4_or_more:
 
     def _field_index(self, name, operator):

sqlalchemy_bigquery/_types.py

@@ -72,25 +72,25 @@
 except NameError:
     pass
 
+STRUCT_FIELD_TYPES = "RECORD", "STRUCT"
 
-def _get_transitive_schema_fields(fields, cur_fields):
+
+def _get_transitive_schema_fields(fields):
     """
     Recurse into record type and return all the nested field names.
     As contributed by @sumedhsakdeo on issue #17
     """
     results = []
     for field in fields:
         results += [field]
-        if field.field_type == "RECORD":
-            cur_fields.append(field)
+        if field.field_type in STRUCT_FIELD_TYPES:
             sub_fields = [
                 SchemaField.from_api_repr(
                     dict(f.to_api_repr(), name=f"{field.name}.{f.name}")
                 )
                 for f in field.fields
             ]
-            results += _get_transitive_schema_fields(sub_fields, cur_fields)
-            cur_fields.pop()
+            results += _get_transitive_schema_fields(sub_fields)
     return results
 
 
@@ -125,7 +125,7 @@ def _get_sqla_column_type(field):
 
 
 def get_columns(bq_schema):
-    fields = _get_transitive_schema_fields(bq_schema, [])
+    fields = _get_transitive_schema_fields(bq_schema)
     return [
         {
             "name": field.name,

tests/unit/test__struct.py

@@ -64,8 +64,10 @@ def _col():
         (_col().NAME, "`t`.`person`.NAME"),
         (_col().children, "`t`.`person`.children"),
         (
-            _col().children[0].label("anon_1"),  # SQLAlchemy doesn't add the label
-            # in this case for some reason
+            # SQLAlchemy doesn't add the label in this case for some reason.
+            # TODO: why?
+            # https://github.com/googleapis/python-bigquery-sqlalchemy/issues/336
+            _col().children[0].label("anon_1"),
             "(`t`.`person`.children)[OFFSET(%(param_1:INT64)s)]",
         ),
         (