Repository: TobikoData/sqlmesh (python)

Issue:
Support mutating BigQuery nested types without dropping
This is a feature request for the BigQuery engine.

Currently, sqlmesh supports nested types (i.e.: `STRUCT`) within models, e.g.:

```sql
SELECT STRUCT(foo, bar) AS my_struct FROM underlying
```

It also supports changes to these, e.g.:

```sql
SELECT STRUCT(foo, bar, new_field) AS my_struct FROM underlying
```

But because BigQuery [does not support mutating nested types using SQL](https://cloud.google.com/bigquery/docs/managing-table-schemas#sql_1), the way this is implemented in sqlmesh requires dropping the column and creating a new one (which truncates the data).

The feature request is to fallback on using the [BigQuery API](https://cloud.google.com/bigquery/docs/managing-table-schemas#api_1) directly for doing these mutations. Specifically, it should be possible to add columns to a nested type without dropping it.

The caveat is that adding columns to a nested type can only be done "at the end" (i.e.: `STRUCT(foo, new_field, bar)` is not possible, only `STRUCT(foo, bar, new_field)`).

Ideally, sqlmesh would allow **appending** fields to a `STRUCT` as long as the query also matches the resulting schema.

Pull request description:
Feat: Support BigQuery nested column additions without dropping
This update adds support for adding fields to a nested type in BigQuery without dropping the existing column by using the ([BigQuery API](https://cloud.google.com/bigquery/docs/managing-table-schemas#api_1)), fixes #2321 

Patch:
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index d370c03a..93a9a761 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import logging
 import typing as t
+from collections import defaultdict
 
 import pandas as pd
 from sqlglot import exp
@@ -51,8 +52,6 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
     MAX_TABLE_COMMENT_LENGTH = 1024
     MAX_COLUMN_COMMENT_LENGTH = 1024
 
-    # SQL is not supported for adding columns to structs: https://cloud.google.com/bigquery/docs/managing-table-schemas#api_1
-    # Can explore doing this with the API in the future
     SCHEMA_DIFFER = SchemaDiffer(
         compatible_types={
             exp.DataType.build("INT64", dialect=DIALECT): {
@@ -83,6 +82,8 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
                 exp.DataType.build("BYTES", dialect=DIALECT).this,
             },
         },
+        support_nested_operations=True,
+        support_nested_drop=False,
     )
 
     @property
@@ -224,6 +225,23 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
                 columns["_PARTITIONDATE"] = exp.DataType.build("DATE")
         return columns
 
+    def alter_table(
+        self,
+        alter_expressions: t.List[exp.Alter],
+    ) -> None:
+        """
+        Performs the alter statements to change the current table into the structure of the target table,
+        and uses the API to add columns to structs, where SQL is not supported.
+        """
+
+        nested_fields, non_nested_expressions = self._split_alter_expressions(alter_expressions)
+
+        if nested_fields:
+            self._update_table_schema_nested_fields(nested_fields, alter_expressions[0].this)
+
+        if non_nested_expressions:
+            super().alter_table(non_nested_expressions)
+
     def fetchone(
         self,
         query: t.Union[exp.Expression, str],
@@ -261,6 +279,64 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
         )
         return list(self._query_data)
 
+    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]]:
+        """
+        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)
+        non_nested_expressions = []
+
+        for alter_expression in alter_expressions:
+            action = alter_expression.args["actions"][0]
+            if (
+                isinstance(action, exp.ColumnDef)
+                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))
+            else:
+                non_nested_expressions.append(alter_expression)
+
+        return nested_fields_to_add, non_nested_expressions
+
+    def _update_table_schema_nested_fields(
+        self, nested_fields_to_add: t.Dict[str, list[t.Tuple[str, str]]], table_name: str
+    ) -> None:
+        """
+        Updates a BigQuery table schema by adding the new nested fields provided.
+        """
+        from google.cloud import bigquery
+
+        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]
+                )
+            if current_fields:
+                new_schema.append(
+                    bigquery.SchemaField(
+                        field.name, "RECORD", mode="NULLABLE", fields=tuple(current_fields)
+                    )
+                )
+            else:
+                new_schema.append(field)
+
+        if new_schema != original_schema:
+            table.schema = new_schema
+            self.client.update_table(table, ["schema"])
+
     def __load_pandas_to_table(
         self,
         table: bigquery.Table,
diff --git a/sqlmesh/core/engine_adapter/databricks.py b/sqlmesh/core/engine_adapter/databricks.py
index bd368ada..da77ff35 100644
--- a/sqlmesh/core/engine_adapter/databricks.py
+++ b/sqlmesh/core/engine_adapter/databricks.py
@@ -39,6 +39,7 @@ class DatabricksEngineAdapter(SparkEngineAdapter):
     SCHEMA_DIFFER = SchemaDiffer(
         support_positional_add=True,
         support_nested_operations=True,
+        support_nested_drop=True,
         array_element_selector="element",
         parameterized_type_defaults={
             exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(10, 0), (0,)],
diff --git a/sqlmesh/core/schema_diff.py b/sqlmesh/core/schema_diff.py
index 7b6123c7..a333b60a 100644
--- a/sqlmesh/core/schema_diff.py
+++ b/sqlmesh/core/schema_diff.py
@@ -301,6 +301,8 @@ class SchemaDiffer(PydanticModel):
             specific position in the set of existing columns.
         support_nested_operations: Whether the engine for which the diff is being computed supports modifications to
             nested data types like STRUCTs and ARRAYs.
+        support_nested_drop: Whether the engine for which the diff is being computed supports removing individual
+            columns of nested STRUCTs.
         compatible_types: Types that are compatible and automatically coerced in actions like UNION ALL. Dict key is data
             type, and value is the set of types that are compatible with it.
         support_coercing_compatible_types: Whether or not the engine for which the diff is being computed supports direct
@@ -321,6 +323,7 @@ class SchemaDiffer(PydanticModel):
 
     support_positional_add: bool = False
     support_nested_operations: bool = False
+    support_nested_drop: bool = False
     array_element_selector: str = ""
     compatible_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = {}
     support_coercing_compatible_types: bool = False
@@ -451,6 +454,15 @@ class SchemaDiffer(PydanticModel):
         )
         return operations
 
+    def _requires_drop_alteration(
+        self, current_struct: exp.DataType, new_struct: exp.DataType
+    ) -> bool:
+        for current_pos, current_kwarg in enumerate(current_struct.expressions.copy()):
+            new_pos, _ = self._get_matching_kwarg(current_kwarg, new_struct, current_pos)
+            if new_pos is None:
+                return True
+        return False
+
     def _resolve_drop_operation(
         self,
         parent_columns: t.List[TableAlterColumn],
@@ -523,12 +535,16 @@ class SchemaDiffer(PydanticModel):
         current_type = exp.DataType.build(current_type, copy=False)
         if self.support_nested_operations:
             if new_type.this == current_type.this == exp.DataType.Type.STRUCT:
-                return self._get_operations(
-                    columns,
-                    current_type,
-                    new_type,
-                    root_struct,
-                )
+                if self.support_nested_drop or not self._requires_drop_alteration(
+                    current_type, new_type
+                ):
+                    return self._get_operations(
+                        columns,
+                        current_type,
+                        new_type,
+                        root_struct,
+                    )
+
             if new_type.this == current_type.this == exp.DataType.Type.ARRAY:
                 # Some engines (i.e. Snowflake) don't support defining types on arrays
                 if not new_type.expressions or not current_type.expressions:
@@ -536,12 +552,15 @@ class SchemaDiffer(PydanticModel):
                 new_array_type = new_type.expressions[0]
                 current_array_type = current_type.expressions[0]
                 if new_array_type.this == current_array_type.this == exp.DataType.Type.STRUCT:
-                    return self._get_operations(
-                        columns,
-                        current_array_type,
-                        new_array_type,
-                        root_struct,
-                    )
+                    if self.support_nested_drop or not self._requires_drop_alteration(
+                        current_array_type, new_array_type
+                    ):
+                        return self._get_operations(
+                            columns,
+                            current_array_type,
+                            new_array_type,
+                            root_struct,
+                        )
         if self._is_coerceable_type(current_type, new_type):
             return []
         elif self._is_compatible_type(current_type, new_type):


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index b95d36e2..863e5793 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -474,6 +474,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
             {
                 "support_positional_add": True,
                 "support_nested_operations": True,
+                "support_nested_drop": True,
                 "array_element_selector": "element",
             },
             {
@@ -630,6 +631,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
             {
                 "support_positional_add": True,
                 "support_nested_operations": True,
+                "support_nested_drop": True,
                 "array_element_selector": "element",
             },
             {
@@ -659,6 +661,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
             {
                 "support_positional_add": True,
                 "support_nested_operations": True,
+                "support_nested_drop": True,
                 "array_element_selector": "element",
             },
             {
@@ -715,6 +718,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         (
             {
                 "support_nested_operations": True,
+                "support_nested_drop": True,
                 "array_element_selector": "element",
             },
             {
@@ -772,6 +776,7 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         (
             {
                 "support_nested_operations": True,
+                "support_nested_drop": True,
                 "array_element_selector": "element",
             },
             {
diff --git a/tests/core/test_schema_diff.py b/tests/core/test_schema_diff.py
index f5581c58..43c0cbc4 100644
--- a/tests/core/test_schema_diff.py
+++ b/tests/core/test_schema_diff.py
@@ -407,7 +407,11 @@ def test_schema_diff_calculate_type_transitions():
                     "INT",
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
         ),
         # Remove a column from the end of a struct
         (
@@ -423,7 +427,11 @@ def test_schema_diff_calculate_type_transitions():
                     "INT",
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
         ),
         # Remove a column from the middle of a struct
         (
@@ -439,7 +447,36 @@ def test_schema_diff_calculate_type_transitions():
                     "INT",
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
+        ),
+        # Remove a column from a struct where nested drop is not supported
+        (
+            "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>>",
+            [
+                TableAlterOperation.drop(
+                    [
+                        TableAlterColumn.struct("info"),
+                    ],
+                    expected_table_struct="STRUCT<id INT>",
+                    column_type="STRUCT<col_a INT, col_b INT, col_c INT>",
+                ),
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.struct("info"),
+                    ],
+                    expected_table_struct="STRUCT<id INT, info STRUCT<col_a INT, col_b INT>>",
+                    column_type="STRUCT<col_a INT, col_b INT>",
+                ),
+            ],
+            dict(
+                support_nested_operations=True,
+                support_nested_drop=False,
+            ),
         ),
         # Remove two columns from the start of a struct
         (
@@ -463,7 +500,11 @@ def test_schema_diff_calculate_type_transitions():
                     "INT",
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
         ),
         # Change a column type in a struct
         (
@@ -534,6 +575,35 @@ def test_schema_diff_calculate_type_transitions():
             dict(
                 support_positional_add=True,
                 support_nested_operations=True,
+                support_nested_drop=True,
+                compatible_types={
+                    exp.DataType.build("INT"): {exp.DataType.build("TEXT")},
+                },
+            ),
+        ),
+        # Add, remove and change a column from a struct where nested drop is not supported
+        (
+            "STRUCT<id INT, info STRUCT<col_a INT, col_b INT, col_c INT>>",
+            "STRUCT<id INT, info STRUCT<col_a INT, col_b TEXT, col_d INT>>",
+            [
+                TableAlterOperation.drop(
+                    [
+                        TableAlterColumn.struct("info"),
+                    ],
+                    expected_table_struct="STRUCT<id INT>",
+                    column_type="STRUCT<col_a INT, col_b INT, col_c INT>",
+                ),
+                TableAlterOperation.add(
+                    [
+                        TableAlterColumn.struct("info"),
+                    ],
+                    expected_table_struct="STRUCT<id INT, info STRUCT<col_a INT, col_b TEXT, col_d INT>>",
+                    column_type="STRUCT<col_a INT, col_b TEXT, col_d INT>",
+                ),
+            ],
+            dict(
+                support_nested_operations=True,
+                support_nested_drop=False,
                 compatible_types={
                     exp.DataType.build("INT"): {exp.DataType.build("TEXT")},
                 },
@@ -581,7 +651,11 @@ def test_schema_diff_calculate_type_transitions():
                     position=TableAlterColumnPosition.last("nest_col_a"),
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
         ),
         # #####################
         # # Array Struct Tests
@@ -617,7 +691,11 @@ def test_schema_diff_calculate_type_transitions():
                     "INT",
                 ),
             ],
-            dict(support_positional_add=True, support_nested_operations=True),
+            dict(
+                support_positional_add=True,
+                support_nested_operations=True,
+                support_nested_drop=True,
+            ),
         ),
         # Alter column type in array of structs
         (
@@ -754,6 +832,7 @@ def test_schema_diff_calculate_type_transitions():
             dict(
                 support_positional_add=True,
                 support_nested_operations=True,
+                support_nested_drop=True,
             ),
         ),
         # Type with precision to same type with no precision and no default is DROP/ADD