Orchestrating a Regulatory Reporting Pipeline with Apache Airflow

"""
dags/corep_pipeline_dag.py — Monthly COREP regulatory reporting pipeline.

Schedule: last day of each month at 01:00 UTC.
Executor:  LocalExecutor (configured in docker-compose.yml).

Design principles:
  1. Every module is a PythonOperator — no BashOperators, no subprocess DAG calls.
  2. Quality failures branch to quarantine — they never silently proceed.
  3. XCom carries only status strings and paths — never DataFrames.
  4. The final audit task runs regardless of branch taken.
  5. SLAs are set on the two slowest tasks (xbrl_gen, xbrl_valid).
"""

from __future__ import annotations

import json
import logging
import os
import subprocess
from datetime import datetime, timedelta, timezone
from pathlib import Path

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from airflow.utils.trigger_rule import TriggerRule
from airflow.utils.dates import days_ago

log = logging.getLogger(__name__)

# ── Project root (mounted into the Airflow container) ────────────────
PROJECT_ROOT = Path(os.environ.get("COREP_PROJECT_ROOT", "/opt/corep"))
DBT_PROJECT  = PROJECT_ROOT / "dbt"

# ── Default task arguments ───────────────────────────────────────────
DEFAULT_ARGS = {
    "owner":            "corep-pipeline",
    "depends_on_past":  False,
    "email_on_failure": True,
    "email_on_retry":   False,
    "email":            [os.environ.get("ALERT_EMAIL", "governance@example.com")],
    "retries":          2,
    "retry_delay":      timedelta(minutes=5),
    "retry_exponential_backoff": True,
}

# ── SLA callback — fires when a task exceeds its SLA ─────────────────
def sla_miss_callback(dag, task_list, blocking_task_list, slas, blocking_tis):
    log.error(
        "[airflow] SLA MISS on DAG=%s tasks=%s — governance team alerted.",
        dag.dag_id,
        [t.task_id for t in task_list],
    )
    # In production: send to PagerDuty / Slack webhook
    _slack_alert(
        f":warning: SLA breach on {dag.dag_id}: "
        + ", ".join(t.task_id for t in task_list)
    )


def _slack_alert(message: str) -> None:
    webhook = os.environ.get("SLACK_WEBHOOK_URL")
    if not webhook:
        log.info("[airflow] No SLACK_WEBHOOK_URL set — alert not sent.")
        return
    try:
        import urllib.request
        payload = json.dumps({"text": message}).encode()
        req = urllib.request.Request(
            webhook, data=payload, headers={"Content-Type": "application/json"}
        )
        urllib.request.urlopen(req, timeout=10)
    except Exception as exc:
        log.warning("[airflow] Slack alert failed (non-fatal): %s", exc)


# ═══════════════════════════════════════════════════════════════════════
# TASK FUNCTIONS
# Each function is a PythonOperator callable.
# Signature must accept **context to receive the Airflow task context dict.
# ═══════════════════════════════════════════════════════════════════════

def task_check_reporting_date(**context) -> None:
    """Derive and validate the reporting date from the DAG run logical_date."""
    logical_date  = context["logical_date"]           # pendulum.DateTime
    reporting_date = logical_date.strftime("%Y-%m-%d")

    # Determine the period: Q1=03-31, Q2=06-30, Q3=09-30, Q4=12-31
    month = logical_date.month
    if month in (3, 6, 9, 12):
        period_type = "QUARTERLY"
    else:
        period_type = "MONTHLY"

    log.info("[check_reporting_date] date=%s period=%s", reporting_date, period_type)
    ti = context["task_instance"]
    ti.xcom_push(key="reporting_date",   value=reporting_date)
    ti.xcom_push(key="reporting_period", value=period_type)


def task_run_ingest(**context) -> None:
    from modules.ingest import IngestModule
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    log.info("[ingest] Running for reporting_date=%s", reporting_date)
    IngestModule(run_context={"reporting_date": reporting_date}).run()


def task_run_quality_layer1(**context) -> None:
    """Run raw-layer quality gates. Push PASS/FAIL status to XCom."""
    from modules.quality import QualityModule, QualityGateError
    ti = context["task_instance"]
    try:
        QualityModule(layer="raw").run()
        ti.xcom_push(key="quality_gate_status", value="PASS")
    except QualityGateError as exc:
        log.error("[quality_layer1] Gate FAILED: %s", exc)
        ti.xcom_push(key="quality_gate_status", value="FAIL")


def branch_quality_layer1(**context) -> str:
    status = context["task_instance"].xcom_pull(
        task_ids="run_quality_layer1", key="quality_gate_status"
    )
    return "run_dbt" if status == "PASS" else "quarantine_raw_failure"


def task_quarantine_raw_failure(**context) -> None:
    """Archive raw tables snapshot to MinIO quarantine on Layer 1 failure."""
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    log.error("[quarantine] Layer 1 quality FAILED for %s — routing to quarantine.", reporting_date)
    _slack_alert(
        f":x: COREP Layer 1 quality gate FAILED for {reporting_date}. "
        "Pipeline halted — check GX data docs in MinIO corep-gx-reports."
    )
    # Write quarantine marker to audit log
    _write_audit_marker(context, "quarantine_raw", "FAIL")


def task_run_dbt(**context) -> None:
    """Run dbt staging → intermediate → mart + dbt test."""
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    env = {
        **os.environ,
        "REPORTING_DATE": reporting_date,
    }
    for cmd in [
        ["dbt", "run",  "--profiles-dir", str(DBT_PROJECT), "--project-dir", str(DBT_PROJECT)],
        ["dbt", "test", "--profiles-dir", str(DBT_PROJECT), "--project-dir", str(DBT_PROJECT)],
    ]:
        result = subprocess.run(cmd, capture_output=True, text=True, env=env, check=False)
        log.info("[dbt] stdout:\n%s", result.stdout[-3000:])
        if result.returncode != 0:
            raise RuntimeError(
                f"[dbt] Command failed: {' '.join(cmd)}\n" + result.stderr[-2000:]
            )
    log.info("[dbt] All models and tests passed.")


def task_run_quality_layer2(**context) -> None:
    from modules.quality import QualityModule, QualityGateError
    ti = context["task_instance"]
    try:
        QualityModule(layer="mart").run()
        ti.xcom_push(key="mart_gate_status", value="PASS")
    except QualityGateError as exc:
        log.error("[quality_layer2] Gate FAILED: %s", exc)
        ti.xcom_push(key="mart_gate_status", value="FAIL")


def branch_quality_layer2(**context) -> str:
    status = context["task_instance"].xcom_pull(
        task_ids="run_quality_layer2", key="mart_gate_status"
    )
    return "run_catalog" if status == "PASS" else "quarantine_mart_failure"


def task_quarantine_mart_failure(**context) -> None:
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    log.error("[quarantine] Layer 2 mart quality FAILED for %s.", reporting_date)
    _slack_alert(
        f":x: COREP mart quality gate FAILED for {reporting_date}. "
        "dbt mart models produced values that failed GX expectations. "
        "Check mart_corep_suite results in MinIO corep-gx-reports."
    )
    _write_audit_marker(context, "quarantine_mart", "FAIL")


def task_run_catalog(**context) -> None:
    from modules.catalog import CatalogModule
    CatalogModule().run()


def task_run_security(**context) -> None:
    from modules.security import SecurityModule
    SecurityModule().run()


def task_refresh_superset(**context) -> None:
    """Refresh Superset schema cache so new mart columns are visible."""
    superset_url  = os.environ.get("SUPERSET_URL", "http://superset:8088")
    admin_user    = os.environ.get("SUPERSET_ADMIN_USER", "admin")
    admin_pass    = os.environ.get("SUPERSET_ADMIN_PASS", "admin")
    try:
        import urllib.request, urllib.parse
        # Login to get CSRF token
        login_data = json.dumps({
            "username": admin_user, "password": admin_pass,
            "provider": "db", "refresh": True,
        }).encode()
        req  = urllib.request.Request(
            f"{superset_url}/api/v1/security/login",
            data=login_data,
            headers={"Content-Type": "application/json"},
        )
        resp = json.loads(urllib.request.urlopen(req, timeout=15).read())
        token = resp["access_token"]
        # Refresh all database metadata
        refresh_req = urllib.request.Request(
            f"{superset_url}/api/v1/database/1/schemas/",
            headers={"Authorization": f"Bearer {token}"},
        )
        urllib.request.urlopen(refresh_req, timeout=15)
        log.info("[superset] Schema cache refreshed.")
    except Exception as exc:
        log.warning("[superset] Refresh failed (non-fatal): %s", exc)


def task_run_xbrl_gen(**context) -> None:
    from modules.xbrl_gen import XbrlGenModule
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    XbrlGenModule(run_context={"reporting_date": reporting_date}).run()


def task_run_xbrl_valid(**context) -> None:
    from modules.xbrl_valid import XbrlValidModule, XbrlValidationError
    ti = context["task_instance"]
    try:
        XbrlValidModule().run()
        ti.xcom_push(key="xbrl_valid_status", value="PASS")
    except XbrlValidationError as exc:
        log.error("[xbrl_valid] Validation FAILED: %s", exc)
        ti.xcom_push(key="xbrl_valid_status", value="FAIL")


def branch_xbrl_valid(**context) -> str:
    status = context["task_instance"].xcom_pull(
        task_ids="run_xbrl_valid", key="xbrl_valid_status"
    )
    return "build_submission_package" if status == "PASS" else "quarantine_xbrl_failure"


def task_quarantine_xbrl_failure(**context) -> None:
    reporting_date = context["task_instance"].xcom_pull(
        task_ids="check_reporting_date", key="reporting_date"
    )
    log.error("[quarantine] XBRL validation FAILED for %s.", reporting_date)
    _slack_alert(
        f":x: COREP XBRL validation FAILED for {reporting_date}. "
        "EBA formula rules not satisfied. Check validation report in MinIO corep-xbrl-output."
    )
    _write_audit_marker(context, "quarantine_xbrl", "FAIL")


def task_build_submission_package(**context) -> None:
    """
    Bundle the XBRL instance + all audit artefacts into a submission package ZIP.
    Uploaded to MinIO corep-xbrl-output/submissions/.
    """
    import zipfile, shutil, tempfile
    from minio import Minio

    ti             = context["task_instance"]
    reporting_date = ti.xcom_pull(task_ids="check_reporting_date", key="reporting_date")
    ts             = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    pkg_name       = f"COREP_{reporting_date}_{ts}.zip"

    xbrl_dir   = Path(os.environ.get("XBRL_OUTPUT_DIR",   "output/xbrl"))
    report_dir = Path(os.environ.get("XBRL_REPORT_DIR",  "output/validation_reports"))
    gx_dir     = Path("gx/uncommitted/data_docs/local_site")

    with tempfile.TemporaryDirectory() as tmp:
        pkg_path = Path(tmp) / pkg_name
        with zipfile.ZipFile(pkg_path, "w", zipfile.ZIP_DEFLATED) as zf:
            # XBRL instance
            for f in sorted(xbrl_dir.glob("*.xbrl")):
                zf.write(f, f"xbrl/{f.name}")
            # Validation report
            for f in sorted(report_dir.glob("*.json")):
                zf.write(f, f"validation/{f.name}")
            # GX data docs index
            index_html = gx_dir / "index.html"
            if index_html.exists():
                zf.write(index_html, "quality/gx_data_docs_index.html")
            # Pipeline manifest (run metadata)
            manifest = {
                "reporting_date": reporting_date,
                "packaged_at":    ts,
                "dag_run_id":     context["run_id"],
                "airflow_version": "2.9",
                "xbrl_files":     [f.name for f in sorted(xbrl_dir.glob("*.xbrl"))],
            }
            zf.writestr("manifest.json", json.dumps(manifest, indent=2))

        # Upload to MinIO
        client = Minio(
            os.environ.get("MINIO_ENDPOINT", "minio:9000"),
            access_key=os.environ.get("MINIO_ROOT_USER",     "minioadmin"),
            secret_key=os.environ.get("MINIO_ROOT_PASSWORD", "minioadmin"),
            secure=False,
        )
        bucket = "corep-xbrl-output"
        if not client.bucket_exists(bucket):
            client.make_bucket(bucket)
        client.fput_object(bucket, f"submissions/{pkg_name}", str(pkg_path))
        log.info("[submission] Package uploaded → minio://%s/submissions/%s", bucket, pkg_name)
        ti.xcom_push(key="submission_package", value=pkg_name)


def task_write_final_audit_log(**context) -> None:
    """Write the complete pipeline run summary to audit.pipeline_run_log."""
    from modules.base import _pg_conn

    ti             = context["task_instance"]
    reporting_date = ti.xcom_pull(task_ids="check_reporting_date", key="reporting_date")
    q1_status      = ti.xcom_pull(task_ids="run_quality_layer1",    key="quality_gate_status") or "SKIP"
    q2_status      = ti.xcom_pull(task_ids="run_quality_layer2",    key="mart_gate_status")    or "SKIP"
    xv_status      = ti.xcom_pull(task_ids="run_xbrl_valid",        key="xbrl_valid_status")   or "SKIP"
    pkg_name       = ti.xcom_pull(task_ids="build_submission_package", key="submission_package") or "N/A"

    overall = "PASS" if xv_status == "PASS" else "FAIL"

    conn = _pg_conn()
    try:
        cur = conn.cursor()
        cur.execute(
            """
            INSERT INTO audit.pipeline_run_log
                (run_id, module_name, status, metadata, ran_at)
            VALUES (%s, 'pipeline_complete', %s, %s, now())
            """,
            (
                context["run_id"],
                overall,
                json.dumps({
                    "reporting_date":    reporting_date,
                    "dag_run_id":        context["run_id"],
                    "quality_layer1":    q1_status,
                    "quality_layer2":    q2_status,
                    "xbrl_validation":   xv_status,
                    "submission_package": pkg_name,
                }),
            ),
        )
        conn.commit()
        log.info("[audit] Final run summary written — overall=%s run_id=%s", overall, context["run_id"])
    finally:
        conn.close()

    if overall == "PASS":
        _slack_alert(
            f":white_check_mark: COREP pipeline PASSED for {reporting_date}. "
            f"Submission package: {pkg_name}"
        )


def _write_audit_marker(context: dict, stage: str, status: str) -> None:
    from modules.base import _pg_conn
    conn = _pg_conn()
    try:
        cur = conn.cursor()
        cur.execute(
            "INSERT INTO audit.pipeline_run_log (run_id, module_name, status, ran_at) VALUES (%s, %s, %s, now())",
            (context["run_id"], stage, status),
        )
        conn.commit()
    finally:
        conn.close()


# ═══════════════════════════════════════════════════════════════════════
# DAG DEFINITION
# ═══════════════════════════════════════════════════════════════════════

with DAG(
    dag_id="corep_monthly_pipeline",
    description="Monthly COREP regulatory reporting: ingest → quality → dbt → catalog → security → XBRL → validate → package",
    default_args=DEFAULT_ARGS,
    schedule_interval="0 1 L * *",   # 01:00 UTC on the last day of every month
    start_date=days_ago(1),
    catchup=False,
    max_active_runs=1,               # Only ONE pipeline run at a time
    tags=["corep", "regulatory", "eba", "governance"],
    sla_miss_callback=sla_miss_callback,
    doc_md="""
    ## COREP Monthly Regulatory Reporting Pipeline

    Runs on the last calendar day of each month at 01:00 UTC.
    Produces a validated XBRL instance for EBA COREP supervisory reporting.

    ### Modules
    | Step | Module | Output |
    |------|--------|--------|
    | 1 | ingest | raw.* tables |
    | 2 | quality_layer1 | GX raw checkpoints |
    | 3 | dbt | staging/intermediate/mart |
    | 4 | quality_layer2 | GX mart checkpoints |
    | 5 | catalog | OpenMetadata indexed |
    | 6 | security | Ranger policies applied |
    | 7 | xbrl_gen | .xbrl instance file |
    | 8 | xbrl_valid | Arelle validation report |
    | 9 | submission | ZIP package in MinIO |
    """,
) as dag:

    t_check_date       = PythonOperator(task_id="check_reporting_date",  python_callable=task_check_reporting_date)
    t_ingest           = PythonOperator(task_id="run_ingest",             python_callable=task_run_ingest)
    t_quality1         = PythonOperator(task_id="run_quality_layer1",     python_callable=task_run_quality_layer1)
    t_branch_q1        = BranchPythonOperator(task_id="branch_quality_layer1", python_callable=branch_quality_layer1)
    t_quarantine_raw   = PythonOperator(task_id="quarantine_raw_failure",  python_callable=task_quarantine_raw_failure)
    t_dbt              = PythonOperator(task_id="run_dbt",                python_callable=task_run_dbt)
    t_quality2         = PythonOperator(task_id="run_quality_layer2",     python_callable=task_run_quality_layer2)
    t_branch_q2        = BranchPythonOperator(task_id="branch_quality_layer2", python_callable=branch_quality_layer2)
    t_quarantine_mart  = PythonOperator(task_id="quarantine_mart_failure", python_callable=task_quarantine_mart_failure)
    t_catalog          = PythonOperator(task_id="run_catalog",            python_callable=task_run_catalog)
    t_security         = PythonOperator(task_id="run_security",           python_callable=task_run_security)
    t_superset         = PythonOperator(task_id="refresh_superset_metadata", python_callable=task_refresh_superset)
    t_xbrl_gen         = PythonOperator(
        task_id="run_xbrl_gen",
        python_callable=task_run_xbrl_gen,
        sla=timedelta(hours=1),    # alert if XBRL generation > 60 min
    )
    t_xbrl_valid       = PythonOperator(
        task_id="run_xbrl_valid",
        python_callable=task_run_xbrl_valid,
        sla=timedelta(minutes=30),  # alert if validation > 30 min
    )
    t_branch_xv        = BranchPythonOperator(task_id="branch_xbrl_valid",     python_callable=branch_xbrl_valid)
    t_quarantine_xbrl  = PythonOperator(task_id="quarantine_xbrl_failure", python_callable=task_quarantine_xbrl_failure)
    t_submission       = PythonOperator(task_id="build_submission_package", python_callable=task_build_submission_package)
    t_audit_final      = PythonOperator(
        task_id="write_final_audit_log",
        python_callable=task_write_final_audit_log,
        trigger_rule=TriggerRule.NONE_FAILED_MIN_ONE_SUCCESS,  # always runs
    )

    # ── Wire dependencies ────────────────────────────────────────────────
    t_check_date >> t_ingest >> t_quality1 >> t_branch_q1
    t_branch_q1  >> [t_dbt, t_quarantine_raw]
    t_dbt        >> t_quality2 >> t_branch_q2
    t_branch_q2  >> [t_catalog, t_quarantine_mart]
    t_catalog    >> [t_security, t_superset]
    [t_security, t_superset] >> t_xbrl_gen
    t_xbrl_gen   >> t_xbrl_valid >> t_branch_xv
    t_branch_xv  >> [t_submission, t_quarantine_xbrl]
    [t_submission, t_quarantine_raw, t_quarantine_mart, t_quarantine_xbrl] >> t_audit_final

# Access Airflow UI at http://localhost:8090
# Default credentials: admin / admin (set in docker-compose.yml)

# Verify the DAG appears — it should auto-detect from dags/ folder
# within 30 seconds of starting the Airflow container

# Check DAG syntax without triggering a run
python -c "from dags.corep_pipeline_dag import dag; print('DAG OK:', dag.dag_id)"

# Enable the DAG (it starts paused for safety)
# In the UI: toggle the pause switch next to "corep_monthly_pipeline"

# Or via Airflow CLI inside the container:
docker exec corep-airflow airflow dags unpause corep_monthly_pipeline

# Trigger via CLI with a specific logical date (simulates end-of-month run)
docker exec corep-airflow airflow dags trigger \
    corep_monthly_pipeline \
    --conf '{"reporting_date": "2026-03-31"}'

# Watch task states in real time
docker exec corep-airflow airflow tasks states-for-dag-run \
    corep_monthly_pipeline \
    --run-id "manual__2026-05-07T10:00:00+00:00"

# Expected output of a passing run:
check_reporting_date    success
run_ingest              success
run_quality_layer1      success
branch_quality_layer1   success
run_dbt                 success
run_quality_layer2      success
branch_quality_layer2   success
run_catalog             success
run_security            success
refresh_superset_metadata success
run_xbrl_gen            success
run_xbrl_valid          success
branch_xbrl_valid       success
build_submission_package success
write_final_audit_log   success
quarantine_raw_failure  skipped
quarantine_mart_failure skipped
quarantine_xbrl_failure skipped

# Branches not taken show as "skipped" — this is correct Airflow behaviour

# When a Layer 1 quality gate fails, the DAG run looks like this:
check_reporting_date    success
run_ingest              success
run_quality_layer1      success  # ← task succeeded; gate FAIL stored in XCom
branch_quality_layer1   success  # ← branch ran and chose quarantine path
quarantine_raw_failure  success  # ← quarantine task ran (status=success but pipeline=FAIL)
run_dbt                 skipped
run_quality_layer2      skipped
... all downstream      skipped
write_final_audit_log   success  # ← always runs due to NONE_FAILED_MIN_ONE_SUCCESS

# The entire DAG run shows as "success" in Airflow because
# no task raised an exception. The FAIL is in the audit log and XCom.
# This is intentional: a quality failure is not a pipeline ERROR,
# it is a DATA QUALITY INCIDENT handled by the designed workflow.

# Query the audit log to confirm:
docker exec corep-postgres psql -U corep_admin -d corep -c \
    "SELECT run_id, module_name, status, metadata->>'reporting_date' AS date
     FROM audit.pipeline_run_log
     ORDER BY ran_at DESC LIMIT 10;"

# openlineage-airflow config in airflow.cfg (or environment variables)
# Set these in docker-compose.yml for the Airflow container:

AIRFLOW__OPENLINEAGE__TRANSPORT=http
AIRFLOW__OPENLINEAGE__URL=http://marquez:5000
AIRFLOW__OPENLINEAGE__NAMESPACE=corep-governance-pipeline

# Verify lineage is being received after a run:
curl http://localhost:5000/api/v1/jobs?namespace=corep-governance-pipeline | python -m json.tool

# Expected: one job entry per task in the DAG
# Each job entry has runs[] with start_at, end_at, and state (COMPLETED/FAILED)

-- Query the full pipeline audit log for a reporting date
SELECT
    module_name,
    status,
    metadata->>'reporting_date'    AS date,
    metadata->>'error_count'       AS errors,
    metadata->>'submission_package' AS package,
    ran_at
FROM audit.pipeline_run_log
WHERE metadata->>'reporting_date' = '2026-03-31'
ORDER BY ran_at;

-- Output:
-- module_name          | status | date       | errors | package                         | ran_at
-- ingest               | PASS   | 2026-03-31 |        |                                 | 2026-03-31 01:02:14
-- quality_raw          | PASS   | 2026-03-31 |        |                                 | 2026-03-31 01:03:45
-- quality_mart         | PASS   | 2026-03-31 |        |                                 | 2026-03-31 01:09:22
-- catalog              | PASS   | 2026-03-31 |        |                                 | 2026-03-31 01:11:05
-- security             | PASS   | 2026-03-31 |        |                                 | 2026-03-31 01:11:38
-- xbrl_valid           | PASS   | 2026-03-31 | 0      |                                 | 2026-03-31 01:18:52
-- pipeline_complete    | PASS   | 2026-03-31 | 0      | COREP_2026-03-31_20260331T011912Z.zip | 2026-03-31 01:19:03