Fix: Avoid type coercion when converting pandas dataframes to agate t…

…ables in the dbt adapter (#2897)
TobikoData · Jul 12, 2024 · 4f1d735 · 4f1d735
1 parent 04db61e
commit 4f1d735
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/sqlmesh/dbt/util.py b/sqlmesh/dbt/util.py
@@ -13,14 +13,14 @@ def _get_dbt_version() -> t.Tuple[int, int]:
 DBT_VERSION = _get_dbt_version()
 
 if DBT_VERSION < (1, 8):
-    from dbt.clients.agate_helper import table_from_data, empty_table, as_matrix  # type: ignore  # noqa: F401
+    from dbt.clients.agate_helper import table_from_data_flat, empty_table, as_matrix  # type: ignore  # noqa: F401
 else:
-    from dbt_common.clients.agate_helper import table_from_data, empty_table, as_matrix  # type: ignore  # noqa: F401
+    from dbt_common.clients.agate_helper import table_from_data_flat, empty_table, as_matrix  # type: ignore  # noqa: F401
 
 
 def pandas_to_agate(df: pd.DataFrame) -> agate.Table:
     """
     Converts a Pandas DataFrame to an Agate Table
     """
 
-    return table_from_data(df.to_dict(orient="records"), df.columns.tolist())
+    return table_from_data_flat(df.to_dict(orient="records"), df.columns.tolist())
diff --git a/tests/dbt/test_util.py b/tests/dbt/test_util.py
@@ -0,0 +1,12 @@
+from __future__ import annotations
+
+import pandas as pd
+from sqlmesh.dbt.util import pandas_to_agate
+
+
+def test_pandas_to_agate_type_coercion():
+    df = pd.DataFrame({"data": ["_2024_01_01", "_2024_01_02", "_2024_01_03"]})
+    agate_rows = pandas_to_agate(df).rows
+
+    values = [v[0] for v in agate_rows.values()]
+    assert values == ["_2024_01_01", "_2024_01_02", "_2024_01_03"]