Repository: TobikoData/sqlmesh
Pull request #3215 (2024-10-03 14:18:41)

Fix: Infer column types when data type is omitted in dbt seeds
This fix addresses an issue with dbt seeds where columns defined solely by a description, or where the data type was specified for only some columns, did not trigger type inference for the rest of the columns. It ensures that explicitly declared data types are set from the declaration while the types of the remaining columns are inferred.

Themiscodes: Thanks for the review @tobymao addressed comments

Patch:
diff --git a/sqlmesh/dbt/seed.py b/sqlmesh/dbt/seed.py
index 9b1ab94f..eaa088d1 100644
--- a/sqlmesh/dbt/seed.py
+++ b/sqlmesh/dbt/seed.py
@@ -41,25 +41,31 @@ class SeedConfig(BaseModelConfig):
         seed_path = self.path.absolute().as_posix()
 
         if column_types := self.column_types:
-            columns = copy.deepcopy(self.columns)
+            column_types_override = copy.deepcopy(self.columns)
             for name, data_type in column_types.items():
-                column = columns.setdefault(name, ColumnConfig(name=name))
+                column = column_types_override.setdefault(name, ColumnConfig(name=name))
                 column.data_type = data_type
                 column.quote = self.quote_columns or column.quote
-                kwargs = self.sqlmesh_model_kwargs(context, columns)
+                kwargs = self.sqlmesh_model_kwargs(context, column_types_override)
         else:
             kwargs = self.sqlmesh_model_kwargs(context)
 
-        if kwargs.get("columns") is None:
+        columns = kwargs.get("columns") or {}
+        descriptions = kwargs.get("column_descriptions") or {}
+        missing_types = (set(descriptions) | set(self.columns)) - set(columns)
+        if not columns or missing_types:
             agate_table = (
                 agate_helper.from_csv(seed_path, [], delimiter=self.delimiter)
                 if SUPPORTS_DELIMITER
                 else agate_helper.from_csv(seed_path, [])
             )
-            kwargs["columns"] = {
+            inferred_types = {
                 name: AGATE_TYPE_MAPPING[tpe.__class__]
                 for name, tpe in zip(agate_table.column_names, agate_table.column_types)
             }
+            for key in missing_types or inferred_types:
+                columns[key] = inferred_types[key]
+            kwargs["columns"] = columns
 
         return create_seed_model(
             self.canonical_name(context),


Test patch:
diff --git a/tests/dbt/test_transformation.py b/tests/dbt/test_transformation.py
index 27767038..1a02cdde 100644
--- a/tests/dbt/test_transformation.py
+++ b/tests/dbt/test_transformation.py
@@ -470,6 +470,52 @@ def test_seed_column_inference(tmp_path):
     }
 
 
+def test_seed_partial_column_inference(tmp_path):
+    seed_csv = tmp_path / "seed.csv"
+    with open(seed_csv, "w", encoding="utf-8") as fd:
+        fd.write("int_col,double_col,datetime_col,boolean_col\n")
+        fd.write("1,1.2,2021-01-01 00:00:00,true\n")
+        fd.write("2,2.3,2021-01-02 00:00:00,false\n")
+        fd.write("null,,null,\n")
+
+    seed = SeedConfig(
+        name="test_model",
+        package="package",
+        path=Path(seed_csv),
+        column_types={
+            "double_col": "double",
+        },
+        columns={
+            "int_col": ColumnConfig(
+                name="int_col", data_type="int", description="Description with type."
+            ),
+            "datetime_col": ColumnConfig(
+                name="datetime_col", description="Description without type."
+            ),
+            "boolean_col": ColumnConfig(name="boolean_col"),
+        },
+    )
+
+    expected_column_types = {
+        "int_col": exp.DataType.build("int"),
+        "double_col": exp.DataType.build("double"),
+        "datetime_col": exp.DataType.build("datetime"),
+        "boolean_col": exp.DataType.build("boolean"),
+    }
+
+    expected_column_descriptions = {
+        "int_col": "Description with type.",
+        "datetime_col": "Description without type.",
+    }
+
+    context = DbtContext()
+    context.project_name = "Foo"
+    context.target = DuckDbConfig(name="target", schema="test")
+    sqlmesh_seed = seed.to_sqlmesh(context)
+    assert sqlmesh_seed.columns_to_types == expected_column_types
+    assert sqlmesh_seed.column_descriptions == expected_column_descriptions
+
+
 def test_agate_integer_cast():
     agate_integer = Integer(null_values=("null", ""))
     assert agate_integer.cast("1") == 1