Source code for util.dag_trigger

"""DAG Trigger Scheduler completion callbacks.

Provides on_success_callback / on_failure_callback functions that target DAGs
use to report completion status back to airflow.dag_trigger_run in Redshift.

Usage in a target DAG:

    from util.dag_trigger import scheduler_callbacks

    @Util.default_dag(
        Config, file_name=__file__, email=EMAIL,
        on_success_callback=scheduler_callbacks.on_success,
        on_failure_callback=scheduler_callbacks.on_failure,
    )
    def my_dag():
        ...

Runs triggered manually (without _scheduler_run_id in conf) are ignored.
"""

from config import Config
from operators.redshift.base_redshift import BaseRedshiftOperator

_REDSHIFT_DB = "dataops"
_REGION = "us-east-1"
_SCHEDULER_RUN_ID_KEY = "_scheduler_run_id"


def _mark_completed(airflow_run_id: str, status: str, error_message: str | None = None) -> None:
    # Accept pending OR triggered (never overwrite a terminal row). This closes
    # the race where a fast-finishing target DAG fires its callback before the
    # scheduler's mark_triggered runs (row still 'pending'), which previously
    # matched zero rows and silently lost the completion.
    error_val = BaseRedshiftOperator.sql_literal(error_message[:65535]) if error_message else "NULL"
    # COALESCE triggered_at: a fast target DAG can fire this callback before the
    # scheduler's mark_triggered runs (row still 'pending', triggered_at NULL).
    # Without this, the run is excluded from find_due_items' last_triggered_at
    # watermark (which requires triggered_at IS NOT NULL), causing the next
    # incremental run to reprocess a wider window / fall back to full refresh.
    # created_at is conservative (<= source-read time), so the cutoff can't skip rows.
    sql = (
        "UPDATE airflow.dag_trigger_run\n"
        f"SET status = {BaseRedshiftOperator.sql_literal(status)},\n"
        f"    triggered_at = COALESCE(triggered_at, created_at),\n"
        f"    finished_at = GETDATE(),\n"
        f"    error_message = {error_val}\n"
        f"WHERE airflow_run_id = {BaseRedshiftOperator.sql_literal(airflow_run_id)}\n"
        f"    AND status IN ('pending', 'triggered');"
    )
    affected = BaseRedshiftOperator.run_statement(
        sql,
        workgroup=Config.REDSHIFT_WORKGROUP,
        database=_REDSHIFT_DB,
        region=_REGION,
    )
    print(f"[dag_trigger] Marked {affected} row(s) '{status}' for {airflow_run_id}")


def _on_success(context) -> None:
    conf = context["dag_run"].conf or {}
    run_id = conf.get(_SCHEDULER_RUN_ID_KEY)
    if not run_id:
        return
    try:
        _mark_completed(run_id, "success")
    except Exception as e:
        print(f"[dag_trigger] Failed to mark success for {run_id}: {e}")


def _on_failure(context) -> None:
    conf = context["dag_run"].conf or {}
    run_id = conf.get(_SCHEDULER_RUN_ID_KEY)
    if not run_id:
        return
    error_msg = str(context.get("exception", ""))[:65535] or None
    try:
        _mark_completed(run_id, "failed", error_msg)
    except Exception as e:
        print(f"[dag_trigger] Failed to mark failure for {run_id}: {e}")




class scheduler_callbacks:
    on_success = staticmethod(_on_success)
    on_failure = staticmethod(_on_failure)