Feat: Support sensors for external tables in Airflow (#1536)

Author: izeigerman

sqlmesh/core/model/definition.py

@@ -363,6 +363,53 @@ def render_post_statements(
         """
         return []
 
+    def render_signals(
+        self,
+        *,
+        start: t.Optional[TimeLike] = None,
+        end: t.Optional[TimeLike] = None,
+        execution_time: t.Optional[TimeLike] = None,
+    ) -> t.List[t.Dict[str, str | int | float | bool]]:
+        """Renders external; signals defined for this model.
+
+        Args:
+            start: The start datetime to render. Defaults to epoch start.
+            end: The end datetime to render. Defaults to epoch start.
+            execution_time: The date/time time reference to use for execution time.
+
+        Returns:
+            The list of rendered expressions.
+        """
+
+        def _create_renderer(expression: exp.Expression) -> ExpressionRenderer:
+            return ExpressionRenderer(
+                expression,
+                self.dialect,
+                [],
+                path=self._path,
+                jinja_macro_registry=self.jinja_macros,
+                python_env=self.python_env,
+                only_execution_time=False,
+            )
+
+        def _render(e: exp.Expression) -> str | int | float | bool:
+            rendered_exprs = _create_renderer(e).render(
+                start=start, end=end, execution_time=execution_time
+            )
+            if len(rendered_exprs) != 1:
+                raise SQLMeshError(f"Expected one expression but got {len(rendered_exprs)}")
+
+            rendered = rendered_exprs[0]
+            if rendered.is_int:
+                return int(rendered.this)
+            if rendered.is_number:
+                return float(rendered.this)
+            if isinstance(rendered, (exp.Literal, exp.Boolean)):
+                return rendered.this
+            return rendered.sql(dialect=self.dialect)
+
+        return [{t.this.name: _render(t.expression) for t in signal} for signal in self.signals]
+
     def ctas_query(self, **render_kwarg: t.Any) -> exp.Subqueryable:
         """Return a dummy query to do a CTAS.
 
@@ -780,6 +827,8 @@ def metadata_hash(self, audits: t.Dict[str, ModelAudit]) -> str:
             else:
                 raise SQLMeshError(f"Unexpected audit name '{audit_name}'.")
 
+        metadata.extend([s.sql(comments=True) for s in self.signals])
+
         # Add comments from the query.
         if self.is_sql:
             rendered_query = self.render_query()
@@ -1957,4 +2006,5 @@ def _refs_to_sql(values: t.Any) -> exp.Expression:
     "table_properties_": lambda value: value,
     "session_properties_": lambda value: value,
     "allow_partials": exp.convert,
+    "signals": lambda values: exp.Tuple(expressions=values),
 }

sqlmesh/core/model/meta.py

@@ -11,6 +11,7 @@
 from sqlmesh.core.model.common import (
     bool_validator,
     parse_expressions,
+    parse_properties,
     properties_validator,
 )
 from sqlmesh.core.model.kind import (
@@ -57,6 +58,7 @@ class ModelMeta(_Node, extra="allow"):
     table_properties_: t.Optional[exp.Tuple] = Field(default=None, alias="table_properties")
     session_properties_: t.Optional[exp.Tuple] = Field(default=None, alias="session_properties")
     allow_partials: bool = False
+    signals: t.List[exp.Tuple] = []
 
     _table_properties: t.Dict[str, exp.Expression] = {}
 
@@ -230,6 +232,28 @@ def _refs_validator(cls, vs: t.Any, values: t.Dict[str, t.Any]) -> t.List[exp.Ex
 
         return refs
 
+    @field_validator("signals", mode="before")
+    @field_validator_v1_args
+    def _signals_validator(cls, v: t.Any, values: t.Dict[str, t.Any]) -> t.Any:
+        if v is None:
+            return []
+
+        if isinstance(v, str):
+            dialect = values.get("dialect")
+            v = d.parse_one(v, dialect=dialect)
+
+        if isinstance(v, (exp.Array, exp.Paren, exp.Tuple)):
+            tuples: t.List[exp.Expression] = (
+                [v.unnest()] if isinstance(v, exp.Paren) else v.expressions
+            )
+            signals = [parse_properties(cls, t, values) for t in tuples]
+        elif isinstance(v, list):
+            signals = [parse_properties(cls, t, values) for t in v]
+        else:
+            raise ConfigError(f"Unexpected signals '{v}'")
+
+        return signals
+
     @model_validator(mode="before")
     def _pre_root_validator(cls, values: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
         grain = values.pop("grain", None)

sqlmesh/schedulers/airflow/dag_generator.py

@@ -9,6 +9,7 @@
 from airflow.models import BaseOperator, baseoperator
 from airflow.operators.empty import EmptyOperator
 from airflow.operators.python import PythonOperator
+from airflow.sensors.base import BaseSensorOperator
 
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo
 from sqlmesh.core.notification_target import NotificationTarget
@@ -21,7 +22,10 @@
 )
 from sqlmesh.schedulers.airflow import common, util
 from sqlmesh.schedulers.airflow.operators import targets
-from sqlmesh.schedulers.airflow.operators.hwm_sensor import HighWaterMarkSensor
+from sqlmesh.schedulers.airflow.operators.hwm_sensor import (
+    HighWaterMarkExternalSensor,
+    HighWaterMarkSensor,
+)
 from sqlmesh.schedulers.airflow.operators.notification import (
     BaseNotificationOperatorProvider,
 )
@@ -58,12 +62,16 @@ def __init__(
         engine_operator_args: t.Optional[t.Dict[str, t.Any]],
         ddl_engine_operator: t.Type[BaseOperator],
         ddl_engine_operator_args: t.Optional[t.Dict[str, t.Any]],
+        external_table_sensor_factory: t.Optional[
+            t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
+        ],
         snapshots: t.Dict[SnapshotId, Snapshot],
     ):
         self._engine_operator = engine_operator
         self._engine_operator_args = engine_operator_args or {}
         self._ddl_engine_operator = ddl_engine_operator
         self._ddl_engine_operator_args = ddl_engine_operator_args or {}
+        self._external_table_sensor_factory = external_table_sensor_factory
         self._snapshots = snapshots
 
     def generate_cadence_dags(self, snapshots: t.Iterable[SnapshotIdLike]) -> t.List[DAG]:
@@ -506,10 +514,12 @@ def _create_snapshot_evaluator_operator(
             task_id=task_id,
         )
 
-    def _create_hwm_sensors(self, snapshot: Snapshot) -> t.List[HighWaterMarkSensor]:
-        output = []
+    def _create_hwm_sensors(self, snapshot: Snapshot) -> t.List[BaseSensorOperator]:
+        output: t.List[BaseSensorOperator] = []
+        depends_on = snapshot.node.depends_on
         for upstream_snapshot_id in snapshot.parents:
             upstream_snapshot = self._snapshots[upstream_snapshot_id]
+            depends_on.discard(upstream_snapshot.name)
             if not upstream_snapshot.is_symbolic and not upstream_snapshot.is_seed:
                 output.append(
                     HighWaterMarkSensor(
@@ -518,6 +528,16 @@ def _create_hwm_sensors(self, snapshot: Snapshot) -> t.List[HighWaterMarkSensor]
                         task_id=f"{sanitize_name(upstream_snapshot.name)}_{upstream_snapshot.version}_high_water_mark_sensor",
                     )
                 )
+
+        if self._external_table_sensor_factory and snapshot.model.signals:
+            output.append(
+                HighWaterMarkExternalSensor(
+                    snapshot=snapshot,
+                    external_table_sensor_factory=self._external_table_sensor_factory,
+                    task_id="external_high_water_mark_sensor",
+                )
+            )
+
         return output
 
 

sqlmesh/schedulers/airflow/integration.py

@@ -7,6 +7,7 @@
 from airflow import DAG
 from airflow.models import BaseOperator, TaskInstance, Variable
 from airflow.operators.python import PythonOperator
+from airflow.sensors.base import BaseSensorOperator
 from airflow.utils.session import provide_session
 from sqlalchemy.orm import Session
 
@@ -58,6 +59,7 @@ class SQLMeshAirflow:
             deletion from Airflow. Default: 1 hour.
         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.
     """
 
     def __init__(
@@ -68,6 +70,9 @@ def __init__(
         ddl_engine_operator_args: t.Optional[t.Dict[str, t.Any]] = None,
         janitor_interval: timedelta = timedelta(hours=1),
         plan_application_dag_ttl: timedelta = timedelta(days=2),
+        external_table_sensor_factory: t.Optional[
+            t.Callable[[t.Dict[str, t.Any]], BaseSensorOperator]
+        ] = None,
     ):
         if isinstance(engine_operator, str):
             if not ddl_engine_operator:
@@ -89,6 +94,7 @@ def __init__(
             self._ddl_engine_operator_args = ddl_engine_operator_args or {}
         self._janitor_interval = janitor_interval
         self._plan_application_dag_ttl = plan_application_dag_ttl
+        self._external_table_sensor_factory = external_table_sensor_factory
 
     @property
     def dags(self) -> t.List[DAG]:
@@ -109,6 +115,7 @@ def dags(self) -> t.List[DAG]:
             self._engine_operator_args,
             self._ddl_engine_operator,
             self._ddl_engine_operator_args,
+            self._external_table_sensor_factory,
             stored_snapshots,
         )
 

sqlmesh/schedulers/airflow/operators/hwm_sensor.py

@@ -3,12 +3,12 @@
 from datetime import datetime
 
 from airflow.models import DagRun
-from airflow.sensors.base import BaseSensorOperator
+from airflow.sensors.base import BaseSensorOperator, PokeReturnValue
 from airflow.utils.context import Context
 
 from sqlmesh.core.snapshot import Snapshot, SnapshotTableInfo
 from sqlmesh.schedulers.airflow import util
-from sqlmesh.utils.date import to_datetime
+from sqlmesh.utils.date import now, to_datetime
 
 logger = logging.getLogger(__name__)
 
@@ -63,3 +63,33 @@ def _compute_target_high_water_mark(
         target_prev = to_datetime(target_snapshot.node.cron_floor(target_date))
         this_prev = to_datetime(self.this_snapshot.node.cron_floor(target_date))
         return min(target_prev, this_prev)
+
+
+class HighWaterMarkExternalSensor(BaseSensorOperator):
+    def __init__(
+        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",
+        **kwargs: t.Any,
+    ):
+        super().__init__(
+            poke_interval=poke_interval,
+            timeout=timeout,
+            mode=mode,
+            **kwargs,
+        )
+        self.snapshot = snapshot
+        self.external_table_sensor_factory = external_table_sensor_factory
+
+    def poke(self, context: Context) -> t.Union[bool, PokeReturnValue]:
+        dag_run = context["dag_run"]
+        signals = self.snapshot.model.render_signals(
+            start=to_datetime(dag_run.data_interval_start),
+            end=to_datetime(dag_run.data_interval_end),
+            execution_time=now(minute_floor=False),
+        )
+        delegates = [self.external_table_sensor_factory(signal) for signal in signals]
+        return all(d.poke(context) for d in delegates)

sqlmesh/utils/date.py

@@ -214,6 +214,7 @@ def date_dict(
         kwargs[f"{prefix}_ts"] = dt.isoformat()
         kwargs[f"{prefix}_epoch"] = millis / 1000
         kwargs[f"{prefix}_millis"] = millis
+        kwargs[f"{prefix}_hour"] = dt.hour
     return kwargs
 
 

tests/core/test_model.py

@@ -2282,3 +2282,68 @@ def test_model_allow_partials():
     assert model.allow_partials
 
     assert "allow_partials TRUE" in model.render_definition()[0].sql()
+
+
+def test_signals():
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+            signals [
+                (
+                    table_name = 'table_a',
+                    ds = @end_ds,
+                ),
+                (
+                    table_name = 'table_b',
+                    ds = @end_ds,
+                    hour = @end_hour,
+                ),
+                (
+                    bool_key = True,
+                    int_key = 1,
+                    float_key = 1.0,
+                    string_key = 'string',
+                )
+            ],
+        );
+        SELECT 1;
+        """
+    )
+
+    model = load_sql_based_model(expressions)
+    assert model.signals == [
+        exp.Tuple(
+            expressions=[
+                exp.to_column("table_name").eq("table_a"),
+                exp.to_column("ds").eq(d.MacroVar(this="end_ds")),
+            ]
+        ),
+        exp.Tuple(
+            expressions=[
+                exp.to_column("table_name").eq("table_b"),
+                exp.to_column("ds").eq(d.MacroVar(this="end_ds")),
+                exp.to_column("hour").eq(d.MacroVar(this="end_hour")),
+            ]
+        ),
+        exp.Tuple(
+            expressions=[
+                exp.to_column("bool_key").eq(True),
+                exp.to_column("int_key").eq(1),
+                exp.to_column("float_key").eq(1.0),
+                exp.to_column("string_key").eq("string"),
+            ]
+        ),
+    ]
+
+    rendered_signals = model.render_signals(start="2023-01-01", end="2023-01-02 15:00:00")
+    assert rendered_signals == [
+        {"table_name": "table_a", "ds": "2023-01-02"},
+        {"table_name": "table_b", "ds": "2023-01-02", "hour": 14},
+        {"bool_key": True, "int_key": 1, "float_key": 1.0, "string_key": "string"},
+    ]
+
+    assert (
+        "signals ((table_name = 'table_a', ds = @end_ds), (table_name = 'table_b', ds = @end_ds, hour = @end_hour), (bool_key = TRUE, int_key = 1, float_key = 1.0, string_key = 'string')"
+        in model.render_definition()[0].sql()
+    )

tests/core/test_snapshot.py

@@ -120,6 +120,7 @@ def test_json(snapshot: Snapshot):
             "references": [],
             "hash_raw_query": False,
             "allow_partials": False,
+            "signals": [],
         },
         "audits": [],
         "name": "name",