Repository: TobikoData/sqlmesh
Pull request #2778 (2024-06-16 16:21:15)

Chore!: bump sqlglot to v25.2.0
Thought I'd post this early, need to deploy sqlglot first. Probably in a few days..

Context: https://github.com/tobymao/sqlglot/commit/468123e4b7612287e128529de62f3a88f4e1d579


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/setup.py b/setup.py
index d9c3d579..051f8f1e 100644
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,7 @@ setup(
         "requests",
         "rich[jupyter]",
         "ruamel.yaml",
-        "sqlglot[rs]~=25.1.0",
+        "sqlglot[rs]~=25.2.0",
     ],
     extras_require={
         "bigquery": [
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index 793bfc90..c838f9f2 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -256,7 +256,7 @@ class VarcharSizeWorkaroundMixin(EngineAdapter):
             # redshift and mssql have a bug where CTAS statements have non determistic types. if a limit
             # is applied to a ctas statement, VARCHAR types default to 1 in some instances.
             select_statement = statement.expression.copy()
-            for select_or_union in select_statement.find_all(exp.Select, exp.Union):
+            for select_or_union in select_statement.find_all(exp.Select, exp.SetOperation):
                 select_or_union.set("limit", None)
                 select_or_union.set("where", None)
 
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 9b2867a8..f5c3e444 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -422,20 +422,20 @@ class _Model(ModelMeta, frozen=True):
         """
         query = self.render_query_or_raise(**render_kwarg).copy()
 
-        for select_or_union in query.find_all(exp.Select, exp.Union):
-            cte = select_or_union.find_ancestor(exp.With, exp.Select, exp.Subquery)
+        for select_or_set_op in query.find_all(exp.Select, exp.SetOperation):
+            cte = select_or_set_op.find_ancestor(exp.With, exp.Select, exp.Subquery)
             skip_limit = isinstance(cte, exp.With) and cte.recursive
 
-            if isinstance(select_or_union, exp.Select) and select_or_union.args.get("from"):
-                select_or_union.where(exp.false(), copy=False)
-                if not skip_limit and not isinstance(select_or_union.parent, exp.Union):
-                    select_or_union.limit(0, copy=False)
+            if isinstance(select_or_set_op, exp.Select) and select_or_set_op.args.get("from"):
+                select_or_set_op.where(exp.false(), copy=False)
+                if not skip_limit and not isinstance(select_or_set_op.parent, exp.SetOperation):
+                    select_or_set_op.limit(0, copy=False)
             elif (
                 not skip_limit
-                and isinstance(select_or_union, exp.Union)
-                and not isinstance(select_or_union.parent, exp.Union)
+                and isinstance(select_or_set_op, exp.SetOperation)
+                and not isinstance(select_or_set_op.parent, exp.SetOperation)
             ):
-                select_or_union.set("limit", exp.Limit(expression=exp.Literal.number(0)))
+                select_or_set_op.set("limit", exp.Limit(expression=exp.Literal.number(0)))
 
         if self.managed_columns:
             query.select(
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"