Repository: TobikoData/sqlmesh
Pull request #2850 (2024-07-02 01:38:07)

feat: allow users to configure high watermark sensor behavior
This PR makes it so that users have the ability to change the behavior for HighWaterMark and HighWaterMarkExternal sensors.

Also removed the default values for `poke_interval` and `timeout` since they were just copies of Airflow's default we will just defer to Airflow to set them. 

Ran the Airflow local tests and confirmed they passed. 


Patch:
diff --git a/docs/guides/scheduling.md b/docs/guides/scheduling.md
index 8f19b5fd..ebc707e8 100644
--- a/docs/guides/scheduling.md
+++ b/docs/guides/scheduling.md
@@ -56,7 +56,7 @@ To integrate with [Airflow](../integrations/airflow.md), ensure that you meet th
         for dag in sqlmesh_airflow.dags:
             globals()[dag.dag_id] = dag
 
-    The example above uses `spark` as the engine of choice. Other engines can be configured instead by providing a corresponding string as an argument to the `SQLMeshAirflow` constructor. Supported strings are `"spark"`, `"databricks"`, `"snowflake"`, `"bigquery"`, `"redshift"`, `"trino"` and `"mssql"`. See the [Airflow Cluster Configuration](../integrations/airflow.md#airflow-cluster-configuration) for full list of arguments and their descriptions.
+    The example above uses `spark` as the engine of choice. Other engines can be configured instead by providing a corresponding string as an argument to the `SQLMeshAirflow` constructor. Supported strings are `"spark"`, `"databricks"`, `"snowflake"`, `"bigquery"`, `"redshift"`, `"trino"`, `"mssql"` and `"mysql"`. See the [Airflow Cluster Configuration](../integrations/airflow.md#airflow-cluster-configuration) for full list of arguments and their descriptions.
 
 After setup is completed, the `sqlmesh_janitor_dag` DAG should become available in the Airflow UI when filtered by the `sqlmesh` tag:
 
diff --git a/docs/integrations/airflow.md b/docs/integrations/airflow.md
index 8fdd4dc4..5275728a 100644
--- a/docs/integrations/airflow.md
+++ b/docs/integrations/airflow.md
@@ -27,18 +27,20 @@ The name of the module file can be arbitrary, but we recommend something descrip
 
 `SQLMeshAirflow` has two required arguments (`engine_operator` and `default_catalog`). Details on these and additional optional arguments below:
 
-| Argument                        | Description                                                                                                                                                                                                     | Type                   | Required |
+| Argument                        | Description                                                                                                                                                                                                     |          Type          | Required |
 |---------------------------------|-----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:----------------------:|:--------:|
-| `engine_operator`               | Name or operator to use for creating models. See [Engine Support](#engine-support) for list of options                                                                                                          | string or BaseOperator | Y        |
-| `default_catalog`               | The default catalog (also called "database" in other engines) to use when models are defined that do not contain a catalog in their name. This should match the default catalog applied by the connection.      | string                 | Y        |
-| `engine_operator_args`          | The dictionary of arguments that will be passed into the evaluate engine operator during its construction. This can be used to customize parameters such as connection ID.                                      | dict                   | N        |
-| `ddl_engine_operator`           | The type of the Airflow operator that will be used for environment management. These operations are SQL only. `engine_operator` is used if not provided                                                         | string or BaseOperator | N        |
-| `ddl_engine_operator_args`      | Args to be passed into just the environment management operator. This can be used to customize parameters such as connection ID.                                                                                | dict                   | N        |
-| `janitor_interval`              | Defines how often the janitor DAG runs. The janitor DAG removes platform-managed DAG instances that are pending deletion from Airflow. Default: 1 hour.                                                         | timedelta              | N        |
-| `plan_application_dag_ttl`      | Determines the time-to-live period for finished plan application DAGs. Once this period is exceeded, finished plan application DAGs are deleted by the janitor. Default: 2 days.                                | timedelta              | N        |
-| `external_table_sensor_factory` | A factory function that creates a sensor operator for a given signal payload. See [External signals](#external-signals) for more info                                                                           | function               | N        |
-| `sensor_mode`                   | The mode to use for SQLMesh sensors. Supported values are "poke" and "reschedule". See https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/sensors.html for more details. Default: "reschedule" | string                 | N        |
-| `generate_cadence_dags`         | Whether to generate cadence DAGs for model versions that are currently deployed to production.                                                                                                                  | bool                   | N        |
+| `engine_operator`               | Name or operator to use for creating models. See [Engine Support](#engine-support) for list of options                                                                                                          | string or BaseOperator |    Y     |
+| `default_catalog`               | The default catalog (also called "database" in other engines) to use when models are defined that do not contain a catalog in their name. This should match the default catalog applied by the connection.      |         string         |    Y     |
+| `engine_operator_args`          | The dictionary of arguments that will be passed into the evaluate engine operator during its construction. This can be used to customize parameters such as connection ID.                                      |          dict          |    N     |
+| `ddl_engine_operator`           | The type of the Airflow operator that will be used for environment management. These operations are SQL only. `engine_operator` is used if not provided                                                         | string or BaseOperator |    N     |
+| `ddl_engine_operator_args`      | Args to be passed into just the environment management operator. This can be used to customize parameters such as connection ID.                                                                                |          dict          |    N     |
+| `janitor_interval`              | Defines how often the janitor DAG runs. The janitor DAG removes platform-managed DAG instances that are pending deletion from Airflow. Default: 1 hour.                                                         |       timedelta        |    N     |
+| `plan_application_dag_ttl`      | Determines the time-to-live period for finished plan application DAGs. Once this period is exceeded, finished plan application DAGs are deleted by the janitor. Default: 2 days.                                |       timedelta        |    N     |
+| `external_table_sensor_factory` | A factory function that creates a sensor operator for a given signal payload. See [External signals](#external-signals) for more info                                                                           |        function        |    N     |
+| `sensor_mode`                   | The mode to use for SQLMesh sensors. Supported values are "poke" and "reschedule". See https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/sensors.html for more details. Default: "reschedule" |         string         |    N     |
+| `high_water_mark_sensor_args`   | The dictionary of arguments that will be passed into the high water mark sensor during its construction.                                                                                                        |          dict          |    N     |
+| `external_sensor_args`          | The dictionary of arguments that will be passed into the external sensor during its construction.                                                                                                               |          dict          |    N     |
+| `generate_cadence_dags`         | Whether to generate cadence DAGs for model versions that are currently deployed to production.                                                                                                                  |          bool          |    N     |
 
 
 ### State connection
@@ -123,6 +125,7 @@ SQLMesh supports a variety of engines in Airflow. Support for each engine is pro
 * [Snowflake](engines/snowflake.md#airflow-scheduler)
 * [Spark](engines/spark.md#airflow-scheduler)
 * [Trino](engines/trino.md#airflow-scheduler)
+* [MySQL](engines/mysql.md#airflow-scheduler)
 
 ## Managed Airflow instances
 
diff --git a/docs/integrations/engines/mysql.md b/docs/integrations/engines/mysql.md
index e8426a3f..acd0d660 100644
--- a/docs/integrations/engines/mysql.md
+++ b/docs/integrations/engines/mysql.md
@@ -19,3 +19,27 @@ pip install "sqlmesh[mysql]"
 | `port`         | The port number of the MySQL server                          | int    | N        |
 | `charset`      | The character set used for the connection                    | string | N        |
 | `ssl_disabled` | Is SSL disabled                                              | bool   | N        |
+
+## Airflow Scheduler
+**Engine Name:** `mysql`
+
+The SQLMesh MySQL Operator is similar to the [MySQLOperator](https://airflow.apache.org/docs/apache-airflow-providers-mysql/stable/index.html), and relies on the same [MySqlHook](https://airflow.apache.org/docs/apache-airflow-providers-mysql/1.0.0/_api/airflow/providers/mysql/hooks/mysql/index.html) implementation.
+
+To enable support for this operator, the Airflow MySQL provider package should be installed on the target Airflow cluster along with SQLMesh with the mysql extra:
+```
+pip install "apache-airflow-providers-mysql"
+pip install "sqlmesh[mysql]"
+```
+
+The operator requires an [Airflow connection](https://airflow.apache.org/docs/apache-airflow/stable/howto/connection.html) to determine the target MySQL account. Refer to [MySQL connection](https://airflow.apache.org/docs/apache-airflow-providers-mysql/stable/connections/mysql.html) for more details.
+
+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"
+sqlmesh_airflow = SQLMeshAirflow(
+    "mysql",
+    default_catalog="<database name>",
+    engine_operator_args={
+        "mysql_conn_id": "<Connection ID>"
+    },
+)
+```
\ No newline at end of file
diff --git a/sqlmesh/core/engine_adapter/bigquery.py b/sqlmesh/core/engine_adapter/bigquery.py
index 7c384d9b..9ef82274 100644
--- a/sqlmesh/core/engine_adapter/bigquery.py
+++ b/sqlmesh/core/engine_adapter/bigquery.py
@@ -65,9 +65,6 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
                 exp.DataType.build("FLOAT64", dialect=DIALECT),
                 exp.DataType.build("BIGNUMERIC", dialect=DIALECT),
             },
-            exp.DataType.build("BIGNUMERIC", dialect=DIALECT): {
-                exp.DataType.build("FLOAT64", dialect=DIALECT),
-            },
             exp.DataType.build("DATE", dialect=DIALECT): {
                 exp.DataType.build("DATETIME", dialect=DIALECT),
             },
diff --git a/sqlmesh/core/snapshot/evaluator.py b/sqlmesh/core/snapshot/evaluator.py
index 17676139..c44349a5 100644
--- a/sqlmesh/core/snapshot/evaluator.py
+++ b/sqlmesh/core/snapshot/evaluator.py
@@ -792,19 +792,19 @@ class SnapshotEvaluator:
     ) -> None:
         snapshot = snapshot.table_info
 
-        table_names = [snapshot.table_name(is_deployable=False)]
+        table_names = [(False, snapshot.table_name(is_deployable=False))]
         if not dev_table_only:
-            table_names.append(snapshot.table_name(is_deployable=True))
+            table_names.append((True, snapshot.table_name(is_deployable=True)))
 
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
-        for table_name in table_names:
+        for is_table_deployable, table_name in table_names:
             table = exp.to_table(table_name)
             if table.db != snapshot.physical_schema:
                 raise SQLMeshError(
                     f"Table '{table_name}' is not a part of the physical schema '{snapshot.physical_schema}' and so can't be dropped."
                 )
-            evaluation_strategy.delete(table_name)
+            evaluation_strategy.delete(table_name, is_table_deployable=is_table_deployable)
 
             if on_complete is not None:
                 on_complete(table_name)
@@ -1776,8 +1776,13 @@ class EngineManagedStrategy(MaterializableStrategy):
         raise ConfigError(f"Cannot mutate managed table: {target_table_name}")
 
     def delete(self, name: str, **kwargs: t.Any) -> None:
-        self.adapter.drop_managed_table(name)
-        logger.info("Dropped managed table '%s'", name)
+        # a dev preview table is created as a normal table, so it needs to be dropped as a normal table
+        if kwargs["is_table_deployable"]:
+            self.adapter.drop_managed_table(name)
+            logger.info("Dropped managed table '%s'", name)
+        else:
+            self.adapter.drop_table(name)
+            logger.info("Dropped dev preview for managed table '%s'", name)
 
 
 def _intervals(snapshot: Snapshot, deployability_index: DeployabilityIndex) -> Intervals:
diff --git a/sqlmesh/schedulers/airflow/dag_generator.py b/sqlmesh/schedulers/airflow/dag_generator.py
index ba33f36c..ce770a5f 100644
--- a/sqlmesh/schedulers/airflow/dag_generator.py
+++ b/sqlmesh/schedulers/airflow/dag_generator.py
@@ -74,6 +74,8 @@ class SnapshotDagGenerator:
             t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
         ],
         sensor_mode: str,
+        high_water_mark_sensor_args: t.Optional[t.Dict[str, t.Any]],
+        external_sensor_args: t.Optional[t.Dict[str, t.Any]],
         state_reader: StateReader,
     ):
         self._engine_operator = engine_operator
@@ -83,6 +85,8 @@ class SnapshotDagGenerator:
         self._external_table_sensor_factory = external_table_sensor_factory
         self._state_reader = state_reader
         self._sensor_mode = sensor_mode
+        self._high_water_mark_sensor_args = high_water_mark_sensor_args or {}
+        self._external_sensor_args = external_sensor_args or {}
 
     def generate_cadence_dags(self, snapshots: t.Iterable[SnapshotIdLike]) -> t.List[DAG]:
         dags = []
@@ -647,12 +651,13 @@ class SnapshotDagGenerator:
                         this_snapshot=snapshot,
                         task_id=f"{sanitize_name(upstream_snapshot.node.name)}_{upstream_snapshot.version}_high_water_mark_sensor",
                         mode=self._sensor_mode,
+                        **self._high_water_mark_sensor_args,
                     )
                 )
 
-        external_sesnor = self._create_hwm_external_sensor(snapshot)
-        if external_sesnor:
-            output.append(external_sesnor)
+        external_sensor = self._create_hwm_external_sensor(snapshot)
+        if external_sensor:
+            output.append(external_sensor)
 
         return output
 
@@ -670,6 +675,7 @@ class SnapshotDagGenerator:
                 mode=self._sensor_mode,
                 start=start,
                 end=end,
+                **self._external_sensor_args,
             )
         return None
 
diff --git a/sqlmesh/schedulers/airflow/integration.py b/sqlmesh/schedulers/airflow/integration.py
index 8ab618d1..06f5db2b 100644
--- a/sqlmesh/schedulers/airflow/integration.py
+++ b/sqlmesh/schedulers/airflow/integration.py
@@ -67,6 +67,8 @@ class SQLMeshAirflow:
         external_table_sensor_factory: A factory function that creates a sensor operator for a given signal payload.
         sensor_mode: The mode to use for SQLMesh sensors. Supported values are "poke" and "reschedule".
             See https://airflow.apache.org/docs/apache-airflow/stable/core-concepts/sensors.html for more details. Default: "reschedule".
+        high_water_mark_sensor_args: The dictionary of arguments that will be passed into the high water mark sensor operator during its construction.
+        external_sensor_args: The dictionary of arguments that will be passed into the external sensor operator during its construction.
         generate_cadence_dags: Whether to generate cadence DAGs for model versions that are currently deployed to production.
     """
 
@@ -83,6 +85,8 @@ class SQLMeshAirflow:
             t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
         ] = None,
         sensor_mode: str = "reschedule",
+        high_water_mark_sensor_args: t.Optional[t.Dict[str, t.Any]] = None,
+        external_sensor_args: t.Optional[t.Dict[str, t.Any]] = None,
         generate_cadence_dags: bool = True,
     ):
         if isinstance(engine_operator, str):
@@ -109,6 +113,8 @@ class SQLMeshAirflow:
         self._generate_cadence_dags = generate_cadence_dags
         self._default_catalog = default_catalog
         self._sensor_mode = sensor_mode
+        self._high_water_mark_sensor_args = high_water_mark_sensor_args or {}
+        self._external_sensor_args = external_sensor_args or {}
 
     @classmethod
     def set_default_catalog(cls, default_catalog: str) -> None:
@@ -193,6 +199,8 @@ class SQLMeshAirflow:
             self._ddl_engine_operator_args,
             self._external_table_sensor_factory,
             self._sensor_mode,
+            self._high_water_mark_sensor_args,
+            self._external_sensor_args,
             state_reader,
         )
 
diff --git a/sqlmesh/schedulers/airflow/operators/hwm_sensor.py b/sqlmesh/schedulers/airflow/operators/hwm_sensor.py
index 8897b4e0..e320f396 100644
--- a/sqlmesh/schedulers/airflow/operators/hwm_sensor.py
+++ b/sqlmesh/schedulers/airflow/operators/hwm_sensor.py
@@ -23,14 +23,10 @@ class HighWaterMarkSensor(BaseSensorOperator):
         self,
         target_snapshot_info: SnapshotTableInfo,
         this_snapshot: Snapshot,
-        poke_interval: float = 60.0,
-        timeout: float = 7.0 * 24.0 * 60.0 * 60.0,  # 7 days
         mode: str = "reschedule",
         **kwargs: t.Any,
     ) -> None:
         super().__init__(
-            poke_interval=poke_interval,
-            timeout=timeout,
             mode=mode,
             **kwargs,
         )
@@ -78,16 +74,12 @@ class HighWaterMarkExternalSensor(BaseSensorOperator):
         self,
         snapshot: Snapshot,
         external_table_sensor_factory: t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator],
-        poke_interval: float = 60.0,
-        timeout: float = 7.0 * 24.0 * 60.0 * 60.0,  # 7 days
         mode: str = "reschedule",
         start: t.Optional[TimeLike] = None,
         end: t.Optional[TimeLike] = None,
         **kwargs: t.Any,
     ):
         super().__init__(
-            poke_interval=poke_interval,
-            timeout=timeout,
             mode=mode,
             **kwargs,
         )
diff --git a/sqlmesh/schedulers/airflow/operators/mysql.py b/sqlmesh/schedulers/airflow/operators/mysql.py
new file mode 100644
index 00000000..424853d2
--- /dev/null
+++ b/sqlmesh/schedulers/airflow/operators/mysql.py
@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import typing as t
+
+from airflow.models import BaseOperator
+from airflow.providers.mysql.hooks.mysql import MySqlHook
+from airflow.utils.context import Context
+
+from sqlmesh.schedulers.airflow.operators.targets import BaseTarget
+
+
+class SQLMeshMySQLOperator(BaseOperator):
+    """The operator that evaluates a SQLMesh model snapshot on a mysql target
+
+    Args:
+        target: The target that will be executed by this operator instance.
+        mysql_conn_id: The Airflow connection id for the mysql target.
+    """
+
+    def __init__(
+        self,
+        *,
+        target: BaseTarget,
+        mysql_conn_id: str = MySqlHook.default_conn_name,
+        **kwargs: t.Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._target = target
+        self._mysql_conn_id = mysql_conn_id
+        self._hook_params = kwargs
+
+    def get_db_hook(self) -> MySqlHook:
+        """Gets the mysql Hook which contains the DB API connection object"""
+        return MySqlHook(self._mysql_conn_id, **self._hook_params)
+
+    def execute(self, context: Context) -> None:
+        """Executes the desired target against the configured mysql connection"""
+        self._target.execute(context, lambda: self.get_db_hook().get_conn(), "mysql")
diff --git a/sqlmesh/schedulers/airflow/util.py b/sqlmesh/schedulers/airflow/util.py
index 9009bd45..cbaf2062 100644
--- a/sqlmesh/schedulers/airflow/util.py
+++ b/sqlmesh/schedulers/airflow/util.py
@@ -173,6 +173,11 @@ def discover_engine_operator(name: str, sql_only: bool = False) -> t.Type[BaseOp
             from sqlmesh.schedulers.airflow.operators.mssql import SQLMeshMsSqlOperator
 
             return SQLMeshMsSqlOperator
+
+        if name == "mysql":
+            from sqlmesh.schedulers.airflow.operators.mysql import SQLMeshMySQLOperator
+
+            return SQLMeshMySQLOperator
     except ImportError:
         raise SQLMeshError(f"Failed to automatically discover an operator for '{name}'.'")
 


Test patch:
diff --git a/tests/core/test_snapshot_evaluator.py b/tests/core/test_snapshot_evaluator.py
index ea95c45e..3dd4b7f9 100644
--- a/tests/core/test_snapshot_evaluator.py
+++ b/tests/core/test_snapshot_evaluator.py
@@ -2365,3 +2365,41 @@ def test_evaluate_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
         table_description=model.description,
         column_descriptions=model.column_descriptions,
     )
+
+
+def test_cleanup_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
+    evaluator = SnapshotEvaluator(adapter_mock)
+
+    model = load_sql_based_model(
+        parse(  # type: ignore
+            """
+            MODEL (
+                name test_schema.test_model,
+                kind MANAGED,
+                physical_properties (
+                    warehouse = 'small',
+                    target_lag = '10 minutes'
+                ),
+                clustered_by a
+            );
+
+            select a, b from foo;
+            """
+        )
+    )
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    adapter_mock.assert_not_called()
+
+    cleanup_task = SnapshotTableCleanupTask(snapshot=snapshot.table_info, dev_table_only=False)
+
+    evaluator.cleanup(target_snapshots=[cleanup_task])
+
+    adapter_mock.drop_table.assert_called_once_with(
+        "sqlmesh__test_schema.test_schema__test_model__1556851963__temp"
+    )
+    adapter_mock.drop_managed_table.assert_called_once_with(
+        "sqlmesh__test_schema.test_schema__test_model__1556851963"
+    )
diff --git a/tests/schedulers/airflow/test_dag_generator.py b/tests/schedulers/airflow/test_dag_generator.py
index 5e3abdc3..a6df3273 100644
--- a/tests/schedulers/airflow/test_dag_generator.py
+++ b/tests/schedulers/airflow/test_dag_generator.py
@@ -70,6 +70,8 @@ def test_generate_plan_application_dag__batch_index_populated(mocker: MockerFixt
         external_table_sensor_factory=None,
         sensor_mode="reschedule",
         state_reader=state_reader_mock,
+        external_sensor_args=None,
+        high_water_mark_sensor_args=None,
     )
 
     environment_name = "test_env"
@@ -157,6 +159,8 @@ def test_sensor_mode_override(mocker: MockerFixture, make_snapshot):
         external_table_sensor_factory=None,
         sensor_mode="poke",
         state_reader=state_reader_mock,
+        external_sensor_args=None,
+        high_water_mark_sensor_args=None,
     )
 
     dags = generator.generate_cadence_dags([snapshot_a, snapshot_b])