semblance-databricks 0.1.0__py3-none-any.whl

semblance_databricks/services/jobs/routes.py

@@ -0,0 +1,168 @@
+"""Jobs 2.2 (primary) and 2.1 representative aliases."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, Request
+
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.http import json_object
+from semblance_databricks.pagination import paginate
+from semblance_databricks.services.deps import mock_from, require_fail_stage
+from semblance_databricks.state import JobRecord, RunRecord
+
+_TERMINAL_RUNS = {"TERMINATED", "SKIPPED", "INTERNAL_ERROR"}
+
+
+def _add_jobs_routes(router: APIRouter, prefix: str) -> None:
+    slug = prefix.replace("/", "_")
+
+    @router.get(f"{prefix}/jobs/list", name=f"list_jobs{slug}")
+    def list_jobs(
+        request: Request,
+        limit: int | None = None,
+        page_token: str | None = None,
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        mock.state.maybe_tick_real(mock.config.clock)
+        jobs = sorted(mock.state.jobs.values(), key=lambda rec: rec.job_id)
+        items = [mock.state.job_json(job) for job in jobs]
+        page, nxt = paginate(
+            items,
+            page_size=limit,
+            page_token=page_token,
+            resource="jobs",
+            codec=mock.page_token_codec,
+            revision=mock.state.revision,
+            default_page_size=mock.config.default_page_size,
+            max_page_size=mock.config.max_page_size,
+        )
+        body: dict[str, Any] = {"jobs": page, "has_more": nxt is not None}
+        if nxt:
+            body["next_page_token"] = nxt
+        return body
+
+    @router.get(f"{prefix}/jobs/get", name=f"get_job{slug}")
+    def get_job(request: Request, job_id: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not job_id:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "job_id required")
+        rec = mock.state.jobs.get(str(job_id))
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Job not found")
+        return mock.state.job_json(rec)
+
+    @router.get(f"{prefix}/jobs/runs/get", name=f"get_run{slug}")
+    def get_run(request: Request, run_id: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        mock.state.maybe_tick_real(mock.config.clock)
+        if not run_id:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "run_id required")
+        rec = mock.state.runs.get(str(run_id))
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Run not found")
+        return mock.state.run_json(rec)
+
+    @router.post(f"{prefix}/jobs/create", name=f"create_job{slug}")
+    async def create_job(request: Request) -> dict[str, Any]:
+        require_fail_stage(request, "before_validate")
+        mock = mock_from(request)
+        body = await json_object(request)
+        require_fail_stage(request, "before_write")
+        jid = mock.state.allocate_numeric_id(mock.state.jobs, 1000)
+        if isinstance(body.get("settings"), dict):
+            settings = dict(body["settings"])
+        else:
+            settings = {key: value for key, value in body.items() if key != "job_id"}
+        if "name" not in settings:
+            settings["name"] = str(body.get("name") or "job")
+        mock.state.jobs[jid] = JobRecord(job_id=jid, settings=settings)
+        mock.state.permissions[("jobs", jid)] = {
+            "object_id": jid,
+            "object_type": "job",
+            "access_control_list": [],
+        }
+        mock.state.bump()
+        require_fail_stage(request, "after_write")
+        return {"job_id": int(jid)}
+
+    @router.post(f"{prefix}/jobs/reset", name=f"reset_job{slug}")
+    async def reset_job(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        jid = str(body.get("job_id", ""))
+        rec = mock.state.jobs.get(jid)
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Job not found")
+        settings_raw = body.get("new_settings") or body.get("settings") or {}
+        if not isinstance(settings_raw, dict):
+            raise DatabricksError(
+                400, "INVALID_PARAMETER_VALUE", "settings must be an object"
+            )
+        rec.settings = dict(settings_raw)
+        mock.state.bump()
+        return {}
+
+    @router.post(f"{prefix}/jobs/delete", name=f"delete_job{slug}")
+    async def delete_job(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        jid = str(body.get("job_id", ""))
+        if jid not in mock.state.jobs:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Job not found")
+        del mock.state.jobs[jid]
+        mock.state.bump()
+        return {}
+
+    @router.post(f"{prefix}/jobs/runs/submit", name=f"submit_run{slug}")
+    async def submit_run(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        rid = mock.state.allocate_numeric_id(mock.state.runs, 2000)
+        job_id = str(body["job_id"]) if body.get("job_id") is not None else None
+        mock.state.runs[rid] = RunRecord(
+            run_id=rid,
+            job_id=job_id,
+            run_name=str(body.get("run_name", "submit")),
+            life_cycle_state="PENDING",
+            ticks_remaining=2,
+            output={"notebook_output": {"result": "pending"}},
+        )
+        mock.state.bump()
+        return {"run_id": int(rid)}
+
+    @router.post(f"{prefix}/jobs/runs/cancel", name=f"cancel_run{slug}")
+    async def cancel_run(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        rid = str(body.get("run_id", ""))
+        rec = mock.state.runs.get(rid)
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Run not found")
+        if rec.life_cycle_state in _TERMINAL_RUNS:
+            raise DatabricksError(400, "INVALID_STATE", "Run is not active")
+        rec.canceling = True
+        rec.life_cycle_state = "TERMINATING"
+        mock.state.bump()
+        return {}
+
+    @router.get(f"{prefix}/jobs/runs/get-output", name=f"get_output{slug}")
+    def get_output(request: Request, run_id: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not run_id:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "run_id required")
+        rec = mock.state.runs.get(str(run_id))
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Run not found")
+        return {
+            "metadata": mock.state.run_json(rec),
+            "notebook_output": rec.output.get("notebook_output", {"result": ""}),
+        }
+
+
+def create_jobs_router() -> APIRouter:
+    router = APIRouter()
+    _add_jobs_routes(router, "/api/2.2")
+    _add_jobs_routes(router, "/api/2.1")
+    return router

semblance_databricks/services/permissions/__init__.py

@@ -0,0 +1,5 @@
+from semblance_databricks.services.permissions.routes import (
+    create_permissions_router,
+)
+
+__all__ = ["create_permissions_router"]

semblance_databricks/services/permissions/routes.py

@@ -0,0 +1,50 @@
+"""Permissions GET/PATCH."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, Request
+
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.http import json_object
+from semblance_databricks.services.deps import mock_from
+
+
+def _key(object_type: str, object_id: str) -> tuple[str, str]:
+    return (object_type, object_id)
+
+
+def create_permissions_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.get("/api/2.0/permissions/{object_type}/{object_id}")
+    def get_permissions(
+        request: Request, object_type: str, object_id: str
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        rec = mock.state.permissions.get(_key(object_type, object_id))
+        if rec is None:
+            raise DatabricksError(
+                404, "RESOURCE_DOES_NOT_EXIST", "Permissions not found"
+            )
+        return rec
+
+    @router.patch("/api/2.0/permissions/{object_type}/{object_id}")
+    async def patch_permissions(
+        request: Request, object_type: str, object_id: str
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        rec = mock.state.permissions.get(_key(object_type, object_id))
+        if rec is None:
+            raise DatabricksError(
+                404, "RESOURCE_DOES_NOT_EXIST", "Permissions not found"
+            )
+        body = await json_object(request)
+        acl = body.get("access_control_list")
+        if isinstance(acl, list):
+            rec["access_control_list"] = acl
+        mock.state.bump()
+        return rec
+
+    return router

semblance_databricks/services/secrets/__init__.py

@@ -0,0 +1,3 @@
+from semblance_databricks.services.secrets.routes import create_secrets_router
+
+__all__ = ["create_secrets_router"]

semblance_databricks/services/secrets/routes.py

@@ -0,0 +1,53 @@
+"""Workspace secrets metadata (values never returned)."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, Request
+
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.http import json_object
+from semblance_databricks.services.deps import mock_from
+from semblance_databricks.state import SecretScope
+
+
+def create_secrets_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.get("/api/2.0/secrets/scopes/list")
+    def list_scopes(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        return {"scopes": [{"name": s.name} for s in mock.state.secret_scopes.values()]}
+
+    @router.get("/api/2.0/secrets/list")
+    def list_secrets(request: Request, scope: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not scope:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "scope required")
+        rec = mock.state.secret_scopes.get(scope)
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Scope not found")
+        return {
+            "secrets": [
+                {"key": key, "last_updated_timestamp": 1_700_000_000_000}
+                for key in rec.keys
+            ]
+        }
+
+    @router.post("/api/2.0/secrets/put")
+    async def put_secret(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        scope = str(body.get("scope", ""))
+        key = str(body.get("key", ""))
+        if not scope or not key:
+            raise DatabricksError(
+                400, "INVALID_PARAMETER_VALUE", "scope and key required"
+            )
+        rec = mock.state.secret_scopes.setdefault(scope, SecretScope(name=scope))
+        rec.keys[key] = str(body.get("string_value", ""))
+        mock.state.bump()
+        return {}
+
+    return router

semblance_databricks/services/workspace/__init__.py

@@ -0,0 +1,3 @@
+from semblance_databricks.services.workspace.routes import create_workspace_router
+
+__all__ = ["create_workspace_router"]

semblance_databricks/services/workspace/routes.py

@@ -0,0 +1,39 @@
+"""Workspace get-status and current user."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import APIRouter, Request
+
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.services.deps import mock_from
+
+
+def create_workspace_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.get("/api/2.0/workspace/get-status")
+    def get_status(request: Request, path: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not path:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "path required")
+        obj = mock.state.workspace_objects.get(path)
+        if obj is None:
+            raise DatabricksError(
+                404,
+                "RESOURCE_DOES_NOT_EXIST",
+                f"Path {path} does not exist",
+            )
+        return {
+            "path": obj.path,
+            "object_type": obj.object_type,
+            "language": obj.language,
+        }
+
+    @router.get("/api/2.0/preview/scim/v2/Me")
+    def me(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        return dict(mock.state.user)
+
+    return router

semblance_databricks/state.py

@@ -0,0 +1,253 @@
+"""Process-local Databricks workspace graph and virtual clock."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from typing import Any
+
+from semblance_databricks.ids import make_id
+
+
+@dataclass
+class ClusterRecord:
+    cluster_id: str
+    cluster_name: str
+    spark_version: str
+    node_type_id: str
+    state: str
+    state_message: str = ""
+    num_workers: int = 1
+    ticks_remaining: int = 0
+    fail_after: bool = False
+    libraries: list[dict[str, Any]] = field(default_factory=list)
+    events: list[dict[str, Any]] = field(default_factory=list)
+    deleted: bool = False
+
+
+@dataclass
+class JobRecord:
+    job_id: str
+    settings: dict[str, Any]
+    created_time: int = 1_700_000_000_000
+    creator_user_name: str = "user@acme.example"
+
+
+@dataclass
+class RunRecord:
+    run_id: str
+    job_id: str | None
+    run_name: str
+    life_cycle_state: str
+    result_state: str | None = None
+    ticks_remaining: int = 0
+    output: dict[str, Any] = field(default_factory=dict)
+    canceling: bool = False
+
+
+@dataclass
+class WarehouseRecord:
+    id: str
+    name: str
+    state: str = "RUNNING"
+    cluster_size: str = "X-Small"
+
+
+@dataclass
+class SecretScope:
+    name: str
+    keys: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass
+class DbfsNode:
+    path: str
+    is_dir: bool
+    data: bytes = b""
+
+
+@dataclass
+class WorkspaceObject:
+    path: str
+    object_type: str = "NOTEBOOK"
+    language: str = "PYTHON"
+
+
+class DatabricksState:
+    def __init__(self, seed: int | None = 42) -> None:
+        self.seed = seed
+        self.revision = 0
+        self.seq = 0
+        self.clusters: dict[str, ClusterRecord] = {}
+        self.jobs: dict[str, JobRecord] = {}
+        self.runs: dict[str, RunRecord] = {}
+        self.warehouses: dict[str, WarehouseRecord] = {}
+        self.secret_scopes: dict[str, SecretScope] = {}
+        self.dbfs: dict[str, DbfsNode] = {}
+        self.workspace_objects: dict[str, WorkspaceObject] = {}
+        self.permissions: dict[tuple[str, str], dict[str, Any]] = {}
+        self.statements: dict[str, dict[str, Any]] = {}
+        self.statement_callbacks: dict[str, Callable[..., Any]] = {}
+        self.user = {
+            "id": "1001",
+            "userName": "user@acme.example",
+            "displayName": "Acme User",
+            "active": True,
+        }
+
+    def bump(self) -> None:
+        self.revision += 1
+
+    def next_seq(self) -> int:
+        self.seq += 1
+        return self.seq
+
+    def allocate_numeric_id(self, existing: dict[str, Any], floor: int) -> str:
+        used = {int(key) for key in existing if str(key).isdigit()}
+        n = max([floor - 1, *used]) + 1
+        while str(n) in existing:
+            n += 1
+        self.seq = max(self.seq, n)
+        return str(n)
+
+    def reindex_seq(self) -> None:
+        nums = [int(k) for k in self.jobs if str(k).isdigit()]
+        nums.extend(int(k) for k in self.runs if str(k).isdigit())
+        if nums:
+            self.seq = max(self.seq, max(nums))
+
+    def clear(self) -> None:
+        self.clusters.clear()
+        self.jobs.clear()
+        self.runs.clear()
+        self.warehouses.clear()
+        self.secret_scopes.clear()
+        self.dbfs.clear()
+        self.workspace_objects.clear()
+        self.permissions.clear()
+        self.statements.clear()
+        self.statement_callbacks.clear()
+        self.revision = 0
+        self.seq = 0
+
+    def tick(self) -> None:
+        for cluster in self.clusters.values():
+            if cluster.deleted:
+                continue
+            if cluster.state in {"PENDING", "RESTARTING"}:
+                if cluster.ticks_remaining > 0:
+                    cluster.ticks_remaining -= 1
+                if cluster.ticks_remaining <= 0:
+                    if cluster.state == "RESTARTING":
+                        cluster.state = "PENDING"
+                        cluster.ticks_remaining = 1
+                        cluster.events.append(
+                            {"type": "PENDING", "message": "Restarting"}
+                        )
+                    elif cluster.fail_after:
+                        cluster.state = "ERROR"
+                        cluster.state_message = "Fixture startup failure"
+                        cluster.events.append(
+                            {"type": "ERROR", "message": cluster.state_message}
+                        )
+                    else:
+                        cluster.state = "RUNNING"
+                        cluster.state_message = ""
+                        cluster.events.append({"type": "RUNNING", "message": "Running"})
+            elif cluster.state == "TERMINATING":
+                cluster.state = "TERMINATED"
+                cluster.events.append({"type": "TERMINATED", "message": "Terminated"})
+        for run in self.runs.values():
+            if run.canceling:
+                if run.life_cycle_state == "TERMINATING":
+                    run.life_cycle_state = "TERMINATED"
+                    run.result_state = "CANCELED"
+                    run.canceling = False
+                else:
+                    run.life_cycle_state = "TERMINATING"
+                continue
+            if (
+                run.life_cycle_state in {"PENDING", "RUNNING"}
+                and run.ticks_remaining > 0
+            ):
+                run.ticks_remaining -= 1
+                if run.life_cycle_state == "PENDING":
+                    run.life_cycle_state = "RUNNING"
+                    run.ticks_remaining = max(run.ticks_remaining, 1)
+                elif run.ticks_remaining <= 0:
+                    run.life_cycle_state = "TERMINATED"
+                    run.result_state = run.result_state or "SUCCESS"
+
+    def maybe_tick_real(self, clock: str) -> None:
+        if clock == "real":
+            self.tick()
+
+    def add_cluster(
+        self,
+        cluster_name: str,
+        *,
+        cluster_id: str | None = None,
+        spark_version: str = "13.3.x-scala2.12",
+        node_type_id: str = "i3.xlarge",
+        state: str = "RUNNING",
+        ticks_remaining: int = 0,
+        fail_after: bool = False,
+        num_workers: int = 1,
+    ) -> ClusterRecord:
+        cid = cluster_id or make_id(
+            "cluster", f"{cluster_name}:{self.next_seq()}", self.seed
+        )
+        if cid in self.clusters:
+            raise ValueError(f"Duplicate cluster_id {cid}")
+        rec = ClusterRecord(
+            cluster_id=cid,
+            cluster_name=cluster_name,
+            spark_version=spark_version,
+            node_type_id=node_type_id,
+            state=state,
+            ticks_remaining=ticks_remaining,
+            fail_after=fail_after,
+            num_workers=num_workers,
+        )
+        rec.events.append({"type": state, "message": f"Cluster {state}"})
+        self.clusters[cid] = rec
+        self.permissions[("clusters", cid)] = {
+            "object_id": cid,
+            "object_type": "clusters",
+            "access_control_list": [],
+        }
+        self.bump()
+        return rec
+
+    def cluster_json(self, rec: ClusterRecord) -> dict[str, Any]:
+        return {
+            "cluster_id": rec.cluster_id,
+            "cluster_name": rec.cluster_name,
+            "spark_version": rec.spark_version,
+            "node_type_id": rec.node_type_id,
+            "state": rec.state,
+            "state_message": rec.state_message,
+            "num_workers": rec.num_workers,
+            "libraries": rec.libraries,
+        }
+
+    def job_json(self, rec: JobRecord) -> dict[str, Any]:
+        return {
+            "job_id": int(rec.job_id) if rec.job_id.isdigit() else rec.job_id,
+            "created_time": rec.created_time,
+            "creator_user_name": rec.creator_user_name,
+            "settings": rec.settings,
+        }
+
+    def run_json(self, rec: RunRecord) -> dict[str, Any]:
+        state: dict[str, Any] = {"life_cycle_state": rec.life_cycle_state}
+        if rec.result_state:
+            state["result_state"] = rec.result_state
+        body: dict[str, Any] = {
+            "run_id": int(rec.run_id) if rec.run_id.isdigit() else rec.run_id,
+            "run_name": rec.run_name,
+            "state": state,
+        }
+        if rec.job_id:
+            body["job_id"] = int(rec.job_id) if rec.job_id.isdigit() else rec.job_id
+        return body

semblance_databricks/testing.py

@@ -0,0 +1,66 @@
+"""Pytest / in-process helpers. Prefer TestClient; no live Databricks."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from pathlib import Path
+from typing import Any
+
+from fastapi.testclient import TestClient
+
+from semblance_databricks.app import DatabricksMock
+from semblance_databricks.config import DatabricksMockConfig
+from semblance_databricks.fixtures.loaders import bundled_acme_path
+
+
+class DatabricksMockContext:
+    """Load a fixture, yield a TestClient, reset state on exit."""
+
+    def __init__(
+        self,
+        *,
+        seed: int | None = 42,
+        fixture: str | Path | None = None,
+        config: DatabricksMockConfig | None = None,
+    ) -> None:
+        if config is None:
+            config = DatabricksMockConfig(seed=seed)
+        self.mock = DatabricksMock(config)
+        path = fixture if fixture is not None else bundled_acme_path()
+        self.mock.load_fixture(path)
+        self._client: TestClient | None = None
+
+    def __enter__(self) -> TestClient:
+        self._client = TestClient(self.mock.as_fastapi())
+        return self._client
+
+    def __exit__(self, *exc: Any) -> None:
+        self.mock.state.clear()
+        if self._client is not None:
+            self._client.close()
+
+
+def databricks_test_client(
+    mock: DatabricksMock | None = None,
+    *,
+    fixture: str | Path | None = None,
+    config: DatabricksMockConfig | None = None,
+) -> TestClient:
+    """Build a TestClient for a DatabricksMock (loads bundled acme if needed)."""
+    if mock is None:
+        mock = DatabricksMock(config or DatabricksMockConfig())
+        mock.load_fixture(fixture if fixture is not None else bundled_acme_path())
+    return TestClient(mock.as_fastapi())
+
+
+def acme_client_iter(
+    config: DatabricksMockConfig | None = None,
+) -> Iterator[tuple[DatabricksMock, TestClient]]:
+    mock = DatabricksMock(config or DatabricksMockConfig(seed=42))
+    mock.load_fixture(bundled_acme_path())
+    client = TestClient(mock.as_fastapi())
+    try:
+        yield mock, client
+    finally:
+        mock.state.clear()
+        client.close()