Data Pipelines

from datetime import datetime, timedelta
from airflow.decorators import dag, task
from airflow.operators.python import PythonOperator
from airflow.providers.postgres.hooks.postgres import PostgresHook

@dag(
    schedule="0 6 * * *",          # 6 AM daily
    start_date=datetime(2026, 1, 1),
    catchup=False,                  # don't backfill missed runs on deploy
    max_active_runs=1,              # prevent overlapping runs
    default_args={
        "retries": 3,
        "retry_delay": timedelta(minutes=5),
        "retry_exponential_backoff": True,
        "email_on_failure": True,
    },
    tags=["finance", "daily"],
)
def daily_revenue_pipeline():

    @task
    def extract_orders(execution_date=None) -> list[dict]:
        hook = PostgresHook(postgres_conn_id="source_db")
        # Use execution_date for idempotent extraction
        rows = hook.get_records(
            "SELECT * FROM orders WHERE date = %s",
            parameters=[execution_date.date()],
        )
        return [dict(r) for r in rows]

    @task
    def transform(orders: list[dict]) -> list[dict]:
        return [
            {**o, "revenue_usd": o["amount"] * o["fx_rate"]}
            for o in orders
            if o["status"] == "completed"
        ]

    @task
    def load(records: list[dict], execution_date=None):
        hook = PostgresHook(postgres_conn_id="warehouse")
        # Idempotent: delete-then-insert for the partition date
        hook.run("DELETE FROM daily_revenue WHERE date = %s", parameters=[execution_date.date()])
        hook.insert_rows("daily_revenue", [[r["date"], r["revenue_usd"]] for r in records])

    orders = extract_orders()
    transformed = transform(orders)
    load(transformed)

dag = daily_revenue_pipeline()

from airflow.operators.bash import BashOperator
from airflow.operators.python import BranchPythonOperator
from airflow.sensors.filesystem import FileSensor
from airflow.sensors.sql import SqlSensor
from airflow.providers.http.sensors.http import HttpSensor

# Wait for a file to appear (S3, GCS, local)
wait_for_export = FileSensor(
    task_id="wait_for_export",
    filepath="/data/exports/{{ ds }}/orders.csv",
    poke_interval=60,    # check every 60s
    timeout=3600,        # fail after 1 hour
    mode="reschedule",   # release worker slot while waiting
)

# Wait for upstream table to be populated
wait_for_source = SqlSensor(
    task_id="wait_for_orders",
    conn_id="source_db",
    sql="SELECT COUNT(*) FROM orders WHERE date = '{{ ds }}' HAVING COUNT(*) > 0",
    poke_interval=120,
    mode="reschedule",
)

# Branch: skip load on weekends
def should_load(**context):
    if context["execution_date"].weekday() >= 5:
        return "skip_load"
    return "load"

branch = BranchPythonOperator(task_id="check_day", python_callable=should_load)

# Push value
@task
def extract() -> dict:
    return {"row_count": 1042, "checksum": "abc123"}  # return value auto-pushes XCom

# Pull value
@task
def validate(stats: dict):   # passed as argument from task dependency
    assert stats["row_count"] > 0, "Empty extract"

# Manual XCom pull (classic operators)
def load(**context):
    stats = context["task_instance"].xcom_pull(task_ids="extract")
    print(stats["row_count"])

@task
def get_regions() -> list[str]:
    return ["us-east", "eu-west", "ap-south"]

@task
def process_region(region: str):
    extract_and_load(region)

# Creates one task instance per region — parallelized automatically
process_region.expand(region=get_regions())

from airflow.hooks.base import BaseHook
from airflow.models import Variable

# Never hardcode credentials — use Connections
conn = BaseHook.get_connection("my_postgres")
dsn = f"postgresql://{conn.login}:{conn.password}@{conn.host}/{conn.schema}"

# Runtime config — use Variables (or better: Airflow Params)
batch_size = int(Variable.get("etl_batch_size", default_var=1000))

dbt_project/
├── models/
│   ├── staging/          # stg_* — raw → typed, renamed, deduplicated
│   │   └── stg_orders.sql
│   ├── intermediate/     # int_* — business logic joins
│   │   └── int_order_items.sql
│   └── marts/            # final — wide tables for BI/downstream
│       └── fct_revenue.sql
├── tests/                # custom SQL tests
├── macros/               # Jinja macros
├── seeds/                # static CSV reference data
└── dbt_project.yml

-- staging/stg_orders.sql
-- Materialization: view (cheap, always fresh)
{{ config(materialized='view') }}

SELECT
    order_id::VARCHAR      AS order_id,
    user_id::VARCHAR       AS user_id,
    created_at::TIMESTAMP  AS created_at,
    amount_cents / 100.0   AS amount_usd,
    status
FROM {{ source('raw', 'orders') }}
WHERE status != 'test'

-- marts/fct_revenue.sql
-- Materialization: table (fast reads, rebuilt on each run)
{{ config(materialized='table') }}

SELECT
    DATE_TRUNC('day', o.created_at) AS date,
    p.name                          AS product_name,
    SUM(oi.quantity)                AS units_sold,
    SUM(oi.quantity * oi.unit_price) AS revenue_usd
FROM {{ ref('stg_orders') }}      o    -- ref() creates dependency
JOIN {{ ref('int_order_items') }} oi ON o.order_id = oi.order_id
JOIN {{ ref('stg_products') }}    p  ON oi.product_id = p.product_id
WHERE o.status = 'completed'
GROUP BY 1, 2

-- Only process new/updated rows — essential for large tables
{{ config(
    materialized='incremental',
    unique_key='order_id',
    incremental_strategy='merge',    -- or 'delete+insert', 'insert_overwrite'
    on_schema_change='append_new_columns',
) }}

SELECT
    order_id,
    user_id,
    amount_usd,
    created_at,
    updated_at
FROM {{ source('raw', 'orders') }}

{% if is_incremental() %}
    -- Only load rows newer than the last run
    WHERE updated_at > (SELECT MAX(updated_at) FROM {{ this }})
{% endif %}

# models/staging/sources.yml
version: 2

sources:
  - name: raw
    database: analytics
    schema: raw_data
    freshness:
      warn_after: {count: 6, period: hour}
      error_after: {count: 24, period: hour}
    loaded_at_field: _loaded_at       # column that holds ingestion timestamp
    tables:
      - name: orders
        description: Raw orders from the transactional database
      - name: products

# Check source freshness in CI
dbt source freshness

# models/staging/stg_orders.yml
version: 2

models:
  - name: stg_orders
    columns:
      - name: order_id
        tests:
          - not_null
          - unique
      - name: status
        tests:
          - accepted_values:
              values: ["pending", "completed", "cancelled", "refunded"]
      - name: user_id
        tests:
          - not_null
          - relationships:
              to: ref('stg_users')
              field: user_id
      - name: amount_usd
        tests:
          - not_null
          - dbt_utils.accepted_range:
              min_value: 0
              max_value: 100000

-- tests/assert_revenue_non_negative.sql — custom SQL test (fails if rows returned)
SELECT date, revenue_usd
FROM {{ ref('fct_revenue') }}
WHERE revenue_usd < 0

-- macros/cents_to_dollars.sql
{% macro cents_to_dollars(column_name) %}
    ({{ column_name }} / 100.0)::NUMERIC(10, 2)
{% endmacro %}

-- Usage in a model
SELECT {{ cents_to_dollars('amount_cents') }} AS amount_usd

-- macros/generate_surrogate_key.sql (or use dbt_utils)
{% macro surrogate_key(fields) %}
    MD5(CONCAT_WS('|', {% for f in fields %}COALESCE(CAST({{ f }} AS VARCHAR), ''){% if not loop.last %}, {% endif %}{% endfor %}))
{% endmacro %}

dbt run                              # run all models
dbt run --select staging             # run a directory
dbt run --select stg_orders+         # run model and all downstream
dbt run --select +fct_revenue        # run model and all upstream
dbt test                             # run all tests
dbt test --select stg_orders         # test one model
dbt build                            # run + test in dependency order
dbt source freshness                 # check source data freshness
dbt docs generate && dbt docs serve  # generate + serve lineage docs
dbt compile                          # render SQL without running

# GOOD: delete-then-insert for a known partition
def load_partition(date: str, records: list[dict]):
    with engine.begin() as conn:
        conn.execute(
            text("DELETE FROM daily_stats WHERE date = :date"),
            {"date": date}
        )
        conn.execute(insert(DailyStats), records)

# GOOD: UPSERT (merge) on unique key
def upsert_orders(records: list[dict]):
    stmt = pg_insert(orders_table).values(records)
    stmt = stmt.on_conflict_do_update(
        index_elements=["order_id"],
        set_={"status": stmt.excluded.status, "updated_at": stmt.excluded.updated_at}
    )
    with engine.begin() as conn:
        conn.execute(stmt)

# BAD: append-only — reruns duplicate data
def load_orders(records):
    engine.execute(insert(orders_table).values(records))  # duplicates on rerun

# Watermark-based incremental (Python)
def get_watermark(conn, table: str) -> datetime:
    row = conn.execute(
        text("SELECT COALESCE(MAX(updated_at), '1970-01-01') FROM :table", bindparams=[bindparam("table")])
    ).fetchone()
    return row[0]

def extract_incremental(source_conn, watermark: datetime) -> list[dict]:
    return source_conn.execute(
        text("SELECT * FROM orders WHERE updated_at > :wm ORDER BY updated_at"),
        {"wm": watermark},
    ).fetchall()

# Generic tests: not_null, unique, accepted_values, relationships
# Package tests: dbt_utils, dbt_expectations (Great Expectations style)
- name: amount_usd
  tests:
    - dbt_expectations.expect_column_values_to_be_between:
        min_value: 0
        max_value: 50000
        row_condition: "status = 'completed'"

import great_expectations as gx

context = gx.get_context()
suite = context.add_expectation_suite("orders_suite")

validator = context.get_validator(
    batch_request=batch_request,
    expectation_suite_name="orders_suite",
)
validator.expect_column_values_to_not_be_null("order_id")
validator.expect_column_values_to_be_unique("order_id")
validator.expect_column_values_to_be_between("amount_usd", min_value=0)
validator.expect_column_pair_values_A_to_be_greater_than_B(
    "completed_at", "created_at"
)

results = validator.validate()
if not results.success:
    raise ValueError(f"Data quality check failed: {results}")

@task
def reconcile(source_count: int, target_count: int, tolerance: float = 0.001):
    delta = abs(source_count - target_count) / max(source_count, 1)
    if delta > tolerance:
        raise ValueError(
            f"Row count mismatch: source={source_count}, target={target_count}, "
            f"delta={delta:.2%} > {tolerance:.2%} tolerance"
        )

# Airflow: SLA miss callback
def sla_miss_callback(dag, task_list, blocking_task_list, slas, blocking_tis):
    send_slack_alert(f"SLA missed for DAG {dag.dag_id}: {task_list}")

@dag(sla_miss_callback=sla_miss_callback)
def my_dag():
    ...

# Airflow: task-level SLA (fail if task exceeds duration)
load = PythonOperator(
    task_id="load",
    python_callable=load_fn,
    sla=timedelta(minutes=30),   # alert if this task takes >30 min
)

# Emit pipeline metrics to Prometheus/StatsD
from airflow.stats import Stats

Stats.incr("pipeline.rows_processed", count=row_count, tags={"dag": dag_id})
Stats.timing("pipeline.duration_ms", value=duration_ms, tags={"dag": dag_id})