dte-airflow 0.1.0__py3-none-any.whl

dte/__init__.py

@@ -0,0 +1,17 @@
+"""DTE. Importing the package applies the one dependency shim below."""
+import sqlglot.expressions as _sge
+
+# Airflow 3.3's constraints pin sqlglot 30, which moved Drop's target from
+# `this` to `tables`. ibis 12 still builds `Drop(this=...)` in every backend,
+# so each DROP (temp views, overwrite, drop_table) lost its name. Map it here,
+# once, for all of them (plan 073).
+# ponytail: global shim; delete when an ibis release builds `Drop(tables=...)`.
+if "tables" in _sge.Drop.arg_types and "this" not in _sge.Drop.arg_types:
+    _drop_init = _sge.Drop.__init__
+
+    def _init(self, **args):
+        if "this" in args and "tables" not in args:
+            args["tables"] = [args.pop("this")]
+        _drop_init(self, **args)
+
+    _sge.Drop.__init__ = _init

dte/airflow/assets.py

@@ -0,0 +1,57 @@
+"""Airflow Asset URIs for what a pipeline reads and writes (plan 089).
+
+Airflow's Assets view is the lineage graph, and an "after upstream" schedule
+triggers on these URIs, so a destination and a source naming the same table
+must produce the same string.
+
+- Files keep their own URL (`s3://…`, `gs://…`); a local path becomes `file://`.
+- Tables are `dte://<conn_id>/<schema>/<table>`. Not `postgres://`: provider
+  URI normalizers (postgres wants database/schema/table, mysql database/table)
+  would reject a conn_id-based URI at parse time and break the DAG, and a
+  source doesn't know its dialect without opening the connection.
+- Only the conn_id ever goes in, never a URL from a connection secret; userinfo
+  in a configured URL is dropped. No airflow import: the API lists these too.
+"""
+from __future__ import annotations
+
+from urllib.parse import urlsplit, urlunsplit
+
+
+def _location(path: str) -> str:
+    if "://" not in path:
+        return "file://" + ("" if path.startswith("/") else "/") + path
+    p = urlsplit(path)
+    return urlunsplit(p._replace(netloc=p.netloc.rpartition("@")[2]))
+
+
+def asset_uri(obj, table: str | None = None) -> str | None:
+    """URI for a `Destination`, or a `Source` plus the table a pipeline reads
+    from it. None when there is nothing addressable (rest_api)."""
+    cfg = obj.config
+    cfg = cfg.model_dump(by_alias=True) if hasattr(cfg, "model_dump") else dict(cfg or {})
+    kind = cfg.get("type") or getattr(obj, "source_type", None)
+    if kind in ("postgres", "sql_database"):
+        return f"dte://{obj.connection_id}/{cfg.get('schema') or 'default'}/{table or cfg.get('table') or obj.table}"
+    if kind == "mongodb":
+        return f"dte://{obj.connection_id}/{cfg['database']}/{cfg['collection']}"
+    if cfg.get("path"):
+        return _location(cfg["path"])
+    if cfg.get("bucket_url"):
+        return _location(f"{cfg['bucket_url'].rstrip('/')}/{cfg.get('file_glob', '')}")
+    return None
+
+
+def lineage(config: dict, sources: dict, destinations: dict) -> tuple[list[str], list[str]]:
+    """(inlets, outlets) URIs for a pipeline config, given Sources and
+    Destinations by id. Missing ones are skipped: the run fails on them, the
+    DAG shouldn't."""
+    from dte.models.pipeline import coerce_source_refs
+
+    inlets = [u for ref in coerce_source_refs(config.get("sources") or [])
+              if (s := sources.get(ref.source_id)) and (u := asset_uri(s, ref.table))]
+    outlets = []
+    if (config.get("pipeline_type") or "transform") == "transform":
+        d = destinations.get(config.get("output_id"))
+        if d and (u := asset_uri(d)):
+            outlets.append(u)
+    return inlets, outlets

dte/airflow/auth.py

@@ -0,0 +1,42 @@
+"""Airflow-backed auth for the DTE API (plan 065).
+
+Login, sessions and users are Airflow's. The caller is resolved with Airflow's
+own `get_user` (Bearer token or the `_token` cookie the UI sets), and DTE
+permissions map onto Airflow DAG access: read -> GET, write -> PUT,
+delete -> DELETE, trigger/test -> POST.
+"""
+from __future__ import annotations
+
+from fastapi import Request
+
+from dte.models.user import User
+
+_METHOD = {"read": "GET", "write": "PUT", "delete": "DELETE", "trigger": "POST", "test": "POST"}
+
+
+def _resolver():
+    from airflow.api_fastapi.app import get_auth_manager
+    from airflow.api_fastapi.core_api.security import get_user
+
+    return get_auth_manager, get_user
+
+
+async def get_current_user(request: Request) -> User:
+    """FastAPI dependency. Airflow imports are lazy so DTE's API still
+    imports (and tests override this) outside an Airflow environment."""
+    from dte.security.permissions import ALL_PERMISSIONS
+
+    get_auth_manager, get_user = _resolver()
+    # Positional: Airflow 3.3 renamed `oauth_token` to `_oauth_token` (unused
+    # at runtime); the order (request, oauth token, bearer) is stable.
+    user = await get_user(request, None, await _bearer(request))
+    mgr = get_auth_manager()
+    allowed = {m: mgr.is_authorized_dag(method=m, user=user) for m in set(_METHOD.values())}
+    perms = [p for p in ALL_PERMISSIONS if allowed[_METHOD[p.split(":")[1]]]]
+    return User(id=str(user.get_id()), email=user.get_name(), permissions=perms)
+
+
+async def _bearer(request: Request):
+    from fastapi.security import HTTPBearer
+
+    return await HTTPBearer(auto_error=False)(request)

dte/airflow/connection_store.py

@@ -0,0 +1,53 @@
+"""ConnectionStore over Airflow Connections (read-only, plan 074).
+
+`get` resolves like any Airflow task does (secrets backends, `AIRFLOW_CONN_*`
+env vars, then the metadata DB). `list` shows the metadata DB plus env-var
+connections; secrets backends can't be enumerated. Only connections whose
+provider is installed are offered.
+"""
+from __future__ import annotations
+
+import asyncio
+import os
+
+from dte.airflow.connections import available_kinds, from_airflow
+from dte.models.pipeline import Connection
+
+
+def _get(conn_id: str):
+    from airflow.hooks.base import BaseHook
+
+    return BaseHook.get_connection(conn_id)
+
+
+def _all_conn_ids() -> list[str]:
+    from airflow.models.connection import Connection as AConn
+    from airflow.utils.session import create_session
+    from sqlalchemy import select
+
+    with create_session() as s:
+        ids = set(s.scalars(select(AConn.conn_id)))
+    ids |= {k[len("AIRFLOW_CONN_"):].lower() for k in os.environ if k.startswith("AIRFLOW_CONN_")}
+    return sorted(ids)
+
+
+def _resolve(conn_id: str) -> Connection | None:
+    try:
+        c = from_airflow(_get(conn_id))
+    except Exception:  # unknown id, unmapped type (redis, ...), or a provider that won't build a hook
+        return None
+    return c if c.connection_type in available_kinds() else None
+
+
+class AirflowConnectionStore:
+    async def get(self, id: str) -> Connection | None:
+        return await asyncio.to_thread(_resolve, id)
+
+    async def list(self, limit: int = 50, offset: int = 0) -> list[Connection]:
+        def _list():
+            rows = (_resolve(i) for i in _all_conn_ids())
+            return [c for c in rows if c][offset:offset + limit]
+
+        return [Connection.model_construct(id=c.id, name=c.name, connection_type=c.connection_type,
+                                           credentials={})
+                for c in await asyncio.to_thread(_list)]

dte/airflow/connections.py

@@ -0,0 +1,205 @@
+"""Airflow Connection -> DTE Connection (plan 074).
+
+The installed provider does the auth. SQL URLs come from the provider hook;
+object storage carries only the conn_id and is opened through the provider's
+filesystem (`dte.sources.object_store`). A kind is offered only when its
+provider (and, for storage, the provider's filesystem) is installed.
+
+`google_cloud_platform` is ambiguous: extra `dte_type` picks "gcs" or "bigquery"
+(default gcs). The google provider's own `gcpbigquery` type is always BigQuery.
+"""
+from __future__ import annotations
+
+import re
+from functools import lru_cache
+
+from dte.models.pipeline import Connection
+from dte.sources.metered import METERED
+
+_SQL = {"postgres": "postgres", "mysql": "mysql", "mssql": "mssql", "sqlite": "sqlite", "snowflake": "snowflake"}
+_STORAGE = {"aws": "s3", "google_cloud_platform": "gcs", "wasb": "azure", "adls": "azure",
+            "azure_data_lake": "azure", "sftp": "sftp"}
+# Airflow connection extra -> credential key the opener reads (plan 093).
+# BigQuery enforces the bytes cap; Snowflake (no bytes cap) the timeout.
+_CAPS = {"dte_max_bytes_billed": "max_bytes_billed",
+         "dte_statement_timeout_seconds": "statement_timeout_seconds",
+         "dte_confirm_above_bytes": "confirm_above_bytes"}
+# The fsspec scheme each storage kind needs the provider to supply.
+STORAGE_SCHEME = {"s3": "s3", "gcs": "gs", "azure": "abfs", "sftp": "sftp"}
+# Per DTE kind: the Airflow conn_types that map to it, and what to install
+# when none is available (shown in the UI).
+KINDS: dict[str, tuple[set[str], str]] = {
+    "local_fs": ({"fs"}, "apache-airflow-providers-standard"),
+    "postgres": ({"postgres"}, "apache-airflow-providers-postgres"),
+    "mysql": ({"mysql"}, "apache-airflow-providers-mysql"),
+    "mssql": ({"mssql"}, "apache-airflow-providers-microsoft-mssql"),
+    "sqlite": ({"sqlite"}, "apache-airflow-providers-sqlite"),
+    "snowflake": ({"snowflake"}, "apache-airflow-providers-snowflake"),
+    "bigquery": ({"google_cloud_platform", "gcpbigquery"}, "apache-airflow-providers-google"),
+    "mongodb": ({"mongo"}, "apache-airflow-providers-mongo"),
+    "rest_api": ({"http"}, "apache-airflow-providers-http"),
+    "s3": ({"aws"}, "apache-airflow-providers-amazon[s3fs]"),
+    "gcs": ({"google_cloud_platform"}, "apache-airflow-providers-google"),
+    "azure": ({"wasb", "adls", "azure_data_lake"}, "apache-airflow-providers-microsoft-azure"),
+    "sftp": ({"sftp"}, "apache-airflow-providers-sftp sshfs"),
+}
+
+
+def _plain_scheme(url: str) -> str:
+    """`postgresql+psycopg2://…` -> `postgresql://…`: DTE's openers and dlt
+    writers want the dialect, not SQLAlchemy's driver."""
+    return re.sub(r"^(\w+)\+\w+://", r"\1://", url)
+
+
+def _hook_url(hook) -> str:
+    try:
+        url = hook.sqlalchemy_url
+    except (AttributeError, NotImplementedError):
+        return _plain_scheme(hook.get_uri())
+    s = url.render_as_string(hide_password=False) if hasattr(url, "render_as_string") else str(url)
+    return _plain_scheme(s)
+
+
+def from_airflow(c) -> Connection:
+    """`c` is an Airflow Connection (or anything with the same attributes)."""
+    x, t = c.extra_dejson, c.conn_type
+    if t in _SQL:
+        kind, creds = _SQL[t], {"url": _hook_url(c.get_hook())}
+        if t == "snowflake":
+            # The provider keeps key-pair and OAuth auth out of the URL; the
+            # opener and writer build their clients from the hook instead.
+            creds["conn_id"] = c.conn_id
+            if x.get("warehouse"):
+                creds["warehouse"] = x["warehouse"]  # named in the run estimate
+    elif t == "mongo":
+        kind, creds = "mongodb", {"url": c.get_hook().uri}
+    elif t == "gcpbigquery" or (t == "google_cloud_platform" and x.get("dte_type") == "bigquery"):
+        # The project comes from the hook; credentials are fetched from it at open time.
+        kind = "bigquery"
+        creds = {"url": f"bigquery://{c.get_hook().project_id}", "conn_id": c.conn_id}
+    elif t in _STORAGE:
+        kind, creds = _STORAGE[t], {"conn_id": c.conn_id}
+    elif t == "http":
+        kind = "rest_api"
+        auth = ({"type": "http_basic", "username": c.login, "password": c.password} if c.login
+                else {"type": "bearer", "token": c.password} if c.password else {})
+        creds = {"auth": auth, "headers": x.get("headers") or {}}
+    elif t == "fs":
+        kind, creds = "local_fs", {}
+    else:
+        raise ValueError(f"airflow connection type {t!r} has no DTE mapping (conn_id={c.conn_id!r})")
+    if kind in METERED:
+        # Plan 093: per-connection caps, applied by the backend_registry openers.
+        for extra_key, cred_key in _CAPS.items():
+            if x.get(extra_key):
+                creds[cred_key] = int(x[extra_key])
+    return Connection.model_construct(id=c.conn_id, name=c.conn_id, connection_type=kind,
+                                      credentials=creds)
+
+
+@lru_cache(maxsize=1)
+def _installed() -> tuple[set[str], set[str]]:
+    """(conn_types with a provider hook, fsspec schemes a provider can open)."""
+    import fsspec
+    from airflow.providers_manager import ProvidersManager
+
+    conn_types = set(ProvidersManager().hooks)
+    schemes = set()
+    for scheme in STORAGE_SCHEME.values():
+        try:
+            fsspec.get_filesystem_class(scheme)  # s3 needs s3fs, gs gcsfs, abfs adlfs, sftp sshfs
+            schemes.add(scheme)
+        except (ImportError, ValueError):
+            pass
+    return conn_types, schemes
+
+
+def available_kinds() -> set[str]:
+    conn_types, schemes = _installed()
+    return {k for k, (types, _) in KINDS.items()
+            if types & conn_types and (k not in STORAGE_SCHEME or STORAGE_SCHEME[k] in schemes)}
+
+
+def to_dte(conn_id: str) -> tuple[Connection, dict]:
+    from airflow.hooks.base import BaseHook  # lazy: airflow is an optional extra
+
+    conn = from_airflow(BaseHook.get_connection(conn_id))
+    return conn, conn.credentials
+
+
+def http_headers(conn_id: str) -> dict[str, str]:
+    """Request headers for an `http_lookup` step (plan 085) from an Airflow
+    connection: `headers` in its extras, plus the password as a Bearer token.
+    A missing connection raises `ModelConnectionMissing`."""
+    from airflow.exceptions import AirflowNotFoundException
+    from airflow.hooks.base import BaseHook
+
+    try:
+        conn = BaseHook.get_connection(conn_id)
+    except AirflowNotFoundException:
+        raise ModelConnectionMissing(
+            f"airflow connection {conn_id!r} not found — create it in Airflow → "
+            "Admin → Connections (type 'http', token as password)"
+        ) from None
+    headers = {str(k): str(v) for k, v in (conn.extra_dejson.get("headers") or {}).items()}
+    if conn.password:
+        headers.setdefault("Authorization", f"Bearer {conn.password}")
+    return headers
+
+
+# Model keys (plan 072). The official providers' hooks own the conn type, the
+# default conn id and client construction (password = key, host = base_url,
+# Bedrock/Vertex via extras); DTE only picks the hook.
+_MODEL_HOOKS = {
+    "claude": ("airflow.providers.anthropic.hooks.anthropic", "AnthropicHook"),
+    "openai": ("airflow.providers.openai.hooks.openai", "OpenAIHook"),
+}
+
+
+class ModelConnectionMissing(RuntimeError):
+    """The step's model connection does not exist in Airflow."""
+
+
+def _hook_class(provider: str):
+    import importlib
+
+    module, name = _MODEL_HOOKS[provider]
+    return getattr(importlib.import_module(module), name)
+
+
+def model_client(provider: str, conn_id: str | None):
+    """The SDK client for `provider`, built from an Airflow connection."""
+    from airflow.exceptions import AirflowNotFoundException
+
+    hook_cls = _hook_class(provider)
+    conn_id = conn_id or hook_cls.default_conn_name
+    try:
+        return hook_cls(conn_id=conn_id).get_conn()
+    except AirflowNotFoundException:
+        raise ModelConnectionMissing(
+            f"airflow connection {conn_id!r} not found — create it in Airflow → "
+            f"Admin → Connections (type {hook_cls.conn_type!r}, API key as password)"
+        ) from None
+
+
+def model_connections(provider: str) -> list[str]:
+    """Conn ids usable by `provider`: every DB connection of its type, plus the
+    default when it resolves elsewhere (env var or secrets backend)."""
+    from airflow.exceptions import AirflowNotFoundException
+    from airflow.hooks.base import BaseHook
+    from airflow.models.connection import Connection as AConn
+    from airflow.utils.session import create_session
+    from sqlalchemy import select
+
+    hook_cls = _hook_class(provider)
+    with create_session() as s:
+        ids = list(s.scalars(
+            select(AConn.conn_id).where(AConn.conn_type == hook_cls.conn_type).order_by(AConn.conn_id)
+        ))
+    if hook_cls.default_conn_name not in ids:
+        try:
+            BaseHook.get_connection(hook_cls.default_conn_name)
+            ids.insert(0, hook_cls.default_conn_name)
+        except AirflowNotFoundException:
+            pass
+    return ids

dte/airflow/dag_factory.py

@@ -0,0 +1,112 @@
+"""One DAG per DTE pipeline, built from its latest revision (plan 068).
+
+Pipeline revision = DAG version: the revision id is a DAG tag, so every save
+changes the serialized DAG and Airflow records a new DagVersion. The
+pipeline's one schedule (if any, and enabled) is the DAG's cron.
+
+`DTE_PIPELINE_STORE=postgres`: read from the `dte` schema at parse time.
+`file` (plan 090): read `DTE_PIPELINES_DIR/*.yaml`, no DB; each file is its
+own revision (id = content hash) and parses in its own try/except, so a
+broken file drops only its own DAG.
+"""
+from __future__ import annotations
+
+import logging
+import re
+from datetime import datetime, timedelta
+from pathlib import Path
+
+from sqlmodel import Session, select
+
+log = logging.getLogger(__name__)
+
+_SAFE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+
+
+def build_dags() -> dict:
+    from dte.airflow.assets import lineage
+    from dte.config import get_settings
+
+    settings = get_settings()
+    if settings.pipeline_store == "file":
+        return _dags_from_files(Path(settings.pipelines_dir))
+    latest, sources, dests = _latest_revisions()
+    dags = (_dag_for(rev, *lineage(rev.config, sources, dests)) for rev in latest if _SAFE.match(rev.pipeline_id))
+    return {d.dag_id: d for d in dags}
+
+
+def _latest_revisions() -> tuple[list, dict, dict]:
+    from dte.airflow.db import sync_engine
+    from dte.models.pipeline import Destination, PipelineRevision, Source
+
+    R = PipelineRevision
+    with Session(sync_engine()) as s:
+        # Revisions cascade-delete with their pipeline, so this is "every pipeline".
+        latest = list(s.exec(select(R).distinct(R.pipeline_id).order_by(R.pipeline_id, R.changed_at.desc())).all())
+        sources = {x.id: x for x in s.exec(select(Source)).all()}
+        dests = {x.id: x for x in s.exec(select(Destination)).all()}
+        s.expunge_all()
+    return latest, sources, dests
+
+
+def _dags_from_files(root: Path) -> dict:
+    from dte.store.files import read_revision
+
+    dags = {}
+    for path in sorted(root.glob("*.yaml")):
+        if not _SAFE.match(path.stem):
+            continue
+        try:
+            # ponytail: no inlets/outlets in file mode (sources/destinations live
+            # in Postgres, and parse reads no DB); fix before defaulting to file.
+            dag = _dag_for(read_revision(path))
+        except Exception:
+            log.exception("DTE pipeline file %s skipped: no DAG until it is fixed", path)
+            continue
+        dags[dag.dag_id] = dag
+    return dags
+
+
+def _dag_for(rev, inlets=(), outlets=()):
+    from airflow.sdk import DAG, Asset
+    from airflow.timetables.trigger import CronTriggerTimetable
+
+    from dte.airflow.notify import callbacks
+    from dte.airflow.operators import DtePipelineOperator
+    from dte.airflow.trigger import dag_id_for, revision_tag
+    from dte.config import get_settings
+    from dte.store.files import dump_yaml
+
+    p, pid = rev.config, rev.pipeline_id
+    sc = p.get("schedule") or {}
+    cron = None
+    if sc and sc.get("enabled", True):
+        if sc.get("mode") == "assets":
+            cron = [Asset(u) for u in sc["assets"]]
+        else:
+            cron = CronTriggerTimetable(sc["cron_expr"], timezone=sc.get("timezone", "UTC"))
+    ptype = p.get("pipeline_type") or "transform"
+    with DAG(
+        dag_id=dag_id_for(pid), schedule=cron, start_date=datetime(2024, 1, 1), catchup=False,
+        max_active_runs=16 if sc.get("overlap") == "allow" else 1,
+        params=dict(sc.get("params") or {}),
+        tags=["dte", f"dte-pipeline:{pid}", f"dte-type:{ptype}", revision_tag(rev.id)],
+        dag_display_name=p.get("name") or pid,
+        # The DAG page's Docs: the pipeline as YAML (what `dte export` writes)
+        # and a way back to the editor. Relative, like the UI's Airflow links,
+        # so Airflow's <base href> resolves it under any subpath.
+        doc_md=f"[Open in DTE](plugin/dte/pipelines/detail?id={pid})\n\n```yaml\n{dump_yaml(p)}```\n",
+        # A paused DAG queues manual runs forever. Per DAG, so an existing
+        # Airflow needn't change its global default (plan 073).
+        is_paused_upon_creation=False,
+    ) as dag:
+        # Manual and scheduled runs share this one task (plan 068), so retries
+        # and notifications apply to both.
+        DtePipelineOperator(
+            task_id="run", pipeline_id=pid, revision_id=rev.id, eval_set_id=sc.get("eval_set_id"),
+            retries=sc.get("retries", 0), retry_delay=timedelta(seconds=sc.get("retry_delay_s", 300)),
+            execution_timeout=timedelta(seconds=sc.get("timeout_s") or get_settings().run_timeout_seconds),
+            inlets=[Asset(u) for u in inlets], outlets=[Asset(u) for u in outlets],
+            **callbacks(p.get("notifications") or []),
+        )
+    return dag

dte/airflow/dag_runs.py

@@ -0,0 +1,32 @@
+"""Dag runs of DTE DAGs, read with Airflow's own session (plan 075).
+
+api-server only: Airflow hides its metadata DB from DAG parsing and tasks.
+Reading through Airflow (not DTE's engine) is what lets `DTE_DATABASE_URL`
+point at a database that has no `dag_run` table.
+"""
+from __future__ import annotations
+
+import asyncio
+
+_COLS = ("run_id", "dag_id", "state", "queued_at", "start_date", "end_date", "run_type", "conf")
+
+
+def _read(dag_ids: list[str] | None, run_id: str | None) -> list[dict]:
+    from airflow.models.dagrun import DagRun
+    from airflow.utils.session import create_session
+    from sqlalchemy import select
+
+    q = select(*(getattr(DagRun, c) for c in _COLS))
+    q = q.where(DagRun.dag_id.in_(dag_ids)) if dag_ids is not None else q.where(DagRun.dag_id.like(r"dte\_\_%", escape="\\"))
+    if run_id is not None:
+        q = q.where(DagRun.run_id == run_id)
+    with create_session() as s:
+        rows = [dict(r._mapping) for r in s.execute(q)]
+    for r in rows:  # DagRunState / DagRunType are str enums
+        for c in ("state", "run_type"):
+            r[c] = getattr(r[c], "value", r[c])
+    return rows
+
+
+async def dag_runs(dag_ids: list[str] | None = None, run_id: str | None = None) -> list[dict]:
+    return await asyncio.to_thread(_read, dag_ids, run_id)

dte/airflow/db.py

@@ -0,0 +1,46 @@
+"""Engines for DTE's `dte` schema (plans 065, 075).
+
+URL: `DTE_DATABASE_URL`, any Postgres (a dedicated `dte` database is the
+recommended setup), else Airflow's `sql_alchemy_conn`. DTE never reads
+Airflow's tables through these engines; dag runs come from `dag_runs.py`.
+"""
+from __future__ import annotations
+
+import os
+import re
+
+from sqlalchemy import create_engine
+from sqlalchemy.ext.asyncio import create_async_engine
+
+SEARCH_PATH = "dte,public"
+_PG = re.compile(r"^postgres(ql)?(\+\w+)?://")
+
+
+def _base_url() -> str:
+    url = os.getenv("DTE_DATABASE_URL")
+    if not url:
+        from airflow.configuration import conf
+
+        url = conf.get("database", "sql_alchemy_conn")
+    if not _PG.match(url):
+        # Airflow 3 blanks its metadata DB URL in DAG parsing and task
+        # processes, which is exactly where DTE reads its schema (plan 073).
+        raise RuntimeError("DTE needs DTE_DATABASE_URL (a Postgres database for DTE's tables) set on every "
+                           "Airflow component; Airflow hides its own URL from DAG parsing and tasks")
+    return url
+
+
+def sync_url() -> str:
+    return _PG.sub("postgresql+psycopg2://", _base_url())
+
+
+def async_url() -> str:
+    return _PG.sub("postgresql+asyncpg://", _base_url())
+
+
+def sync_engine():
+    return create_engine(sync_url(), connect_args={"options": f"-csearch_path={SEARCH_PATH}"})
+
+
+def async_engine():
+    return create_async_engine(async_url(), connect_args={"server_settings": {"search_path": SEARCH_PATH}})

dte/airflow/notify.py

@@ -0,0 +1,102 @@
+"""Pipeline notifications through Airflow's own notifiers (plan 070).
+
+A notification is `{notifier, params, on}`: `notifier` is the class path of
+any notifier an installed provider registers (Slack, SMTP, Teams, SNS,
+PagerDuty, ...), `params` its constructor kwargs, `on` the run events that
+fire it. Templated params can use `{{ dte_message }}` and `{{ dte_event }}`.
+
+Only registered notifiers are ever imported, so a saved class path can't
+reach arbitrary code. Sending never raises: a broken notification must not
+mask the run's own outcome.
+"""
+from __future__ import annotations
+
+import importlib
+import inspect
+import logging
+from functools import partial
+
+log = logging.getLogger(__name__)
+
+_CALLBACK = {"failure": "on_failure_callback", "success": "on_success_callback", "retry": "on_retry_callback"}
+
+
+def available_notifiers() -> list[str]:
+    from airflow.providers_manager import ProvidersManager
+
+    return sorted(ProvidersManager().notification)
+
+
+def get_connection(conn_id: str):
+    from airflow.hooks.base import BaseHook
+
+    return BaseHook.get_connection(conn_id)
+
+
+def _load(path: str):
+    if path not in available_notifiers():
+        raise ValueError(f"{path!r} is not an installed Airflow notifier")
+    module, _, name = path.rpartition(".")
+    return getattr(importlib.import_module(module), name)
+
+
+def describe(path: str) -> dict:
+    """The notifier's constructor params, for the UI to prefill."""
+    cls = _load(path)
+    params = [
+        {"name": p.name, "required": p.default is inspect.Parameter.empty,
+         "default": None if p.default is inspect.Parameter.empty else repr(p.default)}
+        for p in list(inspect.signature(cls.__init__).parameters.values())[1:]
+        if p.kind not in (p.VAR_POSITIONAL, p.VAR_KEYWORD)
+    ]
+    return {"notifier": path, "name": cls.__name__, "params": params,
+            "doc": (inspect.getdoc(cls) or "").split("\n\n")[0]}
+
+
+def problem(n: dict) -> str | None:
+    """Why notification `n` can't send, or None. Checked at save time."""
+    from airflow.exceptions import AirflowNotFoundException
+
+    try:
+        cls = _load(n["notifier"])
+    except (ValueError, ImportError) as e:
+        return str(e)
+    params = n.get("params") or {}
+    try:
+        cls(**params)  # notifiers only store their args; this catches missing/unknown ones
+    except TypeError as e:
+        return f"{cls.__name__}: {e}"
+    for k, v in params.items():
+        if k.endswith("_conn_id") and isinstance(v, str) and "{{" not in v:
+            try:
+                get_connection(v)
+            except AirflowNotFoundException:
+                return f"no Airflow connection {v!r} (param {k}); create it under Admin > Connections"
+    return None
+
+
+def _message(event: str, context) -> str:
+    dag = context["dag"]
+    name = getattr(dag, "dag_display_name", None) or dag.dag_id
+    msg = f"DTE pipeline {name}: {event}\nDAG: {dag.dag_id}\nRun: {context['run_id']}"
+    if event != "success" and context.get("exception"):
+        msg += f"\nError: {context['exception']}"
+    return msg
+
+
+def send(n: dict, event: str, context) -> None:
+    try:
+        notifier = _load(n["notifier"])(**(n.get("params") or {}))
+        # Passed as variables, so text in the error is never itself rendered.
+        notifier({**context, "dte_message": _message(event, context), "dte_event": event})
+    except Exception as e:
+        log.warning("notification %s on %s failed: %s", n.get("notifier"), event, e)
+
+
+def callbacks(notifications: list[dict]) -> dict:
+    """Operator kwargs: `on_<event>_callback` lists for the given notifications."""
+    out: dict[str, list] = {}
+    for n in notifications or []:
+        for event in n.get("on") or ["failure"]:
+            out.setdefault(_CALLBACK[event], []).append(partial(send, n, event))
+    return out