Repository: TobikoData/sqlmesh
Pull request #3190 (2024-09-26 16:10:09)

Fix: Bigquery support of complex nested types
Fix issue of constructing schema fields for nested and array types, by using the function `__get_bq_schemafield` from https://github.com/TobikoData/sqlmesh/pull/3185 when constructing a record or array schema field and the function `_build_nested_fields` to recursively build nested fields, when appending to a record
fixes: #3184 

plaflamme: I don't think this will work if the field being modified is `ARRAY<STRUCT<>>`, I believe [this branch](https://github.com/TobikoData/sqlmesh/blob/340d493134f372fb2cfee2a7d541c64e26ec7d78/sqlmesh/core/engine_adapter/bigquery.py#L327-L332) needs to deal with this case.

Also, I was hoping some unit tests would be added as part of this work, the cases we have are:

* modifying a nested `STRUCT<>`
* modifying a nested `ARRAY<STRUCT<...>>` (this is very common in our models)


Themiscodes: You're right @plaflamme `ARRAY<STRUCT<...>>` could be fixed by passing the field.mode instead of nullable in the branch you mentioned, but going over two levels of nesting (e.g. `STRUCT<ARRAY<STRUCT<ARRAY<...>>>>`) would still not be able to get resolved correctly. I will look into how to handle these cases as well and come up with a fix and tests for them. Thanks for your example use cases they're very helpful!
plaflamme: @Themiscodes Yeah, this is why I refrained from trying to make the fix, it's a bit tricky and I wasn't sure I understood the current code fully.

But from what I understand, if the table looks like this:

```
id INT64,
complex STRUCT<foo ARRAY<STRUCT<bar STRING>>>
```

and we try to alter the nested `foo` to something like this:

```
id INT64,
complex STRUCT<foo ARRAY<STRUCT<bar STRING, baz INT64>>>
```

Then, as long as we run the new method from the top-level `complex` field, it _should_ just work. The part that wasn't clear to me is what's in the `nested_fields_to_add` argument. If that contains `foo`, then yeah, this is much harder, but if it contains `complex`, then it should just work.

Patch:
diff --git a/sqlmesh/core/dialect.py b/sqlmesh/core/dialect.py
index de895b69..9f44199a 100644
--- a/sqlmesh/core/dialect.py
+++ b/sqlmesh/core/dialect.py
@@ -1166,7 +1166,7 @@ def _unquote_schema(schema: t.Dict) -> t.Dict:
 
 @contextmanager
 def normalize_and_quote(
-    query: E, dialect: str, default_catalog: t.Optional[str], quote: bool = True
+    query: E, dialect: DialectType, default_catalog: t.Optional[str], quote: bool = True
 ) -> t.Iterator[E]:
     qualify_tables(query, catalog=default_catalog, dialect=dialect)
     normalize_identifiers(query, dialect=dialect)
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index 8861199a..20753d81 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -37,6 +37,10 @@ if t.TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+NestedField = t.Tuple[str, str, t.List[str]]
+NestedFieldsDict = t.Dict[str, t.List[NestedField]]
+
+
 @set_catalog()
 class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin):
     """
@@ -282,11 +286,11 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin):
     def _split_alter_expressions(
         self,
         alter_expressions: t.List[exp.Alter],
-    ) -> t.Tuple[t.Dict[str, list[t.Tuple[str, str]]], t.List[exp.Alter]]:
+    ) -> t.Tuple[NestedFieldsDict, t.List[exp.Alter]]:
         """
         Returns a dictionary of the nested fields to add and a list of the non-nested alter expressions.
         """
-        nested_fields_to_add: t.Dict[str, list[t.Tuple[str, str]]] = defaultdict(list)
+        nested_fields_to_add: NestedFieldsDict = defaultdict(list)
         non_nested_expressions = []
 
         for alter_expression in alter_expressions:
@@ -296,17 +300,58 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin):
                 and isinstance(action.this, exp.Dot)
                 and isinstance(action.kind, exp.DataType)
             ):
-                record = action.this.this.sql(dialect="bigquery")
-                field = action.this.expression.sql(dialect="bigquery")
-                data_type = action.kind.sql(dialect="bigquery")
-                nested_fields_to_add[record].append((data_type, field))
+                root_field, *leaf_fields = action.this.this.sql(dialect=self.dialect).split(".")
+                new_field = action.this.expression.sql(dialect=self.dialect)
+                data_type = action.kind.sql(dialect=self.dialect)
+                nested_fields_to_add[root_field].append((new_field, data_type, leaf_fields))
             else:
                 non_nested_expressions.append(alter_expression)
 
         return nested_fields_to_add, non_nested_expressions
 
+    def _build_nested_fields(
+        self,
+        current_fields: t.List[bigquery.SchemaField],
+        fields_to_add: t.List[NestedField],
+    ) -> t.List[bigquery.SchemaField]:
+        """
+        Recursively builds and updates the schema fields with the new nested fields.
+        """
+        from google.cloud import bigquery
+
+        new_fields = []
+        root: t.List[t.Tuple[str, str]] = []
+        leaves: NestedFieldsDict = defaultdict(list)
+        for new_field, data_type, leaf_fields in fields_to_add:
+            if leaf_fields:
+                leaves[leaf_fields[0]].append((new_field, data_type, leaf_fields[1:]))
+            else:
+                root.append((new_field, data_type))
+
+        for field in current_fields:
+            # If the new fields are nested, we need to recursively build them
+            if field.name in leaves:
+                subfields = list(field.fields)
+                subfields = self._build_nested_fields(subfields, leaves[field.name])
+                new_fields.append(
+                    bigquery.SchemaField(
+                        field.name, "RECORD", mode=field.mode, fields=tuple(subfields)
+                    )
+                )
+            else:
+                new_fields.append(field)
+
+        # Build and append the new root-level fields
+        new_fields.extend(
+            self.__get_bq_schemafield(
+                new_field[0], exp.DataType.build(new_field[1], dialect=self.dialect)
+            )
+            for new_field in root
+        )
+        return new_fields
+
     def _update_table_schema_nested_fields(
-        self, nested_fields_to_add: t.Dict[str, list[t.Tuple[str, str]]], table_name: str
+        self, nested_fields_to_add: NestedFieldsDict, table_name: str
     ) -> None:
         """
         Updates a BigQuery table schema by adding the new nested fields provided.
@@ -316,18 +361,17 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin):
         table = self._get_table(table_name)
         original_schema = table.schema
         new_schema = []
-
         for field in original_schema:
-            current_fields = list(field.fields)
             if field.name in nested_fields_to_add:
-                current_fields.extend(
-                    bigquery.SchemaField(new_field[1], new_field[0], mode="NULLABLE")
-                    for new_field in nested_fields_to_add[field.name]
+                fields = self._build_nested_fields(
+                    list(field.fields), nested_fields_to_add[field.name]
                 )
-            if current_fields:
                 new_schema.append(
                     bigquery.SchemaField(
-                        field.name, "RECORD", mode="NULLABLE", fields=tuple(current_fields)
+                        field.name,
+                        "RECORD",
+                        mode=field.mode,
+                        fields=tuple(fields),
                     )
                 )
             else:
diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index ea3b5557..975b81c6 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -44,6 +44,7 @@ from sqlmesh.utils.jinja import JinjaMacroRegistry, has_jinja
 from sqlmesh.utils.metaprogramming import Executable, prepare_env, print_exception
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.engine_adapter import EngineAdapter
     from sqlmesh.core.snapshot import Snapshot
@@ -144,7 +145,7 @@ class MacroEvaluator:
 
     def __init__(
         self,
-        dialect: str = "",
+        dialect: DialectType = "",
         python_env: t.Optional[t.Dict[str, Executable]] = None,
         jinja_env: t.Optional[Environment] = None,
         schema: t.Optional[MappingSchema] = None,
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 592bb63f..46ae495f 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -46,6 +46,7 @@ from sqlmesh.utils.metaprogramming import (
 )
 
 if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.audit import ModelAudit, Audit
     from sqlmesh.core.context import ExecutionContext
@@ -1600,24 +1601,15 @@ def load_sql_based_model(
         if prop.name.lower() == "audits":
             model_audits = prop.args.get("value")
 
-    meta_python_env = _python_env(
-        expressions=meta,
-        jinja_macro_references=None,
+    meta_renderer = _meta_renderer(
+        expression=meta,
         module_path=module_path,
-        macros=macros or macro.get_registry(),
+        macros=macros,
+        jinja_macros=jinja_macros,
         variables=variables,
         path=path,
-    )
-    meta_renderer = ExpressionRenderer(
-        meta,
-        dialect,
-        [],
-        path=path,
-        jinja_macro_registry=jinja_macros,
-        python_env=meta_python_env,
+        dialect=dialect,
         default_catalog=default_catalog,
-        quote_identifiers=False,
-        normalize_identifiers=False,
     )
 
     rendered_meta_exprs = meta_renderer.render()
@@ -1954,6 +1946,19 @@ def create_python_model(
             )
         )
 
+    dialect = kwargs.get("dialect")
+    name_renderer = _meta_renderer(
+        expression=d.parse_one(name, dialect=dialect),
+        module_path=module_path,
+        macros=macros,
+        jinja_macros=jinja_macros,
+        variables=variables,
+        path=path,
+        dialect=dialect,
+        default_catalog=kwargs.get("default_catalog"),
+    )
+    name = t.cast(t.List[exp.Expression], name_renderer.render())[0].sql(dialect=dialect)
+
     parsed_depends_on, referenced_variables = (
         _parse_dependencies(python_env, entrypoint) if python_env is not None else (set(), set())
     )
@@ -2385,6 +2390,37 @@ def _refs_to_sql(values: t.Any) -> exp.Expression:
     return exp.Tuple(expressions=values)
 
 
+def _meta_renderer(
+    expression: exp.Expression,
+    module_path: Path,
+    path: Path,
+    jinja_macros: t.Optional[JinjaMacroRegistry] = None,
+    macros: t.Optional[MacroRegistry] = None,
+    dialect: DialectType = None,
+    variables: t.Optional[t.Dict[str, t.Any]] = None,
+    default_catalog: t.Optional[str] = None,
+) -> ExpressionRenderer:
+    meta_python_env = _python_env(
+        expressions=expression,
+        jinja_macro_references=None,
+        module_path=module_path,
+        macros=macros or macro.get_registry(),
+        variables=variables,
+        path=path,
+    )
+    return ExpressionRenderer(
+        expression,
+        dialect,
+        [],
+        path=path,
+        jinja_macro_registry=jinja_macros,
+        python_env=meta_python_env,
+        default_catalog=default_catalog,
+        quote_identifiers=False,
+        normalize_identifiers=False,
+    )
+
+
 META_FIELD_CONVERTER: t.Dict[str, t.Callable] = {
     "start": lambda value: exp.Literal.string(value),
     "cron": lambda value: exp.Literal.string(value),
diff --git a/sqlmesh/core/plan/builder.py b/sqlmesh/core/plan/builder.py
index 8dcf5b96..48eb7ede 100644
--- a/sqlmesh/core/plan/builder.py
+++ b/sqlmesh/core/plan/builder.py
@@ -532,11 +532,18 @@ class PlanBuilder:
                 # In case of the forward only plan any modifications result in reuse of the
                 # previous version for non-seed models.
                 # New snapshots of seed models are considered non-breaking ones.
-                snapshot.categorize_as(
+                category = (
                     SnapshotChangeCategory.NON_BREAKING
                     if snapshot.is_seed
                     else SnapshotChangeCategory.FORWARD_ONLY
                 )
+                # If the model kind changes mark as breaking
+                if snapshot.is_model and snapshot.name in self._context_diff.modified_snapshots:
+                    _, old = self._context_diff.modified_snapshots[snapshot.name]
+                    if old.model.kind.name != snapshot.model.kind.name:
+                        category = SnapshotChangeCategory.BREAKING
+
+                snapshot.categorize_as(category)
             elif s_id.name in self._context_diff.modified_snapshots:
                 self._categorize_snapshot(snapshot, dag, indirectly_modified)
 
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index c466ed62..46ef207b 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -28,6 +28,7 @@ from sqlmesh.utils.metaprogramming import Executable, prepare_env
 
 if t.TYPE_CHECKING:
     from sqlglot._typing import E
+    from sqlglot.dialects.dialect import DialectType
 
     from sqlmesh.core.snapshot import DeployabilityIndex, Snapshot
 
@@ -39,7 +40,7 @@ class BaseExpressionRenderer:
     def __init__(
         self,
         expression: exp.Expression,
-        dialect: str,
+        dialect: DialectType,
         macro_definitions: t.List[d.MacroDef],
         path: Path = Path(),
         jinja_macro_registry: t.Optional[JinjaMacroRegistry] = None,


Test patch:
diff --git a/tests/core/engine_adapter/test_bigquery.py b/tests/core/engine_adapter/test_bigquery.py
index da7d648b..20d7c711 100644
--- a/tests/core/engine_adapter/test_bigquery.py
+++ b/tests/core/engine_adapter/test_bigquery.py
@@ -805,3 +805,94 @@ def test_materialized_view_properties(
     assert sql_calls == [
         "CREATE OR REPLACE MATERIALIZED VIEW `test_table` PARTITION BY `ds` CLUSTER BY `a` AS SELECT 1",
     ]
+
+
+def test_nested_fields_update(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(BigQueryEngineAdapter)
+
+    current_schema = [
+        bigquery.SchemaField(
+            "user",
+            "RECORD",
+            "NULLABLE",
+            fields=(
+                bigquery.SchemaField("name", "STRING", "NULLABLE"),
+                bigquery.SchemaField(
+                    "orders",
+                    "RECORD",
+                    "REPEATED",
+                    fields=([bigquery.SchemaField("id", "INT64", "NULLABLE")]),
+                ),
+            ),
+        )
+    ]
+    new_nested_fields = [("year", "INT64", ["user", "orders"]), ("active", "BOOL", ["user"])]
+    expected = [
+        bigquery.SchemaField(
+            "user",
+            "RECORD",
+            "NULLABLE",
+            fields=(
+                bigquery.SchemaField("name", "STRING", "NULLABLE"),
+                bigquery.SchemaField(
+                    "orders",
+                    "RECORD",
+                    "REPEATED",
+                    fields=(
+                        bigquery.SchemaField("id", "INT64", "NULLABLE"),
+                        bigquery.SchemaField("year", "INT64", "NULLABLE"),
+                    ),
+                ),
+                bigquery.SchemaField("active", "BOOL", "NULLABLE"),
+            ),
+        )
+    ]
+    assert adapter._build_nested_fields(current_schema, new_nested_fields) == expected
+
+    current_schema = [
+        bigquery.SchemaField(
+            "users",
+            "RECORD",
+            "REPEATED",
+            fields=(
+                [
+                    bigquery.SchemaField(
+                        "user",
+                        "RECORD",
+                        "NULLABLE",
+                        fields=(bigquery.SchemaField("name", "STRING", "NULLABLE"),),
+                    )
+                ]
+            ),
+        )
+    ]
+    new_nested_fields = [
+        ("orders", "ARRAY<INT64>", ["users", "user"]),
+        ("tags", "STRING", ["users"]),
+        ("details", "ARRAY<STRING>", []),
+    ]
+    expected = [
+        bigquery.SchemaField(
+            "users",
+            "RECORD",
+            "REPEATED",
+            fields=(
+                bigquery.SchemaField(
+                    "user",
+                    "RECORD",
+                    "NULLABLE",
+                    fields=(
+                        bigquery.SchemaField("name", "STRING", "NULLABLE"),
+                        bigquery.SchemaField("orders", "INT64", "REPEATED"),
+                    ),
+                ),
+                bigquery.SchemaField(
+                    "tags",
+                    "STRING",
+                    "NULLABLE",
+                ),
+            ),
+        ),
+        bigquery.SchemaField("details", "STRING", "REPEATED"),
+    ]
+    assert adapter._build_nested_fields(current_schema, new_nested_fields) == expected
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index 8dd74f06..b62abf4e 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -4565,7 +4565,7 @@ def test_variables_jinja():
 
 def test_variables_python_model(mocker: MockerFixture) -> None:
     @model(
-        "test_variables_python_model",
+        "foo_@{bar}",
         kind="full",
         columns={"a": "string", "b": "string", "c": "string"},
     )
@@ -4580,12 +4580,13 @@ def test_variables_python_model(mocker: MockerFixture) -> None:
             ]
         )
 
-    python_model = model.get_registry()["test_variables_python_model"].model(
+    python_model = model.get_registry()["foo_@{bar}"].model(
         module_path=Path("."),
         path=Path("."),
-        variables={"test_var_a": "test_value", "test_var_unused": 2},
+        variables={"test_var_a": "test_value", "test_var_unused": 2, "bar": "suffix"},
     )
 
+    assert python_model.name == "foo_suffix"
     assert python_model.python_env[c.SQLMESH_VARS] == Executable.value({"test_var_a": "test_value"})
 
     context = ExecutionContext(mocker.Mock(), {}, None, None)
diff --git a/tests/core/test_plan.py b/tests/core/test_plan.py
index 664991ec..6675f223 100644
--- a/tests/core/test_plan.py
+++ b/tests/core/test_plan.py
@@ -245,6 +245,49 @@ def test_forward_only_plan_added_models(make_snapshot, mocker: MockerFixture):
     assert snapshot_b.change_category == SnapshotChangeCategory.BREAKING
 
 
+def test_forward_only_plan_categorizes_change_model_kind_as_breaking(
+    make_snapshot, mocker: MockerFixture
+):
+    snapshot_old = make_snapshot(
+        SqlModel(
+            name="a",
+            dialect="duckdb",
+            query=parse_one("select 1, ds"),
+            kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+        )
+    )
+
+    # Simulate a change in the model kind.
+    updated_snapshot = make_snapshot(
+        SqlModel(
+            **{
+                **snapshot_old.model.dict(),
+                "kind": dict(name=ModelKindName.VIEW),
+            }
+        )
+    )
+
+    context_diff = ContextDiff(
+        environment="prod",
+        is_new_environment=False,
+        is_unfinalized_environment=False,
+        normalize_environment_name=True,
+        create_from="prod",
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={updated_snapshot.name: (updated_snapshot, snapshot_old)},
+        snapshots={updated_snapshot.snapshot_id: updated_snapshot},
+        new_snapshots={updated_snapshot.snapshot_id: updated_snapshot},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+    )
+
+    PlanBuilder(context_diff, DuckDBEngineAdapter.SCHEMA_DIFFER, forward_only=True).build()
+
+    assert updated_snapshot.change_category == SnapshotChangeCategory.BREAKING
+
+
 def test_paused_forward_only_parent(make_snapshot, mocker: MockerFixture):
     snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
     snapshot_a.previous_versions = (
diff --git a/tests/core/test_schema_diff.py b/tests/core/test_schema_diff.py
index 43c0cbc4..153e3987 100644
--- a/tests/core/test_schema_diff.py
+++ b/tests/core/test_schema_diff.py
@@ -393,6 +393,29 @@ def test_schema_diff_calculate_type_transitions():
             ],
             dict(support_positional_add=True, support_nested_operations=True),
         ),
+        # Add columns in different levels of nesting of structs
+        (
+            "STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT>>",
+            "STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT, col_d INT>, txt TEXT>",
+            [
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.primitive("txt"),
+                    ],
+                    "TEXT",
+                    expected_table_struct="STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT>, txt TEXT>",
+                ),
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.struct("info"),
+                        TableAlterColumn.primitive("col_d"),
+                    ],
+                    "INT",
+                    expected_table_struct="STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT, col_d INT>, txt TEXT>",
+                ),
+            ],
+            dict(support_positional_add=False, support_nested_operations=True),
+        ),
         # Remove a column from the start of a struct
         (
             "STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT>>",
@@ -721,6 +744,29 @@ def test_schema_diff_calculate_type_transitions():
                 },
             ),
         ),
+        # Add columns to struct of array within different nesting levels
+        (
+            "STRUCT<id INT, infos ARRAY<STRUCT<col_a INT, col_b INT, col_c INT>>>",
+            "STRUCT<id INT, infos ARRAY<STRUCT<col_a INT, col_b INT, col_c INT, col_d INT >>, col_e INT>",
+            [
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.primitive("col_e"),
+                    ],
+                    "INT",
+                    expected_table_struct="STRUCT<id INT, infos ARRAY<STRUCT<col_a INT, col_b INT, col_c INT>>, col_e INT>",
+                ),
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.array_of_struct("infos"),
+                        TableAlterColumn.primitive("col_d"),
+                    ],
+                    "INT",
+                    expected_table_struct="STRUCT<id INT, infos ARRAY<STRUCT<col_a INT, col_b INT,  col_c INT, col_d INT>>, col_e INT>",
+                ),
+            ],
+            dict(support_positional_add=False, support_nested_operations=True),
+        ),
         # Add an array of primitives
         (
             "STRUCT<id INT, infos ARRAY<STRUCT<col_a INT, col_b INT, col_c INT>>>",