Feat!: Infer column types in the same manner as dbt does when converting dbt seeds (#2328)

Author: izeigerman

setup.py

@@ -63,6 +63,7 @@
             "beautifulsoup4",
             "black==24.1.1",
             "cryptography~=41.0.7",
+            "dbt-common",
             "dbt-core",
             "dbt-duckdb>=1.7.1",
             "Faker",
@@ -113,6 +114,7 @@
         ],
         "dbt": [
             "dbt-core<2",
+            "dbt-common",
         ],
         "gcppostgres": [
             "cloud-sql-python-connector[pg8000]",

sqlmesh/dbt/seed.py

@@ -2,6 +2,10 @@
 
 import typing as t
 
+import agate
+from dbt_common.clients import agate_helper
+from sqlglot import exp
+
 from sqlmesh.core.model import Model, SeedKind, create_seed_model
 from sqlmesh.dbt.basemodel import BaseModelConfig
 
@@ -18,11 +22,51 @@ class SeedConfig(BaseModelConfig):
     General propreties, General configs, and For seeds sections.
     """
 
+    delimiter: str = ","
+
     def to_sqlmesh(self, context: DbtContext) -> Model:
         """Converts the dbt seed into a SQLMesh model."""
+        seed_path = self.path.absolute().as_posix()
+        kwargs = self.sqlmesh_model_kwargs(context)
+        if kwargs.get("columns") is None:
+            agate_table = agate_helper.from_csv(seed_path, [], delimiter=self.delimiter)
+            kwargs["columns"] = {
+                name: AGATE_TYPE_MAPPING[tpe.__class__]
+                for name, tpe in zip(agate_table.column_names, agate_table.column_types)
+            }
+
         return create_seed_model(
             self.canonical_name(context),
-            SeedKind(path=self.path.absolute().as_posix()),
+            SeedKind(path=seed_path),
             dialect=context.dialect,
-            **self.sqlmesh_model_kwargs(context),
+            **kwargs,
         )
+
+
+class Integer(agate.data_types.DataType):
+    def cast(self, d: str) -> t.Optional[int]:
+        if d is None:
+            return d
+        try:
+            return int(d)
+        except ValueError:
+            raise agate.exceptions.CastError('Can not parse value "%s" as Integer.' % d)
+
+    def jsonify(self, d: str) -> str:
+        return d
+
+
+# The dbt version has a bug in which they check whether the type of the input value
+# is int, while the input value is actually always a string.
+agate_helper.Integer = Integer  # type: ignore
+
+
+AGATE_TYPE_MAPPING = {
+    agate_helper.Integer: exp.DataType.build("int"),
+    agate_helper.Number: exp.DataType.build("double"),
+    agate_helper.ISODateTime: exp.DataType.build("datetime"),
+    agate.Date: exp.DataType.build("date"),
+    agate.DateTime: exp.DataType.build("datetime"),
+    agate.Boolean: exp.DataType.build("boolean"),
+    agate.Text: exp.DataType.build("text"),
+}

tests/dbt/test_transformation.py

@@ -294,6 +294,33 @@ def test_seed_columns():
     assert sqlmesh_seed.column_descriptions == expected_column_descriptions
 
 
+def test_seed_column_inference(tmp_path):
+    seed_csv = tmp_path / "seed.csv"
+    with open(seed_csv, "w") as fd:
+        fd.write("int_col,double_col,datetime_col,date_col,boolean_col,text_col\n")
+        fd.write("1,1.2,2021-01-01 00:00:00,2021-01-01,true,foo\n")
+        fd.write("2,2.3,2021-01-02 00:00:00,2021-01-02,false,bar\n")
+
+    seed = SeedConfig(
+        name="test_model",
+        package="package",
+        path=Path(seed_csv),
+    )
+
+    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 == {
+        "int_col": exp.DataType.build("int"),
+        "double_col": exp.DataType.build("double"),
+        "datetime_col": exp.DataType.build("datetime"),
+        "date_col": exp.DataType.build("date"),
+        "boolean_col": exp.DataType.build("boolean"),
+        "text_col": exp.DataType.build("text"),
+    }
+
+
 @pytest.mark.xdist_group("dbt_manifest")
 @pytest.mark.parametrize(
     "model_fqn", ['"memory"."sushi"."waiters"', '"memory"."sushi"."waiter_names"']