Repository: TobikoData/sqlmesh
Pull request #2897 (2024-07-12 02:00:43)

Fix: Avoid type coercion when converting pandas dataframes to agate tables in the dbt adapter
The previous method was coercing string values into more specific types which was messing up dynamic queries in dbt models for some users.


Patch:
diff --git a/docs/integrations/engines/mysql.md b/docs/integrations/engines/mysql.md
index acd0d660..77bd96d4 100644
--- a/docs/integrations/engines/mysql.md
+++ b/docs/integrations/engines/mysql.md
@@ -35,11 +35,15 @@ The operator requires an [Airflow connection](https://airflow.apache.org/docs/ap
 
 By default, the connection ID is set to `mysql_default`, but can be overridden using the `engine_operator_args` parameter to the `SQLMeshAirflow` instance as in the example below:
 ```python linenums="1"
+from sqlmesh.schedulers.airflow import NO_DEFAULT_CATALOG
+
 sqlmesh_airflow = SQLMeshAirflow(
     "mysql",
-    default_catalog="<database name>",
+    default_catalog=NO_DEFAULT_CATALOG,
     engine_operator_args={
         "mysql_conn_id": "<Connection ID>"
     },
 )
-```
\ No newline at end of file
+```
+
+Note: `NO_DEFAULT_CATALOG` is required for MySQL since MySQL doesn't support catalogs. 
\ No newline at end of file
diff --git a/sqlmesh/core/config/scheduler.py b/sqlmesh/core/config/scheduler.py
index 14159f65..b8bf1c1f 100644
--- a/sqlmesh/core/config/scheduler.py
+++ b/sqlmesh/core/config/scheduler.py
@@ -192,7 +192,6 @@ class _BaseAirflowSchedulerConfig(_EngineAdapterStateSyncSchedulerConfig):
         )
 
     def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
-        # The default catalog must still be set on the Airflow side.
         default_catalog = self.get_client(context.console).default_catalog
         return self.default_catalog_override or default_catalog
 
@@ -368,7 +367,6 @@ class MWAASchedulerConfig(_EngineAdapterStateSyncSchedulerConfig, BaseConfig):
         )
 
     def get_default_catalog(self, context: GenericContext) -> t.Optional[str]:
-        # The default catalog must still be set on the Airflow side.
         default_catalog = self.get_client(context.console).default_catalog
         return self.default_catalog_override or default_catalog
 
diff --git a/sqlmesh/dbt/util.py b/sqlmesh/dbt/util.py
index a06f098d..8fc6c6ec 100644
--- a/sqlmesh/dbt/util.py
+++ b/sqlmesh/dbt/util.py
@@ -13,9 +13,9 @@ 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:
@@ -23,4 +23,4 @@ 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/sqlmesh/schedulers/airflow/__init__.py b/sqlmesh/schedulers/airflow/__init__.py
index e69de29b..f5c62217 100644
--- a/sqlmesh/schedulers/airflow/__init__.py
+++ b/sqlmesh/schedulers/airflow/__init__.py
@@ -0,0 +1,1 @@
+NO_DEFAULT_CATALOG = "NO_DEFAULT_CATALOG"
diff --git a/sqlmesh/schedulers/airflow/client.py b/sqlmesh/schedulers/airflow/client.py
index e0d93be0..68ff64d0 100644
--- a/sqlmesh/schedulers/airflow/client.py
+++ b/sqlmesh/schedulers/airflow/client.py
@@ -14,7 +14,7 @@ from sqlmesh.core.snapshot import Snapshot, SnapshotId
 from sqlmesh.core.snapshot.definition import Interval
 from sqlmesh.core.state_sync import Versions
 from sqlmesh.core.user import User
-from sqlmesh.schedulers.airflow import common
+from sqlmesh.schedulers.airflow import common, NO_DEFAULT_CATALOG
 from sqlmesh.utils import unique
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import (
@@ -45,12 +45,14 @@ class BaseAirflowClient(abc.ABC):
         self._console = console
 
     @property
-    def default_catalog(self) -> str:
+    def default_catalog(self) -> t.Optional[str]:
         default_catalog = self.get_variable(common.DEFAULT_CATALOG_VARIABLE_NAME)
         if not default_catalog:
             raise SQLMeshError(
                 "Must define `default_catalog` when creating `SQLMeshAirflow` object. See docs for more info: https://sqlmesh.readthedocs.io/en/stable/integrations/airflow/#airflow-cluster-configuration"
             )
+        if default_catalog == NO_DEFAULT_CATALOG:
+            return None
         return default_catalog
 
     def print_tracking_url(self, dag_id: str, dag_run_id: str, op_name: str) -> None:
diff --git a/sqlmesh/schedulers/airflow/operators/targets.py b/sqlmesh/schedulers/airflow/operators/targets.py
index cecd9809..b9e924b1 100644
--- a/sqlmesh/schedulers/airflow/operators/targets.py
+++ b/sqlmesh/schedulers/airflow/operators/targets.py
@@ -17,7 +17,7 @@ from sqlmesh.core.snapshot import (
     SnapshotTableInfo,
 )
 from sqlmesh.engines import commands
-from sqlmesh.schedulers.airflow import common, util
+from sqlmesh.schedulers.airflow import common, util, NO_DEFAULT_CATALOG
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.pydantic import PydanticModel
 
@@ -30,8 +30,11 @@ class BaseTarget(abc.ABC, t.Generic[CP]):
     ddl_concurrent_tasks: int
 
     @property
-    def default_catalog(self) -> str:
-        return Variable.get(common.DEFAULT_CATALOG_VARIABLE_NAME)
+    def default_catalog(self) -> t.Optional[str]:
+        default_catalog = Variable.get(common.DEFAULT_CATALOG_VARIABLE_NAME)
+        if default_catalog == NO_DEFAULT_CATALOG:
+            return None
+        return default_catalog
 
     def serialized_command_payload(self, context: Context) -> str:
         """Returns the serialized command payload for the Spark application.
diff --git a/sqlmesh/schedulers/airflow/plugin.py b/sqlmesh/schedulers/airflow/plugin.py
index ac64d116..3f7e303e 100644
--- a/sqlmesh/schedulers/airflow/plugin.py
+++ b/sqlmesh/schedulers/airflow/plugin.py
@@ -9,9 +9,11 @@ from airflow.models import Variable
 from airflow.plugins_manager import AirflowPlugin
 
 from sqlmesh.core import constants as c
-from sqlmesh.schedulers.airflow import util
+from sqlmesh.schedulers.airflow import util, NO_DEFAULT_CATALOG
 from sqlmesh.schedulers.airflow.api import sqlmesh_api_v1
-from sqlmesh.schedulers.airflow.common import DEFAULT_CATALOG_VARIABLE_NAME
+from sqlmesh.schedulers.airflow.common import (
+    DEFAULT_CATALOG_VARIABLE_NAME,
+)
 from sqlmesh.utils.errors import SQLMeshError
 
 logger = logging.getLogger(__name__)
@@ -28,14 +30,21 @@ class SqlmeshAirflowPlugin(AirflowPlugin):
             logger.info("MWAA Webserver instance detected. Skipping SQLMesh state migration...")
             return
 
-        default_catalog = Variable.get(DEFAULT_CATALOG_VARIABLE_NAME, default_var=None)
+        # We want to different an expected None default catalog (where the user set `NO_DEFAULT_CATALOG`)
+        # and where the default catalog is not set at all.
+        default_catalog = Variable.get(
+            DEFAULT_CATALOG_VARIABLE_NAME, default_var="MISSING_REQUIRED_CATALOG"
+        )
+        if default_catalog == NO_DEFAULT_CATALOG:
+            # If the user explicitly set `NO_DEFAULT_CATALOG` we want to set the default catalog to None.
+            default_catalog = None
 
         with util.scoped_state_sync() as state_sync:
             try:
-                # If default catalog is not defined then we want to raise an error unless
+                # If default catalog is required but missing (and not explicitly set to None) we want to raise unless
                 # this is a fresh install since we know nothing needs to be migrated and
                 # the client will prevent making any changes until the default catalog is set.
-                if not default_catalog:
+                if default_catalog == "MISSING_REQUIRED_CATALOG":
                     versions = state_sync.get_versions(validate=False)
                     if versions.schema_version != 0:
                         raise SQLMeshError(


Test patch:
diff --git a/tests/dbt/test_util.py b/tests/dbt/test_util.py
new file mode 100644
index 00000000..67c8ba8e
--- /dev/null
+++ 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"]