Repository: TobikoData/sqlmesh
Pull request #2786 (2024-06-17 23:23:30)

Feat: Make sensor mode configurable for SQLMesh sensors in Airflow
The default mode (`reschedule`) is not always the best fit, especially if the Airflow cluster is configured to use Kubernetes workers. With such a setup, a new pod gets allocated each time on each poll, which adds significantly to the task latency.

This update allows users to set the desired mode themselves.


Patch:
diff --git a/docs/integrations/airflow.md b/docs/integrations/airflow.md
index 22422729..8fdd4dc4 100644
--- a/docs/integrations/airflow.md
+++ b/docs/integrations/airflow.md
@@ -27,17 +27,18 @@ 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 |
-|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|:----------------------:|:--------:|
-| `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     |
-| `generate_cadence_dags`         | Whether to generate cadence DAGs for model versions that are currently deployed to production.                                                                                                             |          bool          |    N     |
+| 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        |
 
 
 ### State connection
diff --git a/sqlmesh/schedulers/airflow/dag_generator.py b/sqlmesh/schedulers/airflow/dag_generator.py
index 9f00c017..ba33f36c 100644
--- a/sqlmesh/schedulers/airflow/dag_generator.py
+++ b/sqlmesh/schedulers/airflow/dag_generator.py
@@ -73,6 +73,7 @@ class SnapshotDagGenerator:
         external_table_sensor_factory: t.Optional[
             t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
         ],
+        sensor_mode: str,
         state_reader: StateReader,
     ):
         self._engine_operator = engine_operator
@@ -81,6 +82,7 @@ class SnapshotDagGenerator:
         self._ddl_engine_operator_args = ddl_engine_operator_args or {}
         self._external_table_sensor_factory = external_table_sensor_factory
         self._state_reader = state_reader
+        self._sensor_mode = sensor_mode
 
     def generate_cadence_dags(self, snapshots: t.Iterable[SnapshotIdLike]) -> t.List[DAG]:
         dags = []
@@ -644,6 +646,7 @@ class SnapshotDagGenerator:
                         target_snapshot_info=upstream_snapshot.table_info,
                         this_snapshot=snapshot,
                         task_id=f"{sanitize_name(upstream_snapshot.node.name)}_{upstream_snapshot.version}_high_water_mark_sensor",
+                        mode=self._sensor_mode,
                     )
                 )
 
@@ -664,6 +667,7 @@ class SnapshotDagGenerator:
                 snapshot=snapshot,
                 external_table_sensor_factory=self._external_table_sensor_factory,
                 task_id="external_high_water_mark_sensor",
+                mode=self._sensor_mode,
                 start=start,
                 end=end,
             )
diff --git a/sqlmesh/schedulers/airflow/integration.py b/sqlmesh/schedulers/airflow/integration.py
index 2abc4682..8ab618d1 100644
--- a/sqlmesh/schedulers/airflow/integration.py
+++ b/sqlmesh/schedulers/airflow/integration.py
@@ -65,6 +65,8 @@ class SQLMeshAirflow:
         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.
         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".
         generate_cadence_dags: Whether to generate cadence DAGs for model versions that are currently deployed to production.
     """
 
@@ -80,6 +82,7 @@ class SQLMeshAirflow:
         external_table_sensor_factory: t.Optional[
             t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
         ] = None,
+        sensor_mode: str = "reschedule",
         generate_cadence_dags: bool = True,
     ):
         if isinstance(engine_operator, str):
@@ -105,6 +108,7 @@ class SQLMeshAirflow:
         self._external_table_sensor_factory = external_table_sensor_factory
         self._generate_cadence_dags = generate_cadence_dags
         self._default_catalog = default_catalog
+        self._sensor_mode = sensor_mode
 
     @classmethod
     def set_default_catalog(cls, default_catalog: str) -> None:
@@ -188,6 +192,7 @@ class SQLMeshAirflow:
             self._ddl_engine_operator,
             self._ddl_engine_operator_args,
             self._external_table_sensor_factory,
+            self._sensor_mode,
             state_reader,
         )
 
diff --git a/sqlmesh/schedulers/airflow/operators/targets.py b/sqlmesh/schedulers/airflow/operators/targets.py
index 76f31ed6..cecd9809 100644
--- a/sqlmesh/schedulers/airflow/operators/targets.py
+++ b/sqlmesh/schedulers/airflow/operators/targets.py
@@ -105,7 +105,7 @@ class BaseTarget(abc.ABC, t.Generic[CP]):
         return payload
 
 
-class SnapshotEvaluationTarget(BaseTarget[commands.EvaluateCommandPayload], PydanticModel):
+class SnapshotEvaluationTarget(PydanticModel, BaseTarget[commands.EvaluateCommandPayload]):
     """The target which contains attributes necessary to evaluate a given snapshot.
 
     Args:
@@ -176,7 +176,7 @@ class SnapshotEvaluationTarget(BaseTarget[commands.EvaluateCommandPayload], Pyda
         return self.execution_time or context["dag_run"].logical_date
 
 
-class SnapshotPromotionTarget(BaseTarget[commands.PromoteCommandPayload], PydanticModel):
+class SnapshotPromotionTarget(PydanticModel, BaseTarget[commands.PromoteCommandPayload]):
     """The target which contains attributes necessary to perform snapshot promotion in a given environment.
 
     The promotion means creation of views associated with the environment which target physical tables
@@ -208,7 +208,7 @@ class SnapshotPromotionTarget(BaseTarget[commands.PromoteCommandPayload], Pydant
         )
 
 
-class SnapshotDemotionTarget(BaseTarget[commands.DemoteCommandPayload], PydanticModel):
+class SnapshotDemotionTarget(PydanticModel, BaseTarget[commands.DemoteCommandPayload]):
     """The target which contains attributes necessary to perform snapshot demotion in a given environment.
 
     The demotion means deletion of views that match names of provided snapshots in the target environment.
@@ -234,7 +234,7 @@ class SnapshotDemotionTarget(BaseTarget[commands.DemoteCommandPayload], Pydantic
         )
 
 
-class SnapshotCleanupTarget(BaseTarget[commands.CleanupCommandPayload], PydanticModel):
+class SnapshotCleanupTarget(PydanticModel, BaseTarget[commands.CleanupCommandPayload]):
     """The target which contains attributes necessary to perform table cleanup of expired snapshots"""
 
     command_type: commands.CommandType = commands.CommandType.CLEANUP
@@ -266,7 +266,7 @@ class SnapshotCleanupTarget(BaseTarget[commands.CleanupCommandPayload], Pydantic
         return command
 
 
-class SnapshotCreateTablesTarget(BaseTarget[commands.CreateTablesCommandPayload], PydanticModel):
+class SnapshotCreateTablesTarget(PydanticModel, BaseTarget[commands.CreateTablesCommandPayload]):
     """The target which creates physical tables for the given set of new snapshots."""
 
     command_type: commands.CommandType = commands.CommandType.CREATE_TABLES
@@ -293,7 +293,7 @@ class SnapshotCreateTablesTarget(BaseTarget[commands.CreateTablesCommandPayload]
         )
 
 
-class SnapshotMigrateTablesTarget(BaseTarget[commands.MigrateTablesCommandPayload], PydanticModel):
+class SnapshotMigrateTablesTarget(PydanticModel, BaseTarget[commands.MigrateTablesCommandPayload]):
     """The target which updates schemas of existing physical tables to bring them in correspondance
     with schemas of target snapshots.
     """


Test patch:
diff --git a/tests/schedulers/airflow/test_dag_generator.py b/tests/schedulers/airflow/test_dag_generator.py
index 15f4a896..5e3abdc3 100644
--- a/tests/schedulers/airflow/test_dag_generator.py
+++ b/tests/schedulers/airflow/test_dag_generator.py
@@ -15,7 +15,8 @@ from sqlmesh.core.snapshot import (
 from sqlmesh.schedulers.airflow.dag_generator import SnapshotDagGenerator
 from sqlmesh.schedulers.airflow import common
 from sqlmesh.schedulers.airflow.operators.targets import BaseTarget, SnapshotEvaluationTarget
-from sqlmesh.utils.date import to_datetime
+from sqlmesh.schedulers.airflow.operators.hwm_sensor import HighWaterMarkSensor
+from sqlmesh.utils.date import to_datetime, to_timestamp
 
 
 class TestSubmitOperator(BaseOperator):
@@ -67,6 +68,7 @@ def test_generate_plan_application_dag__batch_index_populated(mocker: MockerFixt
         ddl_engine_operator=TestSubmitOperator,
         ddl_engine_operator_args={},
         external_table_sensor_factory=None,
+        sensor_mode="reschedule",
         state_reader=state_reader_mock,
     )
 
@@ -125,3 +127,45 @@ def test_generate_plan_application_dag__batch_index_populated(mocker: MockerFixt
         assert command is not None
         assert target.batch_index == batch_idx
         assert command.batch_index == batch_idx
+
+
+def test_sensor_mode_override(mocker: MockerFixture, make_snapshot):
+    snapshot_a = make_snapshot(
+        SqlModel(name="a", kind=dict(name="FULL"), query=parse_one("select 1 as a, ds")),
+    )
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_a.unpaused_ts = to_timestamp("2024-01-01")
+
+    snapshot_b = make_snapshot(
+        SqlModel(name="b", kind=dict(name="FULL"), query=parse_one("select a, ds from a")),
+        nodes={snapshot_a.name: snapshot_a.node},
+    )
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+    snapshot_b.unpaused_ts = to_timestamp("2024-01-01")
+
+    state_reader_mock = mocker.Mock()
+    state_reader_mock.get_snapshots.return_value = {
+        snapshot_a.snapshot_id: snapshot_a,
+        snapshot_b.snapshot_id: snapshot_b,
+    }
+
+    generator = SnapshotDagGenerator(
+        engine_operator=TestSubmitOperator,
+        engine_operator_args={},
+        ddl_engine_operator=TestSubmitOperator,
+        ddl_engine_operator_args={},
+        external_table_sensor_factory=None,
+        sensor_mode="poke",
+        state_reader=state_reader_mock,
+    )
+
+    dags = generator.generate_cadence_dags([snapshot_a, snapshot_b])
+    assert len(dags) == 2
+
+    assert len(dags[0].tasks) == 1
+    assert isinstance(dags[0].tasks[0], TestSubmitOperator)
+
+    assert len(dags[1].tasks) == 2
+    assert isinstance(dags[1].tasks[0], HighWaterMarkSensor)
+    assert isinstance(dags[1].tasks[1], TestSubmitOperator)
+    assert dags[1].tasks[0].mode == "poke"