Skip to content

Features & roadmap — pinky-airflow

Update date : 2026-05-31 13:45

MVP

Component Status Description
SnowflakeTaskOperator planned Executes a Snowflake root Task, injects Airflow context via USING CONFIG
SnowflakeTaskSensor planned Polls COMPLETE_GRAPH_TASKS until graph SUCCEEDED or FAILED

SnowflakeTaskOperator

Triggers a Snowflake Task graph and injects the Airflow run context as Task config. The Snowpark SP receives the context via load_args() and sets QUERY_TAG for cost attribution.

from airflow.models import BaseOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook
import json

class PinkyTaskOperator(BaseOperator):
    def __init__(self, task_fqn: str, snowflake_conn_id: str = "snowflake_default", **kwargs):
        super().__init__(**kwargs)
        self.task_fqn = task_fqn
        self.snowflake_conn_id = snowflake_conn_id

    def execute(self, context):
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
        result = hook.get_records(f"SHOW TASKS LIKE '{self.task_fqn.split('.')[-1]}'")
        existing_config = json.loads(result[0]["config"]) if result and result[0]["config"] else {}
        existing_config["airflow"] = {
            "dag_id": context["dag_id"],
            "run_id": context["run_id"],
            "execution_date": str(context["execution_date"]),
            "task_id": context["task_id"],
        }
        hook.run(f"EXECUTE TASK {self.task_fqn} USING CONFIG = '{json.dumps(existing_config)}'")

SnowflakeTaskSensor

Polls graph-level completion (not individual task status). mode='reschedule' frees the Airflow worker between polls.

from airflow.sensors.base import BaseSensorOperator
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook

class PinkyTaskSensor(BaseSensorOperator):
    def __init__(self, task_fqn: str, snowflake_conn_id: str = "snowflake_default", **kwargs):
        super().__init__(mode="reschedule", poke_interval=30, **kwargs)
        self.task_fqn = task_fqn
        self.snowflake_conn_id = snowflake_conn_id

    def poke(self, context) -> bool:
        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

Meta DAG pattern — A → (B ∥ C) → D

from pinky_airflow import SnowflakeTaskOperator, SnowflakeTaskSensor

DB, SCHEMA = "MY_DB", "MY_SCHEMA"

run_a  = SnowflakeTaskOperator(task_id="run_a",  task_fqn=f"{DB}.{SCHEMA}.DAG_A")
wait_a = SnowflakeTaskSensor(task_id="wait_a",  task_fqn=f"{DB}.{SCHEMA}.DAG_A")
run_b  = SnowflakeTaskOperator(task_id="run_b",  task_fqn=f"{DB}.{SCHEMA}.DAG_B")
run_c  = SnowflakeTaskOperator(task_id="run_c",  task_fqn=f"{DB}.{SCHEMA}.DAG_C")
wait_b = SnowflakeTaskSensor(task_id="wait_b",  task_fqn=f"{DB}.{SCHEMA}.DAG_B")
wait_c = SnowflakeTaskSensor(task_id="wait_c",  task_fqn=f"{DB}.{SCHEMA}.DAG_C")
run_d  = SnowflakeTaskOperator(task_id="run_d",  task_fqn=f"{DB}.{SCHEMA}.DAG_D")

run_a >> wait_a >> [run_b, run_c]
run_b >> wait_b
run_c >> wait_c
[wait_b, wait_c] >> run_d

Post-MVP — ordered by priority

SnowflakeDagTriggerSP — Snowflake → Airflow trigger

Snowpark stored procedure that calls the Airflow REST API to trigger a DAG run from inside a Snowflake Task. Airflow base URL + API token stored as a Snowflake GENERIC_STRING secret, resolved via pinky-connect http.py.

DAG run dedup via a custom run_id derived from the triggering Task's QUERY_ID to prevent duplicate DAG runs on Task retry.

XCom — Snowflake Task metadata in Airflow

Push Task run metadata (query ID, duration, rows processed) as Airflow XCom so downstream tasks can branch on actual outcomes.

SnowflakeTaskRunStatus — enum

Mirror Snowflake Task run states for clean conditional branching: SUCCEEDED, FAILED, RUNNING, SKIPPED, SCHEDULED.


Out of scope

  • Data transformation in operators — that is Snowpark's job
  • Replacing Snowflake Tasks for internal orchestration
  • capture_versions migration