Skip to content

operators

Airflow operator for triggering Snowflake Task graphs.

SnowflakeTaskOperator

Bases: BaseOperator

Execute a Snowflake root Task and inject the Airflow run context.

Triggers the Task graph via EXECUTE TASK … USING CONFIG and passes the Airflow context (dag_id, run_id, execution_date, task_id) so that downstream Snowpark SPs can set QUERY_TAG for per-client cost attribution.

Parameters:

Name Type Description Default
task_fqn str

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

 required
snowflake_conn_id str

Airflow connection ID for the Snowflake hook.

 'snowflake_default'
**kwargs Any

Passed through to BaseOperator.

 {}
Example 
from pinky_airflow import SnowflakeTaskOperator

run = SnowflakeTaskOperator(
    task_id="run_pipeline",
    task_fqn="MY_DB.MY_SCHEMA.DAG_ROOT",
)
Source code in src/pinky_airflow/operators.py 
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
61
62
63
64
class SnowflakeTaskOperator(BaseOperator):
    """Execute a Snowflake root Task and inject the Airflow run context.

    Triggers the Task graph via ``EXECUTE TASK … USING CONFIG`` and passes
    the Airflow context (dag_id, run_id, execution_date, task_id) so that
    downstream Snowpark SPs can set ``QUERY_TAG`` for per-client cost attribution.

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

    Example:
        ```python
        from pinky_airflow import SnowflakeTaskOperator

        run = SnowflakeTaskOperator(
            task_id="run_pipeline",
            task_fqn="MY_DB.MY_SCHEMA.DAG_ROOT",
        )
        ```
    """

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

    def execute(self, context: Any) -> None:
        """Trigger the Snowflake Task graph with the current Airflow context injected."""
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
        result = hook.get_records(
            f"SHOW TASKS LIKE '{self.task_fqn.split('.')[-1]}'"
        )
        existing_config: dict[str, Any] = (
            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)}'"
        )

execute(context)

Trigger the Snowflake Task graph with the current Airflow context injected.

Source code in src/pinky_airflow/operators.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
def execute(self, context: Any) -> None:
    """Trigger the Snowflake Task graph with the current Airflow context injected."""
    hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
    result = hook.get_records(
        f"SHOW TASKS LIKE '{self.task_fqn.split('.')[-1]}'"
    )
    existing_config: dict[str, Any] = (
        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)}'"
    )