import sys, os
from urllib.parse import quote
from datetime import datetime, timedelta, timezone
from zoneinfo import ZoneInfo
from airflow import DAG
from airflow.decorators import dag
from airflow.models.dagrun import DagRun as AirflowDagRun
from airflow.models.log import Log
from airflow.operators.email import EmailOperator
from airflow.operators.python import PythonOperator
try:
from airflow.providers.fab.auth_manager.models import User as FabUser
except ImportError:
FabUser = None
from airflow.utils.session import create_session
from config import Config
from operators.communications.notify_slack import NotifySlackOperator
_SSO_USERS: dict = {}
try:
with open(
os.path.join(os.path.dirname(os.path.dirname(__file__)), ".agents", "airflow-sso-users.json")
) as _f:
import json as _json
_SSO_USERS = _json.load(_f)
except (FileNotFoundError, ValueError):
pass
def _get_slack_id(email: str | None) -> str | None:
if not email:
return None
entry = _SSO_USERS.get(email)
return entry.get("slack_user_id") if entry else None
_TRIGGER_EVENTS = frozenset({"trigger", "api.post_dag_run"})
_CLEAR_EVENTS = frozenset({"dagrun_clear"})
_ET = ZoneInfo("America/New_York")
# Per-DAG Slack opt-out registry. Populated at DAG parse time by _dag(); read at
# callback time by notify_failure. Avoids putting a non-BaseOperator key in
# default_args, which Airflow propagates to every task and warns about.
_DAG_NOTIFY_SLACK: dict[str, bool] = {}
# Sentinel for "caller did not pass this kwarg". Needed for `interval` because
# None is a meaningful value (an unscheduled, manual-trigger-only DAG) and so
# cannot double as "unset".
_UNSET = object()
def _now_str() -> str:
now_utc = datetime.now(timezone.utc)
now_et = now_utc.astimezone(_ET)
et_label = (
now_et.strftime("%H:%M ET")
if now_et.date() == now_utc.date()
else now_et.strftime("%Y-%m-%d %H:%M ET")
)
return now_utc.strftime(f"%Y-%m-%d %H:%M UTC ({et_label})")
def _format_duration(start_date) -> str:
elapsed = datetime.now(timezone.utc) - start_date
total_secs = int(elapsed.total_seconds())
hrs, remainder = divmod(total_secs, 3600)
mins, secs = divmod(remainder, 60)
if hrs:
return f"{hrs}h {mins}m {secs}s"
if mins:
return f"{mins}m {secs}s"
return f"{secs}s"
def _append_dag_run_note(dag_id: str, run_id: str, note: str):
"""Prepend a timestamped note to the DagRun (newest first) — visible in Airflow UI Details tab."""
with create_session() as session:
dag_run_obj = (
session.query(AirflowDagRun)
.filter(AirflowDagRun.dag_id == dag_id, AirflowDagRun.run_id == run_id)
.first()
)
if dag_run_obj:
if dag_run_obj.dag_run_note:
existing = dag_run_obj.dag_run_note.content or ""
dag_run_obj.dag_run_note.content = (
f"{note}\n\n---\n\n{existing}".strip()
)
else:
dag_run_obj.note = note
def _build_user_dict(session, entry, action: str) -> dict:
email = None
if FabUser is not None:
fab_row = (
session.query(FabUser.email).filter(FabUser.username == entry.owner).first()
)
email = fab_row.email if fab_row else None
return {
"username": entry.owner,
"display_name": entry.owner_display_name or entry.owner,
"email": email,
"action": action,
}
def _get_trigger_user(dag_id: str, run_id: str) -> dict | None:
"""
Returns {"username", "display_name", "email", "action"} for whoever most
recently initiated this run (initial trigger OR clear/rerun), or None for
scheduled runs.
action == "initially triggered" → first trigger (UI or API)
action == "rerun" → someone used "Clear Existing Tasks"
"""
trigger_time = None
if run_id.startswith("manual__"):
try:
trigger_time = datetime.fromisoformat(run_id[len("manual__") :])
except ValueError:
pass
with create_session() as session:
clear_entry = (
session.query(Log.owner, Log.owner_display_name)
.filter(
Log.dag_id == dag_id,
Log.run_id == run_id,
Log.event.in_(_CLEAR_EVENTS),
)
.order_by(Log.dttm.desc())
.first()
)
if clear_entry:
return _build_user_dict(session, clear_entry, "rerun")
q = session.query(Log.owner, Log.owner_display_name).filter(
Log.dag_id == dag_id,
Log.event.in_(_TRIGGER_EVENTS),
)
if trigger_time:
q = q.filter(
Log.dttm >= trigger_time - timedelta(seconds=30),
Log.dttm <= trigger_time + timedelta(seconds=90),
)
else:
q = q.filter(Log.run_id == run_id)
entry = q.order_by(Log.dttm.desc()).first()
if not entry:
return None
return _build_user_dict(session, entry, "initially triggered")
def notify_failure(context):
dag_id = context["task_instance"].dag_id
run_id = context["dag_run"].run_id
duration = _format_duration(context["dag_run"].start_date)
# Best-effort trigger detection — fallback to None so email always sends below
trigger = None
try:
trigger = _get_trigger_user(dag_id, run_id)
except Exception:
pass
action = trigger["action"] if trigger else "triggered"
actor = trigger["display_name"] if trigger else "scheduler"
# Email: runs unconditionally; CC trigger user if different from DAG owner
SERVER = "Prod" if Config.PRODUCTION else "Dev"
trigger_label = "↺ Rerun" if action == "rerun" else "▶ Initially triggered"
CONTENT = (
f'<strong>DAG:</strong> <a href="{Config.AIRFLOW_URL}/dags/{dag_id}/grid?dag_run_id={quote(run_id, safe="")}">{dag_id}</a>'
f'<br><br><strong>TASK:</strong> {context["task_instance"].task_id}'
f'<br><br><strong>STATUS:</strong> ❌ FAILED'
f'<br><br><strong>TRIGGERED BY:</strong> {trigger_label} by {actor}'
f'<br><br><strong>DURATION:</strong> {duration}'
f'<br><br><strong>TIME:</strong> {_now_str()}'
f'<br><br><strong>OWNER:</strong> {context["dag"].owner}'
f'<br><br><strong>EXCEPTION:</strong><br>{context["exception"]}<br>'
)
owner_emails = context["dag"].default_args["email"]
owner_set = set(owner_emails) if isinstance(owner_emails, list) else {owner_emails}
cc_emails = (
[trigger["email"]]
if trigger and trigger["email"] and trigger["email"] not in owner_set
else []
)
EmailOperator(
task_id="failure_notification",
conn_id=Config.SENDGRID_CONN,
to=owner_emails,
cc=cc_emails,
subject=f"[{SERVER}] ❌ FAILED | {dag_id} | {trigger_label} by {actor}",
html_content=CONTENT,
).execute(context)
# DAG Run note + Slack: best-effort, must not suppress the email above
# Slack can be disabled per-DAG via notify_slack=False in default_dag(...) / _dag(...)
_notify_slack = _DAG_NOTIFY_SLACK.get(dag_id, True)
try:
_append_dag_run_note(
dag_id,
run_id,
(
f"### ❌ FAILED · {'↺ rerun' if action == 'rerun' else '▶ initially triggered'}"
f" by: {actor}\n"
f"`{context['task_instance'].task_id}` · {_now_str()} · duration: {duration}"
),
)
if _notify_slack:
_slack_id = _get_slack_id(trigger["email"] if trigger else None)
_mention = f"<@{_slack_id}> " if _slack_id else ""
NotifySlackOperator(
task_id="failure_slack_notification",
channel=Config.SLACK_ALERTS_CHANNEL,
text=(
f":red_circle: *FAILED* | <{Config.AIRFLOW_URL}/dags/{dag_id}/grid?dag_run_id={quote(run_id, safe='')}|{dag_id}> · `{context['task_instance'].task_id}`\n"
f"{'↺ rerun' if action == 'rerun' else '▶ initially triggered'} by: {_mention}*{actor}* | {_now_str()} | duration: {duration}"
),
).execute(context)
except Exception:
pass
[docs]
class DagUtil:
"""
Central dag utilities that apply to initialization, behavior of dags
"""
[docs]
@staticmethod
def get_dag_id(constant_object, path: str = "") -> str:
"""
Get dag id from constant object
:param constant_object: the constant object
:param path: the path to the file
:return: the dag id
"""
try:
dag_id = constant_object.DAG_ID
except AttributeError:
file_path = (
os.path.abspath(sys.modules[constant_object.__module__].__file__)
if not path
else path
)
dag_id = (
file_path.split(constant_object.SPLIT_PATH, 1)[-1]
.rsplit("/", 1)[0]
.replace("/", "_")
)
if dag_id.startswith(constant_object.GENERAL_PREFIX):
dag_id = dag_id[len(constant_object.GENERAL_PREFIX) :]
return dag_id
@staticmethod
def _dag(
constant_object,
params: dict = {},
decorator: bool = True,
file_name: str = "",
email: str = "",
template: str = "",
tags: list[str] = [],
notify_slack: bool = True,
on_success_callback=None,
on_failure_callback=None,
dag_id: str = "",
interval=_UNSET,
start_date=None,
):
"""
Cut back on repeated code for initializing dags by providing sane defaults and easily allowing override.
Uses central params dict to pass params to all tasks in dag, can return decorator or dag object
:param constant_object: the constant object
:param params: the params to pass to the dag
:param decorator: whether or not to return a decorator
:param file_name: the file name
:param email: the email to send to
:param template: the template to use
:param tags: the tags to use
:param notify_slack: if False, suppress Slack alerts from notify_failure (default True)
:param on_success_callback: DAG-level callback on successful run completion (additive — does not replace task-level notify_failure)
:param on_failure_callback: DAG-level callback on failed run completion (additive — does not replace task-level notify_failure)
:param dag_id: explicit dag_id. Defaults to the id auto-built from file_name.
Required when one file emits several DAGs (per-state DAG factories),
since every DAG in the file would otherwise get the same derived id.
:param interval: explicit schedule. Defaults to constant_object.INTERVAL.
Pass None for an unscheduled (manual-trigger-only) DAG.
:param start_date: explicit start_date. Defaults to constant_object.START_DATE.
:return: the dag
"""
tgt_email = email or constant_object.EMAIL
dag_id = dag_id or DagUtil.get_dag_id(constant_object, file_name)
_DAG_NOTIFY_SLACK[dag_id] = notify_slack
dag_args = {
"default_args": {
**constant_object.DEFAULT_ARGS,
**{
"email": [tgt_email],
"on_failure_callback": notify_failure,
# "on_success_callback": notify_task_success, # this works but it's way too gd annoying
"owner": constant_object.DEFAULT_ARGS.get(
"owner", tgt_email.split("@")[0]
),
"start_date": start_date
or constant_object.DEFAULT_ARGS.get(
"start_date", constant_object.START_DATE
),
},
},
"render_template_as_native_obj": constant_object.RENDER_TEMPLATE_AS_NATIVE_OBJ,
"dag_id": dag_id,
"catchup": constant_object.CATCHUP,
"schedule_interval": (
constant_object.INTERVAL if interval is _UNSET else interval
),
"params": params,
"concurrency": constant_object.MAX_ACTIVE_TASKS,
"max_active_runs": constant_object.MAX_ACTIVE_RUNS,
"template_searchpath": (
f"{constant_object.DAG_ROOT}/{template}"
if template
else constant_object.TEMPLATE_SEARCHPATH
),
"tags": tags,
}
if on_success_callback:
dag_args["on_success_callback"] = on_success_callback
if on_failure_callback:
dag_args["on_failure_callback"] = on_failure_callback
return dag(**dag_args) if decorator else DAG(**dag_args)
[docs]
@staticmethod
def default_dag(
constant_object,
params: dict = {},
file_name: str = "",
email: str = "",
template: str = "",
tags: list[str] = [],
notify_slack: bool = True,
on_success_callback=None,
on_failure_callback=None,
dag_id: str = "",
interval=_UNSET,
start_date=None,
):
"""
This is an interface over _dag that returns a decorator
:param constant_object: the constant object
:param params: the params to pass to the dag
:param file_name: the file name
:param email: the email to send to
:param template: the template to use
:param tags: the tags to use
:param notify_slack: if False, suppress Slack alerts from notify_failure (default True)
:param on_success_callback: DAG-level callback on successful run completion
:param on_failure_callback: DAG-level callback on failed run completion
:param dag_id: explicit dag_id (defaults to the id auto-built from file_name)
:param interval: explicit schedule (defaults to constant_object.INTERVAL;
pass None for an unscheduled, manual-trigger-only DAG)
:param start_date: explicit start_date (defaults to constant_object.START_DATE)
:return: the dag
"""
return DagUtil._dag(
constant_object,
params,
decorator=True,
file_name=file_name,
email=email,
template=template,
tags=tags,
notify_slack=notify_slack,
on_success_callback=on_success_callback,
on_failure_callback=on_failure_callback,
dag_id=dag_id,
interval=interval,
start_date=start_date,
)
[docs]
@staticmethod
def create_default_dag(
constant_object,
params: dict = {},
file_name: str = "",
email: str = "",
template: str = "",
tags: list[str] = [],
notify_slack: bool = True,
on_success_callback=None,
on_failure_callback=None,
dag_id: str = "",
interval=_UNSET,
start_date=None,
):
"""
This interface returns a dag object
:param constant_object: the constant object
:param params: the params to pass to the dag
:param file_name: the file name
:param email: the email to send to
:param template: the template to use
:param tags: the tags to use
:param notify_slack: if False, suppress Slack alerts from notify_failure (default True)
:param on_success_callback: DAG-level callback on successful run completion
:param on_failure_callback: DAG-level callback on failed run completion
:param dag_id: explicit dag_id (defaults to the id auto-built from file_name)
:param interval: explicit schedule (defaults to constant_object.INTERVAL;
pass None for an unscheduled, manual-trigger-only DAG)
:param start_date: explicit start_date (defaults to constant_object.START_DATE)
:return: the dag
"""
return DagUtil._dag(
constant_object,
params,
decorator=False,
file_name=file_name,
email=email,
template=template,
tags=tags,
notify_slack=notify_slack,
on_success_callback=on_success_callback,
on_failure_callback=on_failure_callback,
dag_id=dag_id,
interval=interval,
start_date=start_date,
)
[docs]
@staticmethod
def notify_success(
email: str = "",
subject: str = "",
content: str = "",
task_id: str = "success_notification",
trigger_rule: str = "all_success", # https://airflow.apache.org/docs/apache-airflow/1.10.9/concepts.html#trigger-rules
slack: bool = False,
):
"""
Success notification task: email + DAG Run Note + optional Slack @mention.
Returns a PythonOperator — place in the DAG graph like any other task.
Does not need to call .execute().
:param email: the email to send to
:param subject: the subject to use
:param content: the content to use
:param task_id: the task id
:param trigger_rule: the trigger rule
:param slack: if True, also post to Config.SLACK_ALERTS_CHANNEL (default False)
:return: PythonOperator
"""
_email = email
_subject = subject
_content = content
_slack = slack
def _run(**context):
dag_id = context["dag"].dag_id
run_id = context["dag_run"].run_id
duration = _format_duration(context["dag_run"].start_date)
trigger = None
try:
trigger = _get_trigger_user(dag_id, run_id)
except Exception:
pass
action = trigger["action"] if trigger else "triggered"
actor = trigger["display_name"] if trigger else "scheduler"
SERVER = "Prod" if Config.PRODUCTION else "Dev"
trigger_label = "↺ Rerun" if action == "rerun" else "▶ Initially triggered"
owner_set = {_email} if isinstance(_email, str) else set(_email)
cc_emails = (
[trigger["email"]]
if trigger and trigger["email"] and trigger["email"] not in owner_set
else []
)
# Intentionally no conn_id — uses Airflow's default email connection,
# matching original notify_success behavior. notify_failure uses
# conn_id=Config.SENDGRID_CONN (SendGrid); both connections were set up
# by the team and work correctly — do not consolidate without verifying both.
EmailOperator(
task_id=f"_{task_id}_email",
to=_email,
cc=cc_emails,
subject=_subject or f"[{SERVER}] ✅ SUCCESS | {dag_id} | {trigger_label} by {actor}",
html_content=_content or (
f'<strong>DAG:</strong> <a href="{Config.AIRFLOW_URL}/dags/{dag_id}/grid?dag_run_id={quote(run_id, safe="")}">{dag_id}</a>'
f'<br><br><strong>STATUS:</strong> ✅ SUCCESS'
f'<br><br><strong>TRIGGERED BY:</strong> {trigger_label} by {actor}'
f'<br><br><strong>DURATION:</strong> {duration}'
f'<br><br><strong>TIME:</strong> {_now_str()}'
f'<br><br><strong>OWNER:</strong> {context["dag"].owner}'
),
).execute(context)
# DAG Run note + Slack: best-effort, must not suppress the email above
try:
_append_dag_run_note(
dag_id,
run_id,
(
f"### ✅ SUCCESS · {'↺ rerun' if action == 'rerun' else '▶ initially triggered'}"
f" by: {actor}\n"
f"{_now_str()} · duration: {duration}"
),
)
if _slack:
_slack_id = _get_slack_id(trigger["email"] if trigger else None)
_mention = f"<@{_slack_id}> " if _slack_id else ""
NotifySlackOperator(
task_id=f"_{task_id}_slack",
channel=Config.SLACK_ALERTS_CHANNEL,
text=(
f":large_green_circle: *SUCCESS* | <{Config.AIRFLOW_URL}/dags/{dag_id}/grid?dag_run_id={quote(run_id, safe='')}|{dag_id}>\n"
f"{'↺ rerun' if action == 'rerun' else '▶ initially triggered'} by: {_mention}*{actor}* | {_now_str()} | duration: {duration}"
),
).execute(context)
except Exception:
pass
return PythonOperator(
task_id=task_id,
python_callable=_run,
trigger_rule=trigger_rule,
)
[docs]
@staticmethod
def email_attachments(
email: str = "",
cc: list[str] = [],
subject_dag_id: str = "",
subject: str = "",
content: str = "",
task_id: str = "email_attachments",
files: list[str] = [],
# **context,
):
"""
Configured to send dataframe as csv attachement in an email, but easy to drop in another operator
Calls operator.execute() because this is required to be a callback function -- https://gist.github.com/mholtzscher/9648cfd27769d1df6a6ed855fdd7bd7a
:param email: the email to send to
:param cc: the cc to send to
:param subject_dag_id: the subject dag id
:param subject: the subject to use
:param content: the content to use
:param task_id: the task id
:param files: the files to attach
:return: the email operator
"""
extended_subject = f"Airflow env={'Prod' if Config.PRODUCTION else 'Dev'} DAG={subject_dag_id}: {subject}"
return EmailOperator(
task_id=task_id,
to=email,
cc=cc,
subject=extended_subject,
html_content=content or Config.EMAIL_CONTENT,
files=files,
) # .execute(context)
[docs]
@staticmethod
def is_prod_instance() -> bool:
"""
A utility that will return a boolean value indicating if the current envrionment is prod
or not
:return: True if prod, False otherwise
"""
return Config.PRODUCTION