pipeloom-databricks 0.5.0__py3-none-any.whl

pipeloom_databricks/testing/fake.py

@@ -0,0 +1,293 @@
+"""`FakeDatabricks`: an in-memory `DatabricksClient` for tests and case files.
+
+No network and no SDK. Tests inspect `jobs`, `files`, `runs` and `calls`, and make a call fail with
+`fail["create_job"] = ApiError(...)`. The scenario functions at the end are what a case file names
+in `--client-factory pipeloom_databricks.testing.fake:<scenario>`.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+from collections.abc import Iterator, Mapping, Sequence
+from dataclasses import replace
+
+from pydantic import JsonValue
+
+from pipeloom_databricks.client import ApiError, DirEntry, JobSummary, RunInfo, TaskRunInfo
+
+JOB_TAG = "pipeloom_job"
+
+
+def _parent(path: str) -> str:
+    return path.rsplit("/", 1)[0] if "/" in path.strip("/") else "/"
+
+
+class FakeDatabricks:
+    host = "https://fake.databricks.test"
+
+    def __init__(self, *, user: str = "me@example.com") -> None:
+        self.user = user
+        self.jobs: dict[int, dict[str, JsonValue]] = {}
+        self.files: dict[str, bytes] = {}
+        self.dirs: set[str] = {"/"}
+        self.modified: dict[str, int] = {}
+        self.volume_files: dict[str, bytes] = {}
+        self.secrets: dict[tuple[str, str], str] = {}
+        """(scope, key) -> the plain value; `get_secret` returns it base64 encoded."""
+        self.runs: dict[int, RunInfo] = {}
+        self.run_tokens: dict[str, int] = {}
+        self.repairs: list[tuple[int, tuple[str, ...], int | None]] = []
+        self.starts: list[tuple[int, dict[str, str], tuple[str, ...] | None, str | None]] = []
+        """Every `run_now`: (job_id, job_parameters, only, idempotency_token)."""
+        self.submissions: list[tuple[dict[str, JsonValue], str | None]] = []
+        """Every `submit_run`: (the spec, idempotency_token)."""
+        self.calls: list[tuple[str, str]] = []
+        """(method, first argument as text), in order."""
+        self.fail: dict[str, ApiError] = {}
+        self.stale_polls = 0
+        """After a repair, this many `get_run` calls still show the run as it was before it."""
+        self._stale: dict[int, tuple[RunInfo, int]] = {}
+        self._next = 100
+        self._tick = 0
+
+    def _enter(self, method: str, arg: object = "") -> None:
+        self.calls.append((method, str(arg)))
+        if method in self.fail:
+            raise self.fail[method]
+
+    def _id(self) -> int:
+        self._next += 1
+        return self._next
+
+    def _touch(self, path: str) -> None:
+        self._tick += 1
+        self.modified[path] = self._tick
+
+    def called(self, method: str) -> list[str]:
+        return [arg for name, arg in self.calls if name == method]
+
+    # jobs
+
+    def list_jobs(self, *, name: str | None = None) -> Iterator[JobSummary]:
+        self._enter("list_jobs", name or "")
+        for job_id, settings in sorted(self.jobs.items()):
+            if name is not None and settings.get("name") != name:
+                continue
+            tags = settings.get("tags")
+            tag_map = {str(k): str(v) for k, v in tags.items()} if isinstance(tags, dict) else {}
+            yield JobSummary(job_id, str(settings.get("name", "")), tag_map)
+
+    def add_job(self, settings: Mapping[str, JsonValue]) -> int:
+        """Put a job in without a call record (test setup)."""
+        job_id = self._id()
+        self.jobs[job_id] = json.loads(json.dumps(dict(settings)))
+        return job_id
+
+    def create_job(self, settings: Mapping[str, JsonValue]) -> int:
+        self._enter("create_job")
+        return self.add_job(settings)
+
+    def reset_job(self, job_id: int, settings: Mapping[str, JsonValue]) -> None:
+        self._enter("reset_job", job_id)
+        if job_id not in self.jobs:
+            raise ApiError(404, f"Job {job_id} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        self.jobs[job_id] = json.loads(json.dumps(dict(settings)))
+
+    # runs
+
+    def _start(self, job_id: int | None, params: Mapping[str, str], token: str | None) -> int:
+        if token is not None and token in self.run_tokens:
+            return self.run_tokens[token]
+        run_id = self._id()
+        self.runs[run_id] = RunInfo(run_id, job_id, "PENDING", None, dict(params))
+        if token is not None:
+            self.run_tokens[token] = run_id
+        return run_id
+
+    def run_now(
+        self,
+        job_id: int,
+        *,
+        job_parameters: Mapping[str, str],
+        only: Sequence[str] | None,
+        idempotency_token: str | None,
+    ) -> int:
+        self._enter("run_now", job_id)
+        if job_id not in self.jobs:
+            raise ApiError(404, f"Job {job_id} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        self.starts.append(
+            (
+                job_id,
+                dict(job_parameters),
+                None if only is None else tuple(only),
+                idempotency_token,
+            )
+        )
+        return self._start(job_id, job_parameters, idempotency_token)
+
+    def submit_run(self, spec: Mapping[str, JsonValue], *, idempotency_token: str | None) -> int:
+        self._enter("submit_run")
+        self.submissions.append((json.loads(json.dumps(dict(spec))), idempotency_token))
+        return self._start(None, {}, idempotency_token)
+
+    def set_run(
+        self,
+        run_id: int,
+        *,
+        job_id: int | None = None,
+        life_cycle_state: str = "TERMINATED",
+        result_state: str | None = None,
+        tasks: Sequence[TaskRunInfo] = (),
+        job_parameters: Mapping[str, str] | None = None,
+        latest_repair_id: int | None = None,
+        run_page_url: str | None = None,
+        repair_states: Mapping[int, str] | None = None,
+    ) -> RunInfo:
+        """Put a run in without a call record (test setup, or the simulator)."""
+        run = RunInfo(
+            run_id,
+            job_id,
+            life_cycle_state,
+            result_state,
+            dict(job_parameters or {}),
+            tuple(tasks),
+            latest_repair_id,
+            run_page_url,
+            repair_states=dict(repair_states or {}),
+        )
+        self.runs[run_id] = run
+        return run
+
+    def get_run(self, run_id: int) -> RunInfo:
+        self._enter("get_run", run_id)
+        if run_id not in self.runs:
+            raise ApiError(404, f"Run {run_id} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        stale = self._stale.get(run_id)
+        if stale is not None and stale[1] > 0:
+            self._stale[run_id] = (stale[0], stale[1] - 1)
+            return stale[0]
+        return self.runs[run_id]
+
+    def list_runs(self, job_id: int, *, active_only: bool = False) -> Iterator[RunInfo]:
+        self._enter("list_runs", job_id)
+        for run in self.runs.values():
+            active = run.life_cycle_state not in ("TERMINATED", "SKIPPED", "INTERNAL_ERROR")
+            if run.job_id == job_id and (active or not active_only):
+                yield run
+
+    def repair_run(
+        self, run_id: int, *, rerun_tasks: Sequence[str], latest_repair_id: int | None
+    ) -> int:
+        self._enter("repair_run", run_id)
+        if run_id not in self.runs:
+            raise ApiError(404, f"Run {run_id} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        if latest_repair_id != self.runs[run_id].latest_repair_id:
+            raise ApiError(
+                400, "latest_repair_id is not the latest repair.", "INVALID_PARAMETER_VALUE"
+            )
+        repair_id = self._id()
+        self.repairs.append((run_id, tuple(rerun_tasks), latest_repair_id))
+        if self.stale_polls:
+            before = replace(self.runs[run_id], latest_repair_id=repair_id)
+            self._stale[run_id] = (before, self.stale_polls)
+        self.runs[run_id] = replace(
+            self.runs[run_id],
+            life_cycle_state="RUNNING",
+            latest_repair_id=repair_id,
+            repair_states={**self.runs[run_id].repair_states, repair_id: "RUNNING"},
+        )
+        return repair_id
+
+    # workspace
+
+    def current_user(self) -> str:
+        self._enter("current_user")
+        return self.user
+
+    def mkdirs(self, path: str) -> None:
+        self._enter("mkdirs", path)
+        here = path.rstrip("/") or "/"
+        while here not in self.dirs:
+            self.dirs.add(here)  # like the real API, a folder has no modification time
+            here = _parent(here)
+
+    def upload(self, path: str, data: bytes) -> None:
+        self._enter("upload", path)
+        if _parent(path) not in self.dirs:
+            raise ApiError(404, f"The parent of {path} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        self.files[path] = data
+        self._touch(path)
+
+    def list_dir(self, path: str) -> list[DirEntry]:
+        self._enter("list_dir", path)
+        here = path.rstrip("/") or "/"
+        if here not in self.dirs:
+            raise ApiError(404, f"{path} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        kids = [(p, True) for p in self.dirs if p != "/" and _parent(p) == here]
+        kids += [(p, False) for p in self.files if _parent(p) == here]
+        return [DirEntry(p, d, self.modified.get(p, 0)) for p, d in sorted(kids)]
+
+    def delete(self, path: str, *, recursive: bool) -> None:
+        self._enter("delete", path)
+        here = path.rstrip("/")
+        inside = [p for p in (*self.dirs, *self.files) if p.startswith(here + "/")]
+        if here not in self.dirs and here not in self.files:
+            raise ApiError(404, f"{path} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        if inside and not recursive:
+            raise ApiError(400, f"{path} is not empty.", "DIRECTORY_NOT_EMPTY")
+        for p in inside:
+            self.dirs.discard(p)
+            self.files.pop(p, None)
+        self.dirs.discard(here)
+        self.files.pop(here, None)
+
+    def read_file(self, path: str) -> bytes:
+        self._enter("read_file", path)
+        if path not in self.volume_files:
+            raise ApiError(404, f"{path} does not exist.", "NOT_FOUND")
+        return self.volume_files[path]
+
+    def get_secret(self, scope: str, key: str) -> str:
+        self._enter("get_secret", f"{scope}/{key}")
+        if (scope, key) not in self.secrets:
+            raise ApiError(404, f"Secret {scope}/{key} does not exist.", "RESOURCE_DOES_NOT_EXIST")
+        return base64.b64encode(self.secrets[(scope, key)].encode("utf-8")).decode("ascii")
+
+
+def empty() -> FakeDatabricks:
+    """A workspace with nothing in it."""
+    return FakeDatabricks()
+
+
+def duplicate_jobs() -> FakeDatabricks:
+    """Two jobs that both claim the job tag `shop.main.prod` (PLD006)."""
+    client = FakeDatabricks()
+    for _ in range(2):
+        client.add_job({"name": "pipeloom_shop_main_prod", "tags": {JOB_TAG: "shop.main.prod"}})
+    return client
+
+
+def stale_job() -> FakeDatabricks:
+    """A job for `shop.main.prod` deployed from other settings than any project renders (PLD008)."""
+    client = FakeDatabricks()
+    tags: dict[str, JsonValue] = {JOB_TAG: "shop.main.prod", "pipeloom_settings_hash": "0" * 64}
+    client.add_job({"name": "pipeloom_shop_main_prod", "tags": tags})
+    return client
+
+
+def active_run() -> FakeDatabricks:
+    """Job run 900 of job 101 is still running (PLD014)."""
+    client = FakeDatabricks()
+    job_id = client.add_job(
+        {"name": "pipeloom_shop_main_prod", "tags": {JOB_TAG: "shop.main.prod"}}
+    )
+    client.set_run(900, job_id=job_id, life_cycle_state="RUNNING")
+    return client
+
+
+def api_down() -> FakeDatabricks:
+    """A workspace whose job creation fails with a 500 (PLD009)."""
+    client = FakeDatabricks()
+    client.fail["create_job"] = ApiError(500, "INTERNAL_ERROR: the service is unavailable")
+    return client

pipeloom_databricks/testing/jobsim.py

@@ -0,0 +1,179 @@
+"""A local stand-in for the Jobs service: run a rendered job through the task wrapper.
+
+It substitutes the dynamic value references, runs the tasks in dependency order through
+`pipeloom_databricks.task.main` in this process, and honors `run_if`, `only` and a repair.
+It starts no cluster and calls no API. The project's engine, state store and artifacts root
+come from its own configuration, so a test points them at a temporary directory.
+"""
+
+from __future__ import annotations
+
+import re
+from collections.abc import Collection, Mapping
+from dataclasses import dataclass, field, replace
+from typing import Any
+
+from pydantic import JsonValue
+
+from pipeloom_databricks.client import (
+    ConnectionNotConfigured,
+    DatabricksClient,
+    TaskRunInfo,
+)
+from pipeloom_databricks.task import main
+from pipeloom_databricks.task_values import TASK_VALUE_KEY
+from pipeloom_databricks.testing.fake import FakeDatabricks
+
+START_MS = 1_700_000_000_000
+"""The job start every simulated run (and repair) reports."""
+_REFERENCE = re.compile(r"\{\{([^{}]+)\}\}")
+
+
+@dataclass
+class SimRun:
+    job_run_id: int
+    task_states: dict[str, str] = field(default_factory=dict)
+    """SUCCESS, FAILED, UPSTREAM_FAILED or EXCLUDED, for the tasks this call ran or skipped."""
+    exit_codes: dict[str, int] = field(default_factory=dict)
+    """The exit code of every task that ran."""
+    task_values: dict[str, Any] = field(default_factory=dict)
+    """The `pipeloom` task value each task set, by task key."""
+
+
+class _Recorder:
+    def __init__(self, run: SimRun, task_key: str) -> None:
+        self._run = run
+        self._key = task_key
+
+    def set(self, key: str, value: JsonValue) -> None:
+        if key == TASK_VALUE_KEY:
+            self._run.task_values[self._key] = value
+
+
+def _order(tasks: Mapping[str, Mapping[str, Any]]) -> list[str]:
+    """Dependencies first, ties by key."""
+    done: list[str] = []
+    waiting = sorted(tasks)
+    while waiting:
+        ready = [
+            k for k in waiting if all(d["task_key"] in done for d in tasks[k].get("depends_on", []))
+        ]
+        if not ready:
+            raise ValueError("the job's tasks form a cycle or depend on an unknown task")
+        done.append(ready[0])
+        waiting.remove(ready[0])
+    return done
+
+
+def _no_client() -> DatabricksClient:
+    raise ConnectionNotConfigured("the simulator was given no client")
+
+
+def _record(client: FakeDatabricks, run: SimRun, job_id: int, params: Mapping[str, str]) -> None:
+    """Put the task states known so far into the fake's run, as the Jobs API would report them
+    to the finalize task. A repair keeps the tasks it does not run again."""
+    old = client.runs.get(run.job_run_id)
+    tasks = {t.task_key: t for t in old.tasks} if old else {}
+    for key, state in run.task_states.items():
+        failed = state == "FAILED"
+        tasks[key] = TaskRunInfo(
+            key,
+            run.job_run_id * 1000 + len(tasks),
+            1,
+            "TERMINATED",
+            state,
+            "RUN_EXECUTION_ERROR" if failed else None,
+            "the task failed" if failed else None,
+        )
+    client.set_run(
+        run.job_run_id,
+        job_id=job_id,
+        life_cycle_state="RUNNING",
+        tasks=[tasks[k] for k in sorted(tasks)],
+        job_parameters=params,
+        latest_repair_id=old.latest_repair_id if old else None,
+    )
+
+
+def simulate_job(
+    settings: Mapping[str, JsonValue],
+    *,
+    job_run_id: int,
+    job_parameters: Mapping[str, str],
+    fail: Collection[str] = (),
+    only: Collection[str] | None = None,
+    rerun: Collection[str] | None = None,
+    job_id: int = 1,
+    start_ms: int = START_MS,
+    client: FakeDatabricks | None = None,
+) -> SimRun:
+    """Run the job once.
+
+    `fail` names tasks that die before they run anything (a lost cluster): exit 1, no attempt
+    file. `only` runs just those tasks; the others are EXCLUDED. `rerun` is a repair: only
+    those tasks run again, the others keep what an earlier call did. With `client`, the task
+    states of the run are kept in the fake as they happen, and finalize reads them from it; without
+    one finalize finds no connection and runs without states.
+    """
+    job: dict[str, Any] = dict(settings)
+    tasks: dict[str, Mapping[str, Any]] = {t["task_key"]: t for t in job["tasks"]}
+    params: dict[str, str] = {}
+    for item in job.get("parameters") or []:
+        params[item["name"]] = item["default"]
+    params.update(job_parameters)
+
+    run = SimRun(job_run_id)
+    factory = _no_client if client is None else (lambda: client)
+    for position, key in enumerate(_order(tasks), start=1):
+        task = tasks[key]
+        if only is not None and key not in only:
+            run.task_states[key] = "EXCLUDED"
+            continue
+        if rerun is not None and key not in rerun:
+            continue
+        if client is not None:
+            _record(client, run, job_id, params)
+        seen = [
+            run.task_states.get(d["task_key"], "SUCCESS")
+            for d in task.get("depends_on", [])
+            if run.task_states.get(d["task_key"]) != "EXCLUDED"
+        ]
+        run_if = task.get("run_if") or "ALL_SUCCESS"
+        if run_if not in ("ALL_SUCCESS", "ALL_DONE"):
+            raise ValueError(f"run_if {run_if} is not simulated")
+        if run_if == "ALL_SUCCESS" and any(s != "SUCCESS" for s in seen):
+            run.task_states[key] = "UPSTREAM_FAILED"
+            continue
+        if key in fail:
+            run.task_states[key], run.exit_codes[key] = "FAILED", 1
+            continue
+        values = {
+            "job.id": str(job_id),
+            "job.run_id": str(job_run_id),
+            "job.start_time.timestamp_ms": str(start_ms),
+            "task.name": key,
+            "task.run_id": str(job_run_id * 1000 + position),
+        }
+
+        def substitute(match: re.Match[str], values: dict[str, str] = values) -> str:
+            name = match.group(1).strip()
+            if name.startswith("job.parameters."):
+                return params.get(name.removeprefix("job.parameters."), match.group(0))
+            return values.get(name, match.group(0))
+
+        argv = [_REFERENCE.sub(substitute, p) for p in task["python_wheel_task"]["parameters"]]
+        try:
+            main(argv, task_values=_Recorder(run, key), client_factory=factory)
+            code = 0
+        except SystemExit as stop:
+            code = stop.code if isinstance(stop.code, int) else 0
+        run.exit_codes[key] = code
+        run.task_states[key] = "SUCCESS" if code == 0 else "FAILED"
+    if client is not None:
+        _record(client, run, job_id, params)
+        done = client.runs[job_run_id]
+        bad = any(t.result_state not in ("SUCCESS", "EXCLUDED") for t in done.tasks)
+        client.runs[job_run_id] = replace(
+            done, life_cycle_state="TERMINATED", result_state="FAILED" if bad else "SUCCESS"
+        )
+    return run

pipeloom_databricks/unity.py

@@ -0,0 +1,23 @@
+"""Unity Catalog: `catalog.schema.name` tables and views through the Spark session."""
+
+from __future__ import annotations
+
+from pipeloom.engine.types import CatalogCapabilities, Namespace
+from pipeloom_spark.catalog import SparkSessionCatalog
+
+
+class UnityCatalog(SparkSessionCatalog):
+    """Always three-part, backticked identifiers. A target whose naming left the catalog empty
+    gets the catalog from the options."""
+
+    capabilities = CatalogCapabilities(views=True)
+
+    def __init__(self, catalog: str) -> None:
+        super().__init__(catalog)
+
+    def ensure_namespace(self, ns: Namespace) -> None:
+        # CREATE SCHEMA IF NOT EXISTS needs CREATE SCHEMA on the catalog even when the schema
+        # exists, so a principal that only holds USE SCHEMA must not run it.
+        if self.spark.catalog.databaseExists(self._parts(ns.catalog, ns.schema)):
+            return
+        super().ensure_namespace(ns)