Implementing Airflow Custom Operators and Hooks for Snowflake and dbt: Extending the Modern Data Stack with Reusable, Production-Grade Integrations

Introduction

You're three months into a new data engineering role. The Airflow DAGs that orchestrate your modern data stack are a sprawling collection of PythonOperator calls, each one containing slightly different logic for connecting to Snowflake, slightly different retry semantics, and slightly different ways of invoking dbt. When something breaks at 2 AM, you're not debugging a pipeline — you're archaeologically excavating someone's one-off implementation choices, scattered across a dozen DAG files with no consistent interface.

This is the situation that custom Airflow operators and hooks are designed to prevent. When you invest the time to build proper, reusable integrations, you're not just cleaning up code — you're creating a contract between your orchestration layer and the underlying systems it touches. Every engineer on your team invokes Snowflake the same way. Every dbt run exposes the same interface. Failures surface through the same error handling paths. Production incidents become investigations instead of treasure hunts.

In this lesson, you'll build production-grade custom Airflow operators and hooks for both Snowflake and dbt from scratch. You'll understand why the abstraction layers are designed the way they are, where the official provider packages make trade-offs that don't fit every use case, and how to extend those abstractions to match the realities of your specific stack.

What you'll learn:

• How Airflow's Hook and Operator class hierarchy works at the level you need to extend it effectively
• How to build a SnowflakeHook extension that handles multi-role authentication, query result streaming, and warehouse resizing
• How to implement a dbtOperator family that wraps the dbt CLI with proper logging, artifact capture, and selective model execution
• How to pass metadata between operators using XCom with type-safe interfaces
• How to structure your custom integrations as an internal package for sharing across multiple DAGs and teams
• How to test custom operators in isolation with mocked connections

---

Prerequisites

You should be comfortable with:

• Airflow DAGs, task dependencies, and the executor model — if you need a refresher, start with Data Pipeline Orchestration with Airflow
• The dbt fundamentals and project structure covered in dbt Fundamentals: Transform Data with SQL in Your Warehouse
• Snowflake concepts including warehouses, roles, and connection parameters
• Python class inheritance and context managers
• Basic familiarity with subprocess management in Python

---

How Airflow's Hook and Operator Architecture Actually Works

Before writing a single line of custom code, you need to understand the design contract that Airflow's provider system is built on. Operators and Hooks are not interchangeable — they have distinct responsibilities that matter enormously once you're maintaining these integrations at scale.

Hooks are connection managers. A Hook's job is to abstract the mechanics of establishing and managing a connection to an external system. It holds no task-level state. It doesn't know what DAG it's running in. It just knows how to talk to a specific system — Snowflake, an S3 bucket, a Slack webhook — and it exposes methods that encapsulate the connection lifecycle.

Operators are task managers. An Operator defines a discrete unit of work: what to do, when to do it, and what to return. Operators use Hooks to actually communicate with external systems. An Operator knows about retries, XCom, task context, and Airflow's execution model. A Hook knows about connection strings, authentication, and API semantics.

This separation matters for testing, for reuse, and for maintenance. When you write a SnowflakeHook, you can test it against a real or mocked Snowflake connection independently of any DAG. When you write a SnowflakeOperator, its test can use a mocked Hook.

Airflow's BaseHook class provides:

• get_connection(conn_id) — retrieves a Connection object from the Airflow metadata database
• Connection caching per-thread
• A standard __init__ signature that accepts conn_id

Airflow's BaseOperator class provides:

• execute(context) — the method you override, receives the full task context dictionary
• Template rendering for template_fields
• Retry logic, timeout handling, and callback hooks
• xcom_push and xcom_pull for inter-task data sharing

The BaseOperator.__init__ signature has grown significantly over Airflow's history. Always call super().__init__(**kwargs) to forward base operator parameters cleanly:

from airflow.models import BaseOperator

class MyOperator(BaseOperator):
    def __init__(self, my_param: str, **kwargs):
        super().__init__(**kwargs)
        self.my_param = my_param
    
    def execute(self, context):
        # Your logic here
        pass

Missing **kwargs forwarding is one of the most common mistakes beginners make with custom operators — parameters like retries, retry_delay, on_failure_callback, and task_id would silently fail to reach BaseOperator.__init__.

---

Building a Production SnowflakeHook

The apache-airflow-providers-snowflake package ships a functional SnowflakeHook, but it makes several choices that don't fit every production environment: it doesn't handle multi-warehouse routing, it doesn't stream large result sets efficiently, and it doesn't expose warehouse management operations. We'll build an extension that adds those capabilities while still inheriting from the official provider.

First, set up your project structure. In a production environment, custom operators live in a separate Python package that's installed into your Airflow environment:

airflow_custom_integrations/
├── __init__.py
├── hooks/
│   ├── __init__.py
│   └── snowflake_hook.py
├── operators/
│   ├── __init__.py
│   ├── snowflake_operator.py
│   └── dbt_operator.py
└── tests/
    ├── __init__.py
    ├── test_snowflake_hook.py
    └── test_dbt_operator.py

Now let's build the hook:

# hooks/snowflake_hook.py

import logging
from contextlib import contextmanager
from typing import Any, Generator, Iterator, Optional, Union

import snowflake.connector
from airflow.providers.snowflake.hooks.snowflake import SnowflakeHook as BaseSnowflakeHook
from snowflake.connector import DictCursor

log = logging.getLogger(__name__)


class SnowflakeHook(BaseSnowflakeHook):
    """
    Extended Snowflake Hook with support for:
    - Role-switching mid-session for privilege escalation
    - Streaming result sets to avoid memory exhaustion on large queries
    - Warehouse resize and resume/suspend operations
    - Query tagging for cost attribution
    """

    def __init__(
        self,
        snowflake_conn_id: str = "snowflake_default",
        warehouse: Optional[str] = None,
        database: Optional[str] = None,
        role: Optional[str] = None,
        schema: Optional[str] = None,
        query_tag: Optional[str] = None,
        autocommit: bool = True,
    ):
        super().__init__(
            snowflake_conn_id=snowflake_conn_id,
            warehouse=warehouse,
            database=database,
            role=role,
            schema=schema,
            autocommit=autocommit,
        )
        self.query_tag = query_tag

    @contextmanager
    def get_tagged_conn(self) -> Generator[snowflake.connector.SnowflakeConnection, None, None]:
        """
        Context manager that yields a connection with a query tag applied.
        Query tags are visible in Snowflake's QUERY_HISTORY and are invaluable
        for cost attribution and debugging.
        """
        conn = self.get_conn()
        try:
            if self.query_tag:
                conn.cursor().execute(
                    f"ALTER SESSION SET QUERY_TAG = '{self.query_tag}'"
                )
                log.info("Set Snowflake query tag: %s", self.query_tag)
            yield conn
        finally:
            conn.close()

    def run_with_streaming(
        self,
        sql: str,
        parameters: Optional[dict] = None,
        chunk_size: int = 10_000,
    ) -> Iterator[list[dict]]:
        """
        Execute a query and yield results in chunks using server-side cursors.
        This prevents loading multi-million-row result sets into memory at once.
        
        Yields chunks as lists of dictionaries (column_name -> value).
        """
        with self.get_tagged_conn() as conn:
            with conn.cursor(DictCursor) as cur:
                log.info("Executing streaming query (chunk_size=%d)", chunk_size)
                cur.execute(sql, parameters or {})
                
                while True:
                    rows = cur.fetchmany(chunk_size)
                    if not rows:
                        break
                    log.debug("Fetched chunk of %d rows", len(rows))
                    yield rows

    def execute_with_role(
        self,
        sql: str,
        role: str,
        parameters: Optional[dict] = None,
    ) -> list[dict]:
        """
        Execute SQL using a specific role, then revert to the session role.
        Useful for tasks that need elevated privileges (e.g., grants, cloning)
        without permanently changing the session's role.
        """
        with self.get_tagged_conn() as conn:
            with conn.cursor(DictCursor) as cur:
                cur.execute(f"USE ROLE {role}")
                log.info("Switched to role: %s", role)
                try:
                    cur.execute(sql, parameters or {})
                    return cur.fetchall()
                finally:
                    # Revert to the connection's configured role
                    original_role = self.role or conn.role
                    cur.execute(f"USE ROLE {original_role}")
                    log.info("Reverted to role: %s", original_role)

    def resize_warehouse(
        self,
        warehouse_name: str,
        size: str,
        admin_role: str = "SYSADMIN",
    ) -> None:
        """
        Resize a Snowflake virtual warehouse. Accepts standard Snowflake
        size strings: X-SMALL, SMALL, MEDIUM, LARGE, X-LARGE, etc.
        
        Requires a role with MODIFY privilege on the warehouse.
        """
        valid_sizes = {
            "X-SMALL", "SMALL", "MEDIUM", "LARGE",
            "X-LARGE", "2X-LARGE", "3X-LARGE", "4X-LARGE",
        }
        size_upper = size.upper()
        if size_upper not in valid_sizes:
            raise ValueError(
                f"Invalid warehouse size '{size}'. Must be one of: {valid_sizes}"
            )
        
        sql = f"ALTER WAREHOUSE {warehouse_name} SET WAREHOUSE_SIZE = '{size_upper}'"
        self.execute_with_role(sql=sql, role=admin_role)
        log.info("Resized warehouse %s to %s", warehouse_name, size_upper)

    def get_query_status(self, query_id: str) -> dict:
        """
        Poll the status of a submitted query by ID.
        Useful when submitting async queries and checking later.
        """
        sql = """
            SELECT 
                query_id,
                query_text,
                execution_status,
                error_message,
                total_elapsed_time,
                bytes_scanned,
                credits_used_cloud_services
            FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY_BY_SESSION())
            WHERE query_id = %(query_id)s
        """
        with self.get_tagged_conn() as conn:
            with conn.cursor(DictCursor) as cur:
                cur.execute(sql, {"query_id": query_id})
                result = cur.fetchone()
                if not result:
                    raise ValueError(f"No query found with ID: {query_id}")
                return result

---

Building the SnowflakeOperator Family

With a solid hook in place, operators become thin orchestration wrappers. Here's how to build a flexible SnowflakeQueryOperator that handles templating, result capture, and cost attribution:

# operators/snowflake_operator.py

import logging
from typing import Any, Optional, Sequence, Union

from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

from airflow_custom_integrations.hooks.snowflake_hook import SnowflakeHook

log = logging.getLogger(__name__)


class SnowflakeQueryOperator(BaseOperator):
    """
    Executes one or more SQL statements against Snowflake.
    
    Features:
    - Jinja2 templating on sql, warehouse, and query_tag fields
    - Optional result capture into XCom (disabled by default for large queries)
    - Per-task warehouse override
    - Automatic query tagging with DAG/task metadata
    """

    template_fields: Sequence[str] = ("sql", "warehouse", "query_tag", "parameters")
    template_ext: Sequence[str] = (".sql",)
    ui_color = "#4a86e8"  # Blue — renders in the Airflow task graph

    @apply_defaults
    def __init__(
        self,
        sql: Union[str, list[str]],
        snowflake_conn_id: str = "snowflake_default",
        warehouse: Optional[str] = None,
        database: Optional[str] = None,
        schema: Optional[str] = None,
        role: Optional[str] = None,
        query_tag: Optional[str] = None,
        parameters: Optional[dict] = None,
        do_xcom_push: bool = False,
        autocommit: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.sql = sql
        self.snowflake_conn_id = snowflake_conn_id
        self.warehouse = warehouse
        self.database = database
        self.schema = schema
        self.role = role
        self.parameters = parameters or {}
        self.do_xcom_push = do_xcom_push
        self.autocommit = autocommit
        
        # Auto-generate a query tag if none provided
        self.query_tag = query_tag

    def execute(self, context: dict) -> Optional[list]:
        # Build a rich query tag if one wasn't provided
        effective_tag = self.query_tag or (
            f"airflow|dag={context['dag'].dag_id}"
            f"|task={context['task_instance'].task_id}"
            f"|run={context['run_id']}"
        )

        hook = SnowflakeHook(
            snowflake_conn_id=self.snowflake_conn_id,
            warehouse=self.warehouse,
            database=self.database,
            schema=self.schema,
            role=self.role,
            query_tag=effective_tag,
            autocommit=self.autocommit,
        )

        statements = self.sql if isinstance(self.sql, list) else [self.sql]
        results = []

        with hook.get_tagged_conn() as conn:
            for statement in statements:
                log.info("Executing SQL:\n%s", statement)
                with conn.cursor() as cur:
                    cur.execute(statement, self.parameters)
                    if cur.description:  # SELECT queries have description
                        rows = cur.fetchall()
                        results.append(rows)
                        log.info("Query returned %d rows", len(rows))
                    else:
                        results.append(None)

        if self.do_xcom_push and results:
            return results[-1]  # Only push the last result set
        return None


class SnowflakeWarehouseOperator(BaseOperator):
    """
    Resize a Snowflake warehouse before a resource-intensive task,
    then revert it afterward. Designed to be used as a pair:
    
        resize_up >> heavy_transform >> resize_down
    """

    template_fields: Sequence[str] = ("warehouse_name", "target_size")
    ui_color = "#ff9800"  # Orange — visually distinct in the DAG graph

    @apply_defaults
    def __init__(
        self,
        warehouse_name: str,
        target_size: str,
        admin_role: str = "SYSADMIN",
        snowflake_conn_id: str = "snowflake_default",
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.warehouse_name = warehouse_name
        self.target_size = target_size
        self.admin_role = admin_role
        self.snowflake_conn_id = snowflake_conn_id

    def execute(self, context: dict) -> None:
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id)
        hook.resize_warehouse(
            warehouse_name=self.warehouse_name,
            size=self.target_size,
            admin_role=self.admin_role,
        )
        log.info(
            "Warehouse %s resized to %s",
            self.warehouse_name,
            self.target_size,
        )

Notice how SnowflakeQueryOperator automatically generates a query tag from Airflow's task context when one isn't provided. Every query that runs through this operator becomes traceable back to the exact DAG run that triggered it. When you're investigating a Snowflake credit spike at the end of the month, this context information is invaluable.

---

Building the dbt Operator Family

The dbt integration is where things get genuinely interesting — and where the official provider packages often fall short. The DbtRunOperator in astronomer-cosmos and apache-airflow-providers-dbt-cloud both make strong architectural choices that may not fit your workflow. Understanding how to build your own gives you the flexibility to handle the specific patterns your team actually uses.

The core of the dbt operator is subprocess management. dbt is a CLI tool, and the cleanest way to invoke it from Airflow is via subprocess. The key challenges are:

1. Streaming stdout/stderr in real time (so logs appear in the Airflow UI as the run progresses)
2. Capturing the exit code and translating non-zero exits to task failures
3. Parsing the run_results.json artifact to extract per-model execution results
4. Injecting the correct environment variables for profile configuration

# operators/dbt_operator.py

import json
import logging
import os
import subprocess
from pathlib import Path
from typing import Any, Optional, Sequence

from airflow.exceptions import AirflowException
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults

log = logging.getLogger(__name__)


class DbtBaseOperator(BaseOperator):
    """
    Base class for dbt operators. Handles subprocess execution,
    log streaming, artifact capture, and environment injection.
    
    Not intended to be instantiated directly.
    """

    template_fields: Sequence[str] = ("dbt_vars", "models", "exclude")

    @apply_defaults
    def __init__(
        self,
        project_dir: str,
        profiles_dir: Optional[str] = None,
        target: str = "prod",
        models: Optional[str] = None,
        exclude: Optional[str] = None,
        dbt_vars: Optional[dict] = None,
        env_vars: Optional[dict] = None,
        do_xcom_push_artifacts: bool = True,
        dbt_executable: str = "dbt",
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.project_dir = project_dir
        self.profiles_dir = profiles_dir
        self.target = target
        self.models = models
        self.exclude = exclude
        self.dbt_vars = dbt_vars or {}
        self.env_vars = env_vars or {}
        self.do_xcom_push_artifacts = do_xcom_push_artifacts
        self.dbt_executable = dbt_executable

    def _build_command(self, subcommand: str, extra_args: Optional[list] = None) -> list[str]:
        """
        Build the dbt CLI command list. Using a list rather than a
        shell string avoids shell injection vulnerabilities and makes
        argument handling more reliable with spaces in paths.
        """
        cmd = [
            self.dbt_executable,
            subcommand,
            "--project-dir", self.project_dir,
            "--target", self.target,
            "--no-use-colors",  # Cleaner logs in Airflow's UI
        ]

        if self.profiles_dir:
            cmd.extend(["--profiles-dir", self.profiles_dir])

        if self.models:
            cmd.extend(["--select", self.models])

        if self.exclude:
            cmd.extend(["--exclude", self.exclude])

        if self.dbt_vars:
            # dbt expects vars as a YAML string
            vars_str = json.dumps(self.dbt_vars)
            cmd.extend(["--vars", vars_str])

        if extra_args:
            cmd.extend(extra_args)

        return cmd

    def _build_env(self) -> dict:
        """
        Build the environment for the subprocess. Starts from the current
        process environment, then layers in any operator-level overrides.
        """
        env = os.environ.copy()
        env.update(self.env_vars)
        return env

    def _run_command(self, cmd: list[str]) -> tuple[int, str]:
        """
        Execute a command via subprocess, streaming output line-by-line
        to the Airflow task log. Returns (returncode, last_error_line).
        """
        log.info("Executing command: %s", " ".join(cmd))
        
        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,  # Merge stderr into stdout
            env=self._build_env(),
            cwd=self.project_dir,
            text=True,
            bufsize=1,  # Line-buffered
        )

        output_lines = []
        last_error_line = ""

        for line in iter(process.stdout.readline, ""):
            line = line.rstrip()
            output_lines.append(line)
            log.info(line)
            
            # Track the last non-empty line as a likely error summary
            if line.strip():
                last_error_line = line

        process.stdout.close()
        return_code = process.wait()

        return return_code, last_error_line

    def _read_run_results(self) -> Optional[dict]:
        """
        Read and parse the dbt run_results.json artifact from the target
        directory. Returns None if the file doesn't exist (e.g., compile
        failures before execution begins).
        """
        results_path = Path(self.project_dir) / "target" / "run_results.json"
        
        if not results_path.exists():
            log.warning("run_results.json not found at %s", results_path)
            return None

        with open(results_path) as f:
            return json.load(f)

    def _parse_run_results(self, run_results: dict) -> dict:
        """
        Extract a summary of model-level results from the artifact.
        Returns a dict with counts and any failed model names.
        """
        summary = {
            "elapsed_time": run_results.get("elapsed_time", 0),
            "success": 0,
            "error": 0,
            "skip": 0,
            "warn": 0,
            "failed_nodes": [],
        }

        for result in run_results.get("results", []):
            status = result.get("status", "").lower()
            node_id = result.get("unique_id", "unknown")

            if status in ("success", "pass"):
                summary["success"] += 1
            elif status == "error":
                summary["error"] += 1
                summary["failed_nodes"].append(node_id)
                log.error(
                    "Node %s failed: %s",
                    node_id,
                    result.get("message", "No message"),
                )
            elif status == "skipped":
                summary["skip"] += 1
            elif status == "warn":
                summary["warn"] += 1
                log.warning("Node %s completed with warnings", node_id)

        return summary

    def execute(self, context: dict) -> Any:
        raise NotImplementedError("Subclasses must implement execute()")


class DbtRunOperator(DbtBaseOperator):
    """
    Execute `dbt run` for a specified set of models.
    
    Captures run_results.json and pushes a results summary to XCom
    under the key 'dbt_run_results' for downstream consumption.
    """

    ui_color = "#f0a030"  # dbt orange

    @apply_defaults
    def __init__(self, full_refresh: bool = False, **kwargs):
        super().__init__(**kwargs)
        self.full_refresh = full_refresh

    def execute(self, context: dict) -> dict:
        extra_args = ["--full-refresh"] if self.full_refresh else []
        cmd = self._build_command("run", extra_args=extra_args)

        return_code, last_error = self._run_command(cmd)

        # Always try to read artifacts, even on failure
        run_results = self._read_run_results()
        summary = self._parse_run_results(run_results) if run_results else {}

        if self.do_xcom_push_artifacts and run_results:
            context["task_instance"].xcom_push(
                key="dbt_run_results",
                value=summary,
            )

        if return_code != 0:
            failed = summary.get("failed_nodes", [])
            raise AirflowException(
                f"dbt run failed with exit code {return_code}. "
                f"Failed nodes: {failed}. Last output: {last_error}"
            )

        log.info(
            "dbt run completed: %d success, %d error, %d skip",
            summary.get("success", 0),
            summary.get("error", 0),
            summary.get("skip", 0),
        )
        return summary


class DbtTestOperator(DbtBaseOperator):
    """
    Execute `dbt test` and fail the task if any tests fail.
    
    Pushes test results to XCom for use by alerting operators downstream.
    """

    ui_color = "#27ae60"  # Green — testing is healthy

    def execute(self, context: dict) -> dict:
        cmd = self._build_command("test")
        return_code, last_error = self._run_command(cmd)

        run_results = self._read_run_results()
        summary = self._parse_run_results(run_results) if run_results else {}

        if self.do_xcom_push_artifacts and run_results:
            context["task_instance"].xcom_push(
                key="dbt_test_results",
                value=summary,
            )

        if return_code != 0:
            failed_tests = summary.get("failed_nodes", [])
            raise AirflowException(
                f"dbt test failed. Failed tests: {failed_tests}"
            )

        log.info(
            "dbt tests passed: %d pass, %d warn, %d fail",
            summary.get("success", 0),
            summary.get("warn", 0),
            summary.get("error", 0),
        )
        return summary


class DbtSourceFreshnessOperator(DbtBaseOperator):
    """
    Execute `dbt source freshness` and optionally fail the task
    if any sources exceed their configured freshness thresholds.
    """

    ui_color = "#8e44ad"  # Purple

    @apply_defaults
    def __init__(
        self,
        fail_on_error: bool = True,
        fail_on_warn: bool = False,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.fail_on_error = fail_on_error
        self.fail_on_warn = fail_on_warn

    def execute(self, context: dict) -> dict:
        cmd = self._build_command("source freshness")
        return_code, last_error = self._run_command(cmd)

        # source freshness writes to sources.json, not run_results.json
        sources_path = (
            Path(self.project_dir) / "target" / "sources.json"
        )
        sources_data = {}

        if sources_path.exists():
            with open(sources_path) as f:
                sources_data = json.load(f)
            
            if self.do_xcom_push_artifacts:
                context["task_instance"].xcom_push(
                    key="dbt_source_freshness",
                    value=sources_data,
                )

        if self.fail_on_error and return_code != 0:
            raise AirflowException(
                f"dbt source freshness check failed: {last_error}"
            )

        return sources_data

---

Wiring It All Together: A Real DAG

Here's how these operators compose into a production DAG for a typical ELT pipeline. This example loads data from Snowflake staging tables, runs dbt transformations, runs tests, and resizes the warehouse around the most expensive transformation step:

# dags/analytics_daily_pipeline.py

from datetime import datetime, timedelta

from airflow import DAG
from airflow.models import Variable

from airflow_custom_integrations.operators.snowflake_operator import (
    SnowflakeQueryOperator,
    SnowflakeWarehouseOperator,
)
from airflow_custom_integrations.operators.dbt_operator import (
    DbtRunOperator,
    DbtTestOperator,
    DbtSourceFreshnessOperator,
)

DBT_PROJECT_DIR = Variable.get("dbt_project_dir", default_var="/opt/dbt/analytics")
DBT_PROFILES_DIR = Variable.get("dbt_profiles_dir", default_var="/opt/dbt/profiles")
TRANSFORM_WAREHOUSE = "TRANSFORM_WH"
SNOWFLAKE_CONN = "snowflake_analytics"

default_args = {
    "owner": "data-engineering",
    "retries": 2,
    "retry_delay": timedelta(minutes=5),
    "retry_exponential_backoff": True,
    "on_failure_callback": lambda context: log_failure_to_slack(context),
}

with DAG(
    dag_id="analytics_daily_pipeline",
    start_date=datetime(2024, 1, 1),
    schedule_interval="0 6 * * *",
    default_args=default_args,
    catchup=False,
    tags=["analytics", "dbt", "snowflake"],
    doc_md="""
    ## Analytics Daily Pipeline
    
    Runs the full analytics transformation pipeline:
    1. Checks source freshness
    2. Runs staging transformations at X-SMALL
    3. Upsizes warehouse and runs heavy intermediate models
    4. Runs mart models
    5. Executes all dbt tests
    6. Downsizes warehouse
    """,
) as dag:

    # Step 1: Check that source data is fresh enough to run
    check_freshness = DbtSourceFreshnessOperator(
        task_id="check_source_freshness",
        project_dir=DBT_PROJECT_DIR,
        profiles_dir=DBT_PROFILES_DIR,
        target="prod",
        fail_on_error=True,
        fail_on_warn=False,
    )

    # Step 2: Run staging layer at small warehouse
    run_staging = DbtRunOperator(
        task_id="run_staging_models",
        project_dir=DBT_PROJECT_DIR,
        profiles_dir=DBT_PROFILES_DIR,
        target="prod",
        models="tag:staging",
        dbt_vars={
            "run_date": "{{ ds }}",
            "dag_run_id": "{{ run_id }}",
        },
        query_tag="airflow|dag=analytics_daily_pipeline|layer=staging",
    )

    # Step 3: Upsize for heavy intermediate processing
    resize_up = SnowflakeWarehouseOperator(
        task_id="resize_warehouse_large",
        warehouse_name=TRANSFORM_WAREHOUSE,
        target_size="LARGE",
        snowflake_conn_id=SNOWFLAKE_CONN,
    )

    # Step 4: Run intermediate models while warehouse is large
    run_intermediate = DbtRunOperator(
        task_id="run_intermediate_models",
        project_dir=DBT_PROJECT_DIR,
        profiles_dir=DBT_PROFILES_DIR,
        target="prod",
        models="tag:intermediate",
        dbt_vars={"run_date": "{{ ds }}"},
    )

    # Step 5: Run mart models
    run_marts = DbtRunOperator(
        task_id="run_mart_models",
        project_dir=DBT_PROJECT_DIR,
        profiles_dir=DBT_PROFILES_DIR,
        target="prod",
        models="tag:mart",
        dbt_vars={"run_date": "{{ ds }}"},
    )

    # Step 6: Downsize after expensive work is done
    resize_down = SnowflakeWarehouseOperator(
        task_id="resize_warehouse_small",
        warehouse_name=TRANSFORM_WAREHOUSE,
        target_size="X-SMALL",
        snowflake_conn_id=SNOWFLAKE_CONN,
        trigger_rule="all_done",  # Downsize even if transforms fail
    )

    # Step 7: Run all dbt tests
    run_tests = DbtTestOperator(
        task_id="run_dbt_tests",
        project_dir=DBT_PROJECT_DIR,
        profiles_dir=DBT_PROFILES_DIR,
        target="prod",
    )

    # Step 8: Record pipeline metadata for observability
    record_completion = SnowflakeQueryOperator(
        task_id="record_pipeline_completion",
        sql="""
            INSERT INTO analytics.pipeline_runs 
                (dag_id, run_id, run_date, completed_at, status)
            VALUES 
                ('analytics_daily_pipeline', '{{ run_id }}', '{{ ds }}', 
                 CURRENT_TIMESTAMP(), 'success')
        """,
        snowflake_conn_id=SNOWFLAKE_CONN,
    )

    # Define the DAG shape
    (
        check_freshness
        >> run_staging
        >> resize_up
        >> run_intermediate
        >> run_marts
        >> resize_down
        >> run_tests
        >> record_completion
    )

---

Passing Results Between Operators with XCom

The dbt_run_results XCom pushed by DbtRunOperator becomes genuinely useful when you build downstream operators that consume it. Here's an alerting pattern that sends a Slack summary only when models fail:

from airflow.operators.python import PythonOperator

def send_dbt_summary(**context):
    ti = context["task_instance"]
    
    # Pull results from the run_marts task
    mart_results = ti.xcom_pull(
        task_ids="run_mart_models",
        key="dbt_run_results",
    )
    
    if not mart_results:
        log.warning("No dbt run results found in XCom")
        return
    
    failed = mart_results.get("failed_nodes", [])
    elapsed = mart_results.get("elapsed_time", 0)
    
    message = (
        f"*dbt mart run complete* — "
        f"{mart_results['success']} success, "
        f"{mart_results['error']} error, "
        f"{mart_results['skip']} skip "
        f"({elapsed:.1f}s)"
    )
    
    if failed:
        message += f"\n*Failed nodes:*\n" + "\n".join(f"• `{n}`" for n in failed)
    
    # Post to Slack (implementation details omitted)
    slack_client.chat_postMessage(channel="#data-alerts", text=message)

---

Structuring Your Package for Production Distribution

Custom operators only deliver their full value when they're shared across your organization's DAG codebase. Here's how to package them properly:

# setup.py
from setuptools import setup, find_packages

setup(
    name="airflow-custom-integrations",
    version="1.4.2",
    packages=find_packages(),
    install_requires=[
        "apache-airflow>=2.5.0",
        "apache-airflow-providers-snowflake>=4.0.0",
        "snowflake-connector-python>=3.0.0",
    ],
    extras_require={
        "dev": [
            "pytest",
            "pytest-mock",
            "apache-airflow[sqlite]",
        ]
    },
    entry_points={
        "apache_airflow_provider": [
            "provider_info = airflow_custom_integrations:get_provider_info",
        ],
    },
)

# airflow_custom_integrations/__init__.py

def get_provider_info():
    """
    Registers this package as an Airflow provider.
    Enables the provider to appear in the Airflow UI under
    Admin > Providers, and ensures connection types are registered.
    """
    return {
        "package-name": "airflow-custom-integrations",
        "name": "Custom Snowflake + dbt Integrations",
        "description": "Production-grade operators and hooks for Snowflake and dbt",
        "versions": ["1.4.2"],
    }

When your Airflow workers install this package, all operators and hooks become available to any DAG that imports from airflow_custom_integrations. Version pinning in your Airflow image's requirements.txt ensures all environments run the same operator behavior.

---

Testing Custom Operators in Isolation

Untested custom operators are a production time bomb. Here's a testing pattern that mocks connections without needing a live Snowflake instance:

# tests/test_snowflake_operator.py

from unittest.mock import MagicMock, patch
import pytest

from airflow.models import DagBag, TaskInstance
from airflow.utils import timezone

from airflow_custom_integrations.operators.snowflake_operator import SnowflakeQueryOperator


@pytest.fixture
def mock_context():
    """Minimal Airflow task context for testing execute()."""
    mock_dag = MagicMock()
    mock_dag.dag_id = "test_dag"
    
    mock_ti = MagicMock()
    mock_ti.task_id = "test_task"
    
    return {
        "dag": mock_dag,
        "task_instance": mock_ti,
        "run_id": "manual__2024-01-01T00:00:00+00:00",
    }


@patch("airflow_custom_integrations.hooks.snowflake_hook.SnowflakeHook.get_conn")
def test_query_operator_executes_sql(mock_get_conn, mock_context):
    """Verify the operator sends the correct SQL to the connection."""
    mock_conn = MagicMock()
    mock_cursor = MagicMock()
    mock_cursor.description = None  # Simulate a DML statement
    mock_conn.cursor.return_value.__enter__ = MagicMock(return_value=mock_cursor)
    mock_conn.cursor.return_value.__exit__ = MagicMock(return_value=False)
    mock_get_conn.return_value.__enter__ = MagicMock(return_value=mock_conn)
    mock_get_conn.return_value.__exit__ = MagicMock(return_value=False)
    
    operator = SnowflakeQueryOperator(
        task_id="test_task",
        sql="INSERT INTO test_table VALUES (1, 'hello')",
        snowflake_conn_id="snowflake_default",
    )
    
    operator.execute(mock_context)
    
    mock_cursor.execute.assert_called_once_with(
        "INSERT INTO test_table VALUES (1, 'hello')",
        {},
    )


@patch("airflow_custom_integrations.operators.dbt_operator.DbtBaseOperator._run_command")
@patch("airflow_custom_integrations.operators.dbt_operator.DbtBaseOperator._read_run_results")
def test_dbt_run_operator_raises_on_failure(
    mock_read_results, mock_run_command, mock_context
):
    """Verify that a non-zero exit code raises AirflowException."""
    from airflow.exceptions import AirflowException
    from airflow_custom_integrations.operators.dbt_operator import DbtRunOperator
    
    mock_run_command.return_value = (1, "Compilation error in model 'fct_orders'")
    mock_read_results.return_value = {
        "elapsed_time": 12.4,
        "results": [
            {
                "unique_id": "model.analytics.fct_orders",
                "status": "error",
                "message": "Compilation error",
            }
        ],
    }
    
    operator = DbtRunOperator(
        task_id="test_dbt_run",
        project_dir="/fake/project",
        models="fct_orders",
    )
    
    with pytest.raises(AirflowException, match="dbt run failed"):
        operator.execute(mock_context)

---

Hands-On Exercise

Extend the DbtRunOperator to support slim CI — running only the models that changed in a Git diff. This is one of the most powerful optimizations available for teams with large dbt projects, and it's described in detail in automating dbt environment promotion with CI/CD pipelines.

Your task:

1. Add a state_path parameter to DbtRunOperator that, when provided, passes --state {state_path} and --select state:modified+ to the dbt CLI. This tells dbt to run only models whose compiled SQL has changed relative to the state artifacts in state_path.

2. Add a DbtCompileOperator that runs dbt compile and pushes the path to the target/ directory as an XCom value, so a downstream DbtRunOperator can use it as its state.

3. Write a test that verifies when state_path is provided, the --state flag appears in the constructed command, and when it's not provided, the command doesn't include it.

Expected code additions:

# In DbtRunOperator.__init__:
self.state_path = state_path  # Optional[str], default None

# In DbtRunOperator._build_command:
if self.state_path:
    extra_args.extend(["--state", self.state_path, "--select", "state:modified+"])

The DbtCompileOperator should follow the same pattern as DbtTestOperator but call _build_command("compile") and push the target directory path to XCom after a successful run.

---

Common Mistakes & Troubleshooting

1. Forgetting **kwargs in __init__

Your operator silently ignores retries, retry_delay, on_success_callback, and other base operator parameters. Always include **kwargs and call super().__init__(**kwargs).

2. Using shell=True in subprocess

# WRONG — shell injection risk, platform-dependent quoting issues
subprocess.run(f"dbt run --project-dir {self.project_dir}", shell=True)

# RIGHT — list form, explicit argument handling
subprocess.run(["dbt", "run", "--project-dir", self.project_dir])

3. Missing trigger_rule on cleanup tasks

As shown in the warehouse resize example, cleanup tasks — downsize operations, temp table drops, lock releases — must use trigger_rule="all_done" or they'll be skipped when upstream tasks fail.

4. Pushing large objects to XCom

The Airflow metadata database has practical limits on XCom size. Full dbt manifest.json files can be 50MB+ for large projects. Push the artifact path to object storage (S3, GCS), not the content.

5. Hardcoding connection IDs in operator logic

# WRONG — brittle, forces every team to use the same connection name
hook = SnowflakeHook(snowflake_conn_id="snowflake_default")

# RIGHT — parameterize it, provide a sensible default
def __init__(self, snowflake_conn_id: str = "snowflake_default", **kwargs):
    self.snowflake_conn_id = snowflake_conn_id

6. Not handling dbt's profile environment variables

dbt resolves profiles from profiles.yml, but in production you often want to inject Snowflake credentials via environment variables rather than mounting a profiles file. Make sure your _build_env method correctly injects DBT_SNOWFLAKE_ACCOUNT, DBT_SNOWFLAKE_USER, DBT_SNOWFLAKE_PASSWORD (or equivalent for key-pair auth) from Airflow's secret backend. This connects tightly with configuring role-based access control in Snowflake — your operator needs to use the right role for each task layer.

7. Subprocess blocking the Airflow worker

For very long-running dbt runs, the subprocess pattern above holds the Airflow worker thread for the entire duration. This is generally fine for typical dbt runs (minutes), but for extreme cases (multi-hour full refreshes), consider using the dbt Cloud API via DbtCloudHook instead, which submits a job and polls rather than blocking.

8. Template fields and non-string types

Airflow renders template_fields as Jinja2 templates, but this only works on strings. If you include a dict field (like dbt_vars) in template_fields, Airflow will try to stringify and render it, which produces correct results for simple cases but breaks when the dict values contain Jinja syntax. Test your template rendering explicitly in CI.

---

Summary & Next Steps

You've now built a complete, production-grade integration layer for Snowflake and dbt in Airflow. The key architectural decisions we made:

• Hooks encapsulate connection management — SnowflakeHook owns the authentication, query tagging, and session management logic, separate from any task-level orchestration
• Operators are thin task wrappers — they define what to do and how to fail gracefully, delegating the actual system interaction to hooks
• Template fields enable DAG-level flexibility — SQL, model selectors, and variables can all be rendered from Airflow context, enabling dynamic pipelines without code changes
• Artifact capture creates observability — by reading run_results.json after every dbt invocation, you have structured data about every model execution, ready for alerting and monitoring
• Package distribution enables team-wide reuse — bundling your operators into an installable package means every team benefits from the same battle-tested integration logic

From here, there are several natural directions:

• Monitoring and alerting: Connect your XCom-captured run results to a monitoring layer. The article on monitoring dbt pipeline failures in production covers how Elementary integrates with this kind of structured output.

• Lineage tracking: The query tagging we implemented connects to Snowflake's QUERY_HISTORY view, which you can join against dbt's manifest.json to build automated lineage. This feeds into the broader pattern described in multi-hop data lineage tracking across the modern data stack.

• Incremental and backfill patterns: With these operators in place, building incremental load pipelines becomes much cleaner. The scheduling and backfilling historical data loads in Airflow article shows how to wire the {{ ds }} variable (which we're already threading through dbt_vars) into proper catchup behavior.

• Testing your dbt models: The data quality layer sits directly on top of the infrastructure you've built here. See implementing dbt tests and data quality checks in production pipelines to complete the reliability picture.

The patterns in this lesson — separation of concerns between Hooks and Operators, artifact capture, environment injection, structured testing — apply well beyond Snowflake and dbt. Once you've built these integrations once, you'll find yourself reaching for the same patterns every time you need to connect Airflow to a new system.