Fix!: Type inference for BQ engine adapter (#3981)

Author: VaggelisD

sqlmesh/core/engine_adapter/bigquery.py

@@ -231,33 +231,45 @@ def columns(
     ) -> t.Dict[str, exp.DataType]:
         """Fetches column names and types for the target table."""
 
-        def dtype_to_sql(dtype: t.Optional[StandardSqlDataType]) -> str:
+        def dtype_to_sql(
+            dtype: t.Optional[StandardSqlDataType], field: bigquery.SchemaField
+        ) -> str:
             assert dtype
+            assert field
 
             kind = dtype.type_kind
             assert kind
 
             # Not using the enum value to preserve compatibility with older versions
             # of the BigQuery library.
             if kind.name == "ARRAY":
-                return f"ARRAY<{dtype_to_sql(dtype.array_element_type)}>"
+                return f"ARRAY<{dtype_to_sql(dtype.array_element_type, field)}>"
             if kind.name == "STRUCT":
                 struct_type = dtype.struct_type
                 assert struct_type
                 fields = ", ".join(
-                    f"{field.name} {dtype_to_sql(field.type)}" for field in struct_type.fields
+                    f"{struct_field.name} {dtype_to_sql(struct_field.type, nested_field)}"
+                    for struct_field, nested_field in zip(struct_type.fields, field.fields)
                 )
                 return f"STRUCT<{fields}>"
             if kind.name == "TYPE_KIND_UNSPECIFIED":
-                return "JSON"
+                field_type = field.field_type
+
+                if field_type == "RANGE":
+                    # If the field is a RANGE then `range_element_type` should be set to
+                    # one of `"DATE"`, `"DATETIME"` or `"TIMESTAMP"`.
+                    return f"RANGE<{field.range_element_type.element_type}>"
+
+                return field_type
+
             return kind.name
 
         def create_mapping_schema(
             schema: t.Sequence[bigquery.SchemaField],
         ) -> t.Dict[str, exp.DataType]:
             return {
                 field.name: exp.DataType.build(
-                    dtype_to_sql(field.to_standard_sql().type), dialect=self.dialect
+                    dtype_to_sql(field.to_standard_sql().type, field), dialect=self.dialect
                 )
                 for field in schema
             }

tests/core/engine_adapter/integration/test_integration_bigquery.py

@@ -369,3 +369,37 @@ def test_compare_nested_values_in_table_diff(ctx: TestContext):
 
     ctx.engine_adapter.drop_table(src_table)
     ctx.engine_adapter.drop_table(target_table)
+
+
+def test_column_types(ctx: TestContext):
+    model_name = ctx.table("test")
+    sqlmesh = ctx.create_context()
+
+    sqlmesh.upsert_model(
+        load_sql_based_model(
+            d.parse(
+                f"""
+                MODEL (
+                    name {model_name},
+                );
+
+                SELECT
+                    RANGE('01-01-1900'::DATE, '01-01-1902'::DATE) AS col1,
+                    JSON '{{"id": 10}}' AS col2,
+                    STRUCT([PARSE_JSON('{{"id": 10}}')] AS arr) AS col3;
+                """
+            )
+        )
+    )
+
+    sqlmesh.plan(auto_apply=True, no_prompts=True)
+
+    columns = sqlmesh.engine_adapter.columns(model_name)
+
+    assert columns["col1"].is_type("RANGE<DATE>")
+    assert columns["col2"].is_type("JSON")
+
+    col3 = columns["col3"]
+    coldef = col3.find(exp.ColumnDef)
+    assert col3.is_type("STRUCT")
+    assert coldef and coldef.kind and coldef.kind.is_type("ARRAY<JSON>")