Repository: dbt-labs/dbt-bigquery
Pull request #806 (2023-07-01 01:29:32)

Contracts: Handle struct column specified both at root and nested levels + arrays of structs
resolves https://github.com/dbt-labs/dbt-bigquery/issues/782
resolves https://github.com/dbt-labs/dbt-bigquery/issues/781

<!---
  Include the number of the issue addressed by this PR above if applicable.
  PRs for code changes without an associated issue *will not be merged*.
  See CONTRIBUTING.md for more information.

  Example:
    resolves #1234
-->

### Description

Handles scenarios where a nested column type may be specified both at the struct-level (e.g. `b.id: int`, `b.name: string`) and parent level `b: struct`. Additionally handles when parent level data_type is used to specify an array of structs. 

This PR resolves two issues -- initially I had aimed to just tackle 782 on its own by ignoring a top-level data_type if specified (preserving its constraints) but the top-level array specification was just a couple additional lines at that point, and figured reviewing / testing it all together would be more straightforward since the issues and affected codepaths were very closely related. 

<!---
  Describe the Pull Request here. Add any references and info to help reviewers
  understand your changes. Include any tradeoffs you considered.
-->

### Checklist

- [x] I have read [the contributing guide](https://github.com/dbt-labs/dbt-bigquery/blob/main/CONTRIBUTING.md) and understand what's expected of me
- [x] I have signed the [CLA](https://docs.getdbt.com/docs/contributor-license-agreements)
- [x] I have run this code in development and it appears to resolve the stated issue
- [x] This PR includes tests, or tests are not required/relevant for this PR
 - scope of changes is entirely covered by unit tests 🎉  
- [ ] I have [opened an issue to add/update docs](https://github.com/dbt-labs/docs.getdbt.com/issues/new/choose), or docs changes are not required/relevant for this PR
- [x] I have run `changie new` to [create a changelog entry](https://github.com/dbt-labs/dbt-bigquery/blob/main/CONTRIBUTING.md#Adding-CHANGELOG-Entry)

🎩 parent array
```
-- schema.yml
  - name: nested_array_field
    config:
      contract: {enforced: true}
      materialized: table
    columns:
      - name: my_array
        data_type: array
        constraints:
          - type: not_null
      - name: my_array.id
        data_type: integer
        description: Desc for field_1
      - name: my_array.name
        data_type: string
        description: Desc for field_2
```
nested_array_field.sql
```
select
  array( select struct(
    1 as id,
    'name' as name)
  ) as my_array

  UNION ALL 

  select null as my_array
```

🎩 parent specification with constraints

```
schema.yml

  - name: nested_fields
    config:
      materialized: table
      contract: {enforced: true}
    columns:
      - name: a
        data_type: string
        description: "field a"
      - name: b.id
        data_type: integer
        description: "nested field b.id"
        tests:
          - not_null
      - name: b.name
        data_type: string
      - name: b
        description: "top-level"
        data_type: struct
        constraints:
          - type: not_null
 ```

nested_fields.sql
```
select
  'string' as a,
  struct(
    1 as id,
    'name' as name
  ) as b

  UNION ALL 

select 'string2' as a, null as b
```

github-actions[bot]: Thank you for your pull request! We could not find a changelog entry for this change. For details on how to document a change, see the [dbt-bigquery contributing guide](https://github.com/dbt-labs/dbt-bigquery/blob/main/CONTRIBUTING.md).
MichelleArk: @colin-rogers-dbt - I've tagged you for review as this is follow-on work from https://github.com/dbt-labs/dbt-bigquery/pull/738.
MichelleArk: 🎩 - Missing data_type: 
```sh
❯ dbt run --select nested_fields --project-dir ~/basic-dbt
17:16:08  Running with dbt=1.6.0-b8
17:16:08  Registered adapter: bigquery=1.6.0-b4
17:16:08  Found 5 models, 1 test, 0 snapshots, 0 analyses, 498 macros, 0 operations, 0 seed files, 0 sources, 0 exposures, 0 metrics, 0 groups
17:16:08  
17:16:10  Concurrency: 1 threads (target='dev')
17:16:10  
17:16:10  1 of 1 START sql table model dbt_marky.nested_fields ........................... [RUN]
17:16:10  1 of 1 ERROR creating sql table model dbt_marky.nested_fields .................. [ERROR in 0.83s]
17:16:10  
17:16:10  Finished running 1 table model in 0 hours 0 minutes and 2.01 seconds (2.01s).
17:16:11  
17:16:11  Completed with 1 error and 0 warnings:
17:16:11  
17:16:11  Compilation Error in model nested_fields (models/nested_fields.sql)
17:16:11    Contracted models require data_type to be defined for each column. Please ensure that the column name and data_type are defined within the YAML configuration for the ['a', 'b.name'] column(s).
17:16:11    
17:16:11    > in macro bigquery__get_empty_schema_sql (macros/utils/get_columns_spec_ddl.sql)
...
```

Patch:
diff --git a/.changes/unreleased/Fixes-20230630-213112.yaml b/.changes/unreleased/Fixes-20230630-213112.yaml
new file mode 100644
index 0000000..7238c0c
--- /dev/null
+++ b/.changes/unreleased/Fixes-20230630-213112.yaml
@@ -0,0 +1,7 @@
+kind: Fixes
+body: 'Contracts: Handle struct column specified both at root and nested levels +
+  arrays of structs'
+time: 2023-06-30T21:31:12.63257-04:00
+custom:
+  Author: michelleark
+  Issue: 781 782
diff --git a/dbt/adapters/bigquery/column.py b/dbt/adapters/bigquery/column.py
index cfa298b..a5a60cf 100644
--- a/dbt/adapters/bigquery/column.py
+++ b/dbt/adapters/bigquery/column.py
@@ -5,6 +5,8 @@ from dbt.adapters.base.column import Column
 
 from google.cloud.bigquery import SchemaField
 
+_PARENT_DATA_TYPE_KEY = "__parent_data_type"
+
 Self = TypeVar("Self", bound="BigQueryColumn")
 
 
@@ -131,7 +133,7 @@ class BigQueryColumn(Column):
 def get_nested_column_data_types(
     columns: Dict[str, Dict[str, Any]],
     constraints: Optional[Dict[str, str]] = None,
-) -> Dict[str, Dict[str, str]]:
+) -> Dict[str, Dict[str, Optional[str]]]:
     """
     columns:
         * Dictionary where keys are of flat columns names and values are dictionary of column attributes
@@ -159,16 +161,16 @@ def get_nested_column_data_types(
     """
     constraints = constraints or {}
 
-    nested_column_data_types: Dict[str, Union[str, Dict]] = {}
+    nested_column_data_types: Dict[str, Optional[Union[str, Dict]]] = {}
     for column in columns.values():
         _update_nested_column_data_types(
             column["name"],
-            column["data_type"],
+            column.get("data_type"),
             constraints.get(column["name"]),
             nested_column_data_types,
         )
 
-    formatted_nested_column_data_types: Dict[str, Dict[str, str]] = {}
+    formatted_nested_column_data_types: Dict[str, Dict[str, Optional[str]]] = {}
     for column_name, unformatted_column_type in nested_column_data_types.items():
         formatted_nested_column_data_types[column_name] = {
             "name": column_name,
@@ -191,9 +193,9 @@ def get_nested_column_data_types(
 
 def _update_nested_column_data_types(
     column_name: str,
-    column_data_type: str,
+    column_data_type: Optional[str],
     column_rendered_constraint: Optional[str],
-    nested_column_data_types: Dict[str, Union[str, Dict]],
+    nested_column_data_types: Dict[str, Optional[Union[str, Dict]]],
 ) -> None:
     """
     Recursively update nested_column_data_types given a column_name, column_data_type, and optional column_rendered_constraint.
@@ -215,15 +217,38 @@ def _update_nested_column_data_types(
 
     if len(column_name_parts) == 1:
         # Base case: column is not nested - store its data_type concatenated with constraint if provided.
-        nested_column_data_types[root_column_name] = (
-            column_data_type
-            if column_rendered_constraint is None
-            else f"{column_data_type} {column_rendered_constraint}"
+        column_data_type_and_constraints = (
+            (
+                column_data_type
+                if column_rendered_constraint is None
+                else f"{column_data_type} {column_rendered_constraint}"
+            )
+            if column_data_type
+            else None
         )
+
+        if existing_nested_column_data_type := nested_column_data_types.get(root_column_name):
+            assert isinstance(existing_nested_column_data_type, dict)  # keeping mypy happy
+            # entry could already exist if this is a parent column -- preserve the parent data type under "_PARENT_DATA_TYPE_KEY"
+            existing_nested_column_data_type.update(
+                {_PARENT_DATA_TYPE_KEY: column_data_type_and_constraints}
+            )
+        else:
+            nested_column_data_types.update({root_column_name: column_data_type_and_constraints})
     else:
-        # Initialize nested dictionary
-        if root_column_name not in nested_column_data_types:
-            nested_column_data_types[root_column_name] = {}
+        parent_data_type = nested_column_data_types.get(root_column_name)
+        if isinstance(parent_data_type, dict):
+            # nested dictionary already initialized
+            pass
+        elif parent_data_type is None:
+            # initialize nested dictionary
+            nested_column_data_types.update({root_column_name: {}})
+        else:
+            # a parent specified its base type -- preserve its data_type and potential rendered constraints
+            # this is used to specify a top-level 'struct' or 'array' field with its own description, constraints, etc
+            nested_column_data_types.update(
+                {root_column_name: {_PARENT_DATA_TYPE_KEY: parent_data_type}}
+            )
 
         # Recursively process rest of remaining column name
         remaining_column_name = ".".join(column_name_parts[1:])
@@ -237,7 +262,9 @@ def _update_nested_column_data_types(
         )
 
 
-def _format_nested_data_type(unformatted_nested_data_type: Union[str, Dict[str, Any]]) -> str:
+def _format_nested_data_type(
+    unformatted_nested_data_type: Optional[Union[str, Dict[str, Any]]]
+) -> Optional[str]:
     """
     Recursively format a (STRUCT) data type given an arbitrarily nested data type structure.
 
@@ -249,11 +276,27 @@ def _format_nested_data_type(unformatted_nested_data_type: Union[str, Dict[str,
     >>> BigQueryAdapter._format_nested_data_type({'c': 'string not_null', 'd': {'e': 'string'}})
     'struct<c string not_null, d struct<e string>>'
     """
-    if isinstance(unformatted_nested_data_type, str):
+    if unformatted_nested_data_type is None:
+        return None
+    elif isinstance(unformatted_nested_data_type, str):
         return unformatted_nested_data_type
     else:
+        parent_data_type, *parent_constraints = unformatted_nested_data_type.pop(
+            _PARENT_DATA_TYPE_KEY, ""
+        ).split() or [None]
+
         formatted_nested_types = [
-            f"{column_name} {_format_nested_data_type(column_type)}"
+            f"{column_name} {_format_nested_data_type(column_type) or ''}".strip()
             for column_name, column_type in unformatted_nested_data_type.items()
         ]
-        return f"""struct<{", ".join(formatted_nested_types)}>"""
+
+        formatted_nested_type = f"""struct<{", ".join(formatted_nested_types)}>"""
+
+        if parent_data_type and parent_data_type.lower() == "array":
+            formatted_nested_type = f"""array<{formatted_nested_type}>"""
+
+        if parent_constraints:
+            parent_constraints = " ".join(parent_constraints)
+            formatted_nested_type = f"""{formatted_nested_type} {parent_constraints}"""
+
+        return formatted_nested_type
diff --git a/dbt/adapters/bigquery/impl.py b/dbt/adapters/bigquery/impl.py
index 5c965ca..353be08 100644
--- a/dbt/adapters/bigquery/impl.py
+++ b/dbt/adapters/bigquery/impl.py
@@ -300,7 +300,7 @@ class BigQueryAdapter(BaseAdapter):
         cls,
         columns: Dict[str, Dict[str, Any]],
         constraints: Optional[Dict[str, str]] = None,
-    ) -> Dict[str, Dict[str, str]]:
+    ) -> Dict[str, Dict[str, Optional[str]]]:
         return get_nested_column_data_types(columns, constraints)
 
     def get_columns_in_relation(self, relation: BigQueryRelation) -> List[BigQueryColumn]:
diff --git a/dbt/include/bigquery/macros/utils/get_columns_spec_ddl.sql b/dbt/include/bigquery/macros/utils/get_columns_spec_ddl.sql
index a1247dc..1a4193c 100644
--- a/dbt/include/bigquery/macros/utils/get_columns_spec_ddl.sql
+++ b/dbt/include/bigquery/macros/utils/get_columns_spec_ddl.sql
@@ -5,6 +5,16 @@
 {%- endmacro -%}
 
 {% macro bigquery__get_empty_schema_sql(columns) %}
+    {%- set col_err = [] -%}
+    {% for col in columns.values() %}
+      {%- if col['data_type'] is not defined -%}
+        {{ col_err.append(col['name']) }}
+      {%- endif -%}
+    {%- endfor -%}
+    {%- if (col_err | length) > 0 -%}
+      {{ exceptions.column_type_missing(column_names=col_err) }}
+    {%- endif -%}
+
     {%- set columns = adapter.nest_column_data_types(columns) -%}
     {{ return(dbt.default__get_empty_schema_sql(columns)) }}
 {% endmacro %}


Test patch:
diff --git a/tests/unit/test_column.py b/tests/unit/test_column.py
index 8c3c645..10f3059 100644
--- a/tests/unit/test_column.py
+++ b/tests/unit/test_column.py
@@ -14,6 +14,12 @@ from dbt.adapters.bigquery.column import get_nested_column_data_types
             None,
             {"a": {"name": "a", "data_type": "string"}},
         ),
+        # Flat column - missing data_type
+        (
+            {"a": {"name": "a"}},
+            None,
+            {"a": {"name": "a", "data_type": None}},
+        ),
         # Flat column - with constraints
         (
             {"a": {"name": "a", "data_type": "string"}},
@@ -32,18 +38,75 @@ from dbt.adapters.bigquery.column import get_nested_column_data_types
             None,
             {"b": {"name": "b", "data_type": "struct<nested string>"}},
         ),
+        # Single nested column, 1 level - missing data_type
+        (
+            {"b.nested": {"name": "b.nested"}},
+            None,
+            {"b": {"name": "b", "data_type": "struct<nested>"}},
+        ),
         # Single nested column, 1 level - with constraints
         (
             {"b.nested": {"name": "b.nested", "data_type": "string"}},
             {"b.nested": "not null"},
             {"b": {"name": "b", "data_type": "struct<nested string not null>"}},
         ),
+        # Single nested column, 1 level - with constraints, missing data_type (constraints not valid without data_type)
+        (
+            {"b.nested": {"name": "b.nested"}},
+            {"b.nested": "not null"},
+            {"b": {"name": "b", "data_type": "struct<nested>"}},
+        ),
         # Single nested column, 1 level - with constraints + other keys
         (
             {"b.nested": {"name": "b.nested", "data_type": "string", "other": "unpreserved"}},
             {"b.nested": "not null"},
             {"b": {"name": "b", "data_type": "struct<nested string not null>"}},
         ),
+        # Single nested column, 1 level - with corresponding parent column
+        (
+            {
+                "b": {"name": "b", "data_type": "struct"},
+                "b.nested": {"name": "b.nested", "data_type": "string"},
+            },
+            None,
+            {"b": {"name": "b", "data_type": "struct<nested string>"}},
+        ),
+        # Single nested column, 1 level - with corresponding parent column specified last
+        (
+            {
+                "b.nested": {"name": "b.nested", "data_type": "string"},
+                "b": {"name": "b", "data_type": "struct"},
+            },
+            None,
+            {"b": {"name": "b", "data_type": "struct<nested string>"}},
+        ),
+        # Single nested column, 1 level - with corresponding parent column + parent constraint
+        (
+            {
+                "b": {"name": "b", "data_type": "struct"},
+                "b.nested": {"name": "b.nested", "data_type": "string"},
+            },
+            {"b": "not null"},
+            {"b": {"name": "b", "data_type": "struct<nested string> not null"}},
+        ),
+        # Single nested column, 1 level - with corresponding parent column as array
+        (
+            {
+                "b": {"name": "b", "data_type": "array"},
+                "b.nested": {"name": "b.nested", "data_type": "string"},
+            },
+            None,
+            {"b": {"name": "b", "data_type": "array<struct<nested string>>"}},
+        ),
+        # Single nested column, 1 level - with corresponding parent column as array + constraint
+        (
+            {
+                "b": {"name": "b", "data_type": "array"},
+                "b.nested": {"name": "b.nested", "data_type": "string"},
+            },
+            {"b": "not null"},
+            {"b": {"name": "b", "data_type": "array<struct<nested string>> not null"}},
+        ),
         # Multiple nested columns, 1 level
         (
             {
@@ -106,6 +169,28 @@ from dbt.adapters.bigquery.column import get_nested_column_data_types
                 },
             },
         ),
+        # Nested columns, multiple levels - missing data_type
+        (
+            {
+                "b.user.name.first": {
+                    "name": "b.user.name.first",
+                    "data_type": "string",
+                },
+                "b.user.name.last": {
+                    "name": "b.user.name.last",
+                    "data_type": "string",
+                },
+                "b.user.id": {"name": "b.user.id", "data_type": "int64"},
+                "b.user.country": {"name": "b.user.country"},  # missing data_type
+            },
+            None,
+            {
+                "b": {
+                    "name": "b",
+                    "data_type": "struct<user struct<name struct<first string, last string>, id int64, country>>",
+                },
+            },
+        ),
         # Nested columns, multiple levels - with constraints!
         (
             {
@@ -128,6 +213,32 @@ from dbt.adapters.bigquery.column import get_nested_column_data_types
                 },
             },
         ),
+        # Nested columns, multiple levels - with parent arrays and constraints!
+        (
+            {
+                "b.user.names": {
+                    "name": "b.user.names",
+                    "data_type": "array",
+                },
+                "b.user.names.first": {
+                    "name": "b.user.names.first",
+                    "data_type": "string",
+                },
+                "b.user.names.last": {
+                    "name": "b.user.names.last",
+                    "data_type": "string",
+                },
+                "b.user.id": {"name": "b.user.id", "data_type": "int64"},
+                "b.user.country": {"name": "b.user.country", "data_type": "string"},
+            },
+            {"b.user.names.first": "not null", "b.user.id": "unique"},
+            {
+                "b": {
+                    "name": "b",
+                    "data_type": "struct<user struct<names array<struct<first string not null, last string>>, id int64 unique, country string>>",
+                },
+            },
+        ),
     ],
 )
 def test_get_nested_column_data_types(columns, constraints, expected_nested_columns):