Skip to content

sensors

Airflow sensor polling Snowflake Task graph completion.

SnowflakeTaskSensor

Bases: BaseSensorOperator

Poll a Snowflake Task graph until it succeeds or fails.

Uses COMPLETE_GRAPH_TASKS for graph-level status — not individual task status. Always runs in mode='reschedule' to free the Airflow worker between polls (see ADR-0002).

Parameters:

Name Type Description Default
task_fqn str

Fully-qualified Snowflake root Task name (DB.SCHEMA.TASK).

 required
snowflake_conn_id str

Airflow connection ID for the Snowflake hook.

 'snowflake_default'
**kwargs Any

Passed through to BaseSensorOperator.

 {}

Raises:

Type Description
Exception

If the graph reaches state FAILED.
Example 
from pinky_airflow import SnowflakeTaskOperator, SnowflakeTaskSensor

run  = SnowflakeTaskOperator(task_id="run_pipeline",  task_fqn="DB.SCHEMA.DAG")
wait = SnowflakeTaskSensor(task_id="wait_pipeline", task_fqn="DB.SCHEMA.DAG")

run >> wait
Source code in src/pinky_airflow/sensors.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
class SnowflakeTaskSensor(BaseSensorOperator):
    """Poll a Snowflake Task graph until it succeeds or fails.

    Uses ``COMPLETE_GRAPH_TASKS`` for graph-level status — not individual task
    status. Always runs in ``mode='reschedule'`` to free the Airflow worker
    between polls (see ADR-0002).

    Args:
        task_fqn: Fully-qualified Snowflake root Task name (``DB.SCHEMA.TASK``).
        snowflake_conn_id: Airflow connection ID for the Snowflake hook.
        **kwargs: Passed through to ``BaseSensorOperator``.

    Raises:
        Exception: If the graph reaches state ``FAILED``.

    Example:
        ```python
        from pinky_airflow import SnowflakeTaskOperator, SnowflakeTaskSensor

        run  = SnowflakeTaskOperator(task_id="run_pipeline",  task_fqn="DB.SCHEMA.DAG")
        wait = SnowflakeTaskSensor(task_id="wait_pipeline", task_fqn="DB.SCHEMA.DAG")

        run >> wait
        ```
    """

    def __init__(
        self,
        task_fqn: str,
        snowflake_conn_id: str = "snowflake_default",
        **kwargs: Any,
    ) -> None:
        super().__init__(mode="reschedule", poke_interval=30, **kwargs)
        self.task_fqn = task_fqn
        self.snowflake_conn_id = snowflake_conn_id

    def poke(self, context: Any) -> bool:
        """Return ``True`` when the graph has completed successfully."""
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
        if hook.get_records(
            f"SELECT * FROM TABLE(INFORMATION_SCHEMA.CURRENT_GRAPH_TASKS('{self.task_fqn}'))"
        ):
            return False
        failed = hook.get_records(
            f"SELECT * FROM TABLE(INFORMATION_SCHEMA.COMPLETE_GRAPH_TASKS('{self.task_fqn}'))"
            f" WHERE STATE = 'FAILED'"
        )
        if failed:
            raise Exception(f"Graph {self.task_fqn} failed: {failed}")
        return True

poke(context)

Return True when the graph has completed successfully.

Source code in src/pinky_airflow/sensors.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
def poke(self, context: Any) -> bool:
    """Return ``True`` when the graph has completed successfully."""
    hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
    if hook.get_records(
        f"SELECT * FROM TABLE(INFORMATION_SCHEMA.CURRENT_GRAPH_TASKS('{self.task_fqn}'))"
    ):
        return False
    failed = hook.get_records(
        f"SELECT * FROM TABLE(INFORMATION_SCHEMA.COMPLETE_GRAPH_TASKS('{self.task_fqn}'))"
        f" WHERE STATE = 'FAILED'"
    )
    if failed:
        raise Exception(f"Graph {self.task_fqn} failed: {failed}")
    return True