pipeloom-airflow 0.6.0__py3-none-any.whl

pipeloom_airflow/__init__.py

@@ -0,0 +1,35 @@
+"""Apache Airflow renderer for pipeloom.
+
+Importing the package imports nothing: names load on first use (PEP 562). Only the modules
+`_compat` and `_compat_k8s`, which arrive with the DAG classes, may import `airflow`.
+"""
+
+from __future__ import annotations
+
+import importlib
+from typing import Any
+
+_LAZY = {
+    "ExecutionConfig": "config",
+    "ExecutionMode": "config",
+    "KubernetesConfig": "config",
+    "LoadMode": "config",
+    "ProjectConfig": "config",
+    "RenderConfig": "config",
+    "validate_configs": "config",
+    "PipeloomAirflowError": "errors",
+    "LoadedPlan": "load",
+    "load_plan": "load",
+    "TaskPlan": "render",
+    "UnitTask": "render",
+    "render_plan": "render",
+}
+
+__all__ = list(_LAZY)
+
+
+def __getattr__(name: str) -> Any:
+    module = _LAZY.get(name)
+    if module is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    return getattr(importlib.import_module(f"{__name__}.{module}"), name)

pipeloom_airflow/_compat.py

@@ -0,0 +1,92 @@
+"""The one module of the plugin that imports `airflow` (with `_compat_k8s`). Every name the plugin
+uses from Airflow comes from here, so a move between Airflow minors is fixed in one place."""
+
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any
+
+from pipeloom_airflow.version import require_airflow
+
+# Before any Airflow import: an unsupported or missing Airflow is PLA005, not an ImportError.
+require_airflow()
+
+from airflow.sdk import (
+    DAG,
+    Asset,
+    BaseHook,
+    BaseOperator,
+    BaseSensorOperator,
+    Param,
+    TaskGroup,
+)
+from airflow.sdk.exceptions import (
+    AirflowException,
+    AirflowFailException,
+    AirflowSkipException,
+)
+from airflow.triggers.base import BaseTrigger, TriggerEvent
+
+if TYPE_CHECKING:
+    from airflow.sdk import Context
+else:
+    Context = Mapping[str, Any]
+
+__all__ = [
+    "DAG",
+    "AirflowException",
+    "AirflowFailException",
+    "AirflowSkipException",
+    "Asset",
+    "BaseHook",
+    "BaseOperator",
+    "BaseSensorOperator",
+    "BaseTrigger",
+    "Context",
+    "Param",
+    "TaskGroup",
+    "TriggerEvent",
+    "aget_connection",
+    "airflow_version",
+    "get_connection",
+    "get_task_states",
+    "mask_secret",
+]
+
+
+def get_connection(conn_id: str) -> Any:
+    """The connection, from a worker."""
+    return BaseHook.get_connection(conn_id)
+
+
+def mask_secret(value: str) -> None:
+    """Hide `value` in task and triggerer logs."""
+    import airflow.sdk
+
+    masker = getattr(airflow.sdk, "mask_secret", None)
+    if masker is not None:
+        masker(value)
+
+
+def airflow_version() -> str:
+    from airflow import __version__
+
+    return str(__version__)
+
+
+def get_task_states(
+    ti: Any, dag_id: str, task_ids: Sequence[str], run_id: str
+) -> dict[str, str | None]:
+    """Task id -> state of the given tasks in one DAG run, as the task SDK reports them."""
+    found = ti.get_task_states(dag_id=dag_id, task_ids=list(task_ids), run_ids=[run_id])
+    states = found.get(run_id, {}) if isinstance(found, Mapping) else {}
+    return {str(k): (None if v is None else str(v)) for k, v in states.items()}
+
+
+async def aget_connection(conn_id: str) -> Any:
+    """The connection, without blocking the event loop of a trigger."""
+    getter = getattr(BaseHook, "aget_connection", None)
+    if getter is not None:
+        return await getter(conn_id)
+    return await asyncio.to_thread(BaseHook.get_connection, conn_id)

pipeloom_airflow/_compat_k8s.py

@@ -0,0 +1,36 @@
+"""The only module besides `_compat` that imports the cncf-kubernetes provider (and the
+`kubernetes` client it brings). Import it only when the extra `kubernetes` is installed."""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Mapping
+from typing import Any
+
+from airflow.providers.cncf.kubernetes.operators.pod import KubernetesPodOperator
+from kubernetes.client import models as k8s
+
+__all__ = ["KubernetesPodOperator", "merge_env_vars", "pod_name"]
+
+POD_NAME_PREFIX = "pipeloom-"
+POD_NAME_BASE_MAX = 54
+"""The provider adds `-` and 8 random characters to a name, and a pod name is 63 at most."""
+
+
+def pod_name(task_id: str) -> str:
+    """`pipeloom-<task id>` as a DNS-1123 label: lower case, `[a-z0-9-]`, at most 54 characters."""
+    name = re.sub(r"[^a-z0-9-]+", "-", f"{POD_NAME_PREFIX}{task_id}".lower())
+    return name[:POD_NAME_BASE_MAX].strip("-")
+
+
+def merge_env_vars(env: Mapping[str, str], existing: Any) -> Any:
+    """`ExecutionConfig.env` plus the `env_vars` the user put in `pod_kwargs` (a dict or a list of
+    `V1EnvVar`); the pod_kwargs entry wins. Returns what `KubernetesPodOperator(env_vars=)` takes."""
+    if not env:
+        return existing
+    if not existing:
+        return dict(env)
+    if isinstance(existing, Mapping):
+        return {**env, **existing}
+    named = {v.name for v in existing}
+    return [*(k8s.V1EnvVar(name=k, value=v) for k, v in env.items() if k not in named), *existing]

pipeloom_airflow/commands.py

@@ -0,0 +1,164 @@
+"""The command lines a task runs, and the run identity every task of a DAG run shares."""
+
+from __future__ import annotations
+
+import ntpath
+import posixpath
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from typing import Literal
+
+from pipeloom.diagnostics import Report
+from pipeloom.platform import check_orchestrator_id
+from pipeloom_airflow import naming
+from pipeloom_airflow.errors import PipeloomAirflowError
+
+
+@dataclass(frozen=True)
+class RunContext:
+    run_id: str
+    run_started_at: datetime
+    full_refresh: bool
+    orchestrator_ids: Mapping[str, str]
+
+
+def _bool(value: object) -> bool | None:
+    if isinstance(value, bool):
+        return value
+    if isinstance(value, str) and value.strip().lower() in ("true", "false"):
+        return value.strip().lower() == "true"
+    return None
+
+
+def run_context(
+    *,
+    dag_id: str,
+    airflow_run_id: str,
+    run_after: datetime,
+    task_id: str,
+    try_number: int,
+    group_id: str | None,
+    full_refresh: object,
+) -> RunContext:
+    """What a task knows about its run. `full_refresh` is a bool or "true"/"false"; anything
+    else is PLA002. A naive `run_after` is read as UTC."""
+    refresh = _bool(full_refresh)
+    if refresh is None:
+        report = Report()
+        report.add(
+            "PLA002",
+            f"full_refresh must be true or false, got {full_refresh!r}",
+            hint="trigger the DAG with the param full_refresh set to a boolean",
+        )
+        raise PipeloomAirflowError(report.sorted())
+    started = run_after.replace(tzinfo=UTC) if run_after.tzinfo is None else run_after
+    ids = {
+        "airflow_dag_id": dag_id,
+        "airflow_run_id": airflow_run_id,
+        "airflow_task_id": task_id,
+        "airflow_try_number": str(try_number),
+    }
+    for key, value in ids.items():
+        check_orchestrator_id(key, value)
+    return RunContext(
+        run_id=naming.airflow_run_id(dag_id, airflow_run_id, group_id),
+        run_started_at=started.astimezone(UTC),
+        full_refresh=refresh,
+        orchestrator_ids=ids,
+    )
+
+
+def _ids(ctx: RunContext) -> list[str]:
+    return [a for k, v in ctx.orchestrator_ids.items() for a in ("--orchestrator-id", f"{k}={v}")]
+
+
+@dataclass(frozen=True)
+class CoreCommand:
+    """The core invocation a task makes. PROJECT mode passes no manifest: every unit compiles
+    and checks the plan hash, so no cache path crosses a process boundary."""
+
+    project_dir: str
+    target: str
+    manifest: str | None
+    plan_hash: str
+    vars_json: str | None
+
+    def _base(self) -> list[str]:
+        return ["--project-dir", self.project_dir, "--target", self.target]
+
+    def _plan(self) -> list[str]:
+        return (["--manifest", self.manifest] if self.manifest else []) + [
+            "--expect-plan-hash",
+            self.plan_hash,
+        ]
+
+    def _vars(self) -> list[str]:
+        return ["--vars", self.vars_json] if self.vars_json else []
+
+    def run_argv(
+        self,
+        python: Sequence[str],
+        unit_id: str,
+        ctx: RunContext,
+        *,
+        log_format: Literal["json", "text"],
+    ) -> list[str]:
+        return [
+            *python,
+            "-m",
+            "pipeloom",
+            "run",
+            "--log-format",
+            log_format,
+            *self._base(),
+            *self._plan(),
+            "--select",
+            f"unit:{unit_id}",
+            "--run-id",
+            ctx.run_id,
+            "--run-started-at",
+            ctx.run_started_at.isoformat(),
+            "--full-refresh",
+            "true" if ctx.full_refresh else "false",
+            *_ids(ctx),
+            *self._vars(),
+        ]
+
+    def finalize_argv(
+        self,
+        python: Sequence[str],
+        ctx: RunContext,
+        *,
+        select: Sequence[str],
+        exclude: Sequence[str],
+        unit_states_path: str | None = None,
+        unit_states_json: str | None = None,
+    ) -> list[str]:
+        states = ["--unit-states", unit_states_path] if unit_states_path is not None else []
+        if unit_states_json is not None:
+            states += ["--unit-states-json", unit_states_json]
+        return [
+            *python,
+            "-m",
+            "pipeloom",
+            "finalize",
+            *self._base(),
+            *self._plan(),
+            "--run-id",
+            ctx.run_id,
+            "--run-started-at",
+            ctx.run_started_at.isoformat(),
+            *[a for s in select for a in ("--select", s)],
+            *[a for s in exclude for a in ("--exclude", s)],
+            *_ids(ctx),
+            *states,
+            *self._vars(),
+        ]
+
+
+def venv_python(venv_dir: str, *, platform: str) -> str:
+    """The interpreter of a virtualenv: `bin/python`, or `Scripts/python.exe` on win32."""
+    if platform == "win32":
+        return ntpath.join(venv_dir, "Scripts", "python.exe")
+    return posixpath.join(venv_dir, "bin", "python")

pipeloom_airflow/config.py

@@ -0,0 +1,182 @@
+"""Configuration of the Airflow plugin. Stdlib and the core only; nothing here touches the disk."""
+
+from __future__ import annotations
+
+import ntpath
+import posixpath
+import string
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from enum import StrEnum
+from typing import Any, Final, Literal
+
+from pydantic import JsonValue
+
+from pipeloom.diagnostics import Report
+from pipeloom_airflow.naming import PLACEHOLDERS
+
+
+class LoadMode(StrEnum):
+    MANIFEST = "manifest"
+    PROJECT = "project"
+
+
+class ExecutionMode(StrEnum):
+    LOCAL = "local"
+    VIRTUALENV = "virtualenv"
+    KUBERNETES = "kubernetes"
+    DATABRICKS = "databricks"
+    DOCKER = "docker"
+
+
+@dataclass(frozen=True)
+class ProjectConfig:
+    """Where the DAG file reads the project."""
+
+    project_dir: str
+    manifest_path: str | None = None
+    vars: Mapping[str, JsonValue] = field(default_factory=dict)
+    cache_dir: str | None = None
+
+
+@dataclass(frozen=True)
+class RenderConfig:
+    load_method: LoadMode = LoadMode.MANIFEST
+    select: tuple[str, ...] = ()
+    exclude: tuple[str, ...] = ()
+    group_by: Literal["none", "folder", "tag"] = "none"
+    """Visual TaskGroups only; every unit is still one task."""
+    emit_assets: bool = False
+    asset_uri: str = "pipeloom://{project}/{target}/{relation}"
+
+
+@dataclass(frozen=True)
+class KubernetesConfig:
+    image: str
+    namespace: str | None = None
+    python: str = "python"
+    pod_kwargs: Mapping[str, Any] = field(default_factory=dict)
+
+
+@dataclass(frozen=True)
+class ExecutionConfig:
+    mode: ExecutionMode = ExecutionMode.LOCAL
+    project_dir: str | None = None
+    manifest_path: str | None = None
+    """Where units execute; the defaults are ProjectConfig's."""
+    virtualenv_dir: str | None = None
+    kubernetes: KubernetesConfig | None = None
+    env: Mapping[str, str] = field(default_factory=dict)
+    databricks_conn_id: str = "databricks_default"
+    databricks_pipeline: str = "main"
+    databricks_ensure_job: bool = False
+    databricks_poll: Literal["per_task", "shared"] = "per_task"
+    """`per_task`: each sensor polls its own task, so results show live. `shared`: the launch
+    waits for the whole run and the sensors read its snapshot."""
+    databricks_poll_interval: float = 30.0
+    databricks_timeout: float = 86400.0
+    """Seconds the launch waits for the job run before it times out (and retries, attaching)."""
+    databricks_client_factory: str | None = None
+    """`module:attr`, for tests only."""
+
+
+OPERATOR_ARGS_RESERVED: Final = (
+    "task_id",
+    "retries",
+    "retry_delay",
+    "trigger_rule",
+    "outlets",
+    "execution_timeout",
+    "deferrable",
+)
+DATABRICKS_ARGS_RESERVED: Final = (
+    "timeout",  # databricks_timeout
+    "poll",
+    "poll_interval",
+    "conn_id",
+    "client_factory",
+    "wait_timeout_s",
+    "task_key",
+    "launch_task_id",
+    "max_attempts",
+    "command",
+    "pipeline",
+    "ensure_job",
+    "select",
+    "exclude",
+    "group_id",
+)
+"""Constructor arguments of the DATABRICKS-mode operators that pipeloom sets itself."""
+POD_KWARGS_RESERVED: Final = (
+    "cmds",
+    "arguments",
+    "image",
+    "name",
+    "do_xcom_push",
+    "task_group",
+    *OPERATOR_ARGS_RESERVED,
+)
+
+
+def _absolute(path: str) -> bool:
+    """Absolute on either platform: the worker that runs a unit may not be the parser's OS."""
+    return posixpath.isabs(path) or ntpath.isabs(path)
+
+
+def validate_configs(
+    project: ProjectConfig,
+    render: RenderConfig,
+    execution: ExecutionConfig,
+    operator_args: Mapping[str, Any],
+    report: Report,
+) -> None:
+    """Report PLA001 and PLA002. Strings only: execution-side paths are never touched."""
+
+    def bad(message: str, hint: str | None = None) -> None:
+        report.add("PLA002", message, hint=hint)
+
+    if execution.mode is ExecutionMode.DOCKER:
+        report.add(
+            "PLA001",
+            "execution mode 'docker' is not supported",
+            hint="use local, virtualenv, kubernetes or databricks",
+        )
+    if render.load_method is LoadMode.MANIFEST and not project.manifest_path:
+        bad("load_method 'manifest' needs project.manifest_path", "or use load_method 'project'")
+    paths = {
+        "project.project_dir": project.project_dir,
+        "project.manifest_path": project.manifest_path,
+        "project.cache_dir": project.cache_dir,
+        "execution.project_dir": execution.project_dir,
+        "execution.manifest_path": execution.manifest_path,
+        "execution.virtualenv_dir": execution.virtualenv_dir,
+    }
+    for name, value in paths.items():
+        if value is not None and not _absolute(value):
+            bad(f"{name} must be an absolute path, got '{value}'")
+    if execution.mode is ExecutionMode.VIRTUALENV and not execution.virtualenv_dir:
+        bad("execution mode 'virtualenv' needs execution.virtualenv_dir")
+    for name in ("databricks_timeout", "databricks_poll_interval"):
+        if not getattr(execution, name) > 0:
+            bad(f"execution.{name} must be a positive number of seconds")
+    if execution.mode is ExecutionMode.KUBERNETES:
+        if execution.kubernetes is None or not execution.kubernetes.image:
+            bad("execution mode 'kubernetes' needs execution.kubernetes.image")
+        else:
+            for key in sorted(set(execution.kubernetes.pod_kwargs) & set(POD_KWARGS_RESERVED)):
+                bad(f"pod_kwargs.{key} is set by pipeloom", "remove it from pod_kwargs")
+            both = set(execution.kubernetes.pod_kwargs) & set(operator_args)
+            for key in sorted(both - set(POD_KWARGS_RESERVED)):
+                bad(f"'{key}' is in both pod_kwargs and operator_args", "keep it in one of them")
+    reserved = set(OPERATOR_ARGS_RESERVED)
+    if execution.mode is ExecutionMode.DATABRICKS:
+        reserved |= set(DATABRICKS_ARGS_RESERVED)
+    for key in sorted(set(operator_args) & reserved):
+        bad(f"operator_args.{key} is set by pipeloom", "remove it from operator_args")
+    try:
+        fields = {f for _, f, _, _ in string.Formatter().parse(render.asset_uri) if f is not None}
+    except ValueError as err:
+        bad(f"asset_uri is not a valid template: {err}")
+        return
+    for name in sorted(fields - set(PLACEHOLDERS)):
+        bad(f"asset_uri has an unknown placeholder '{{{name}}}'", f"use: {', '.join(PLACEHOLDERS)}")