semblance-databricks 0.1.0__py3-none-any.whl

semblance_databricks/config.py

@@ -0,0 +1,52 @@
+"""Typed configuration for DatabricksMock."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+AuthMode = Literal["disabled", "optional", "strict"]
+ClockMode = Literal["virtual", "real"]
+FailStage = Literal["before_validate", "before_write", "after_write"] | None
+
+
+@dataclass(frozen=True)
+class TokenGrant:
+    """A token accepted in strict auth mode. Never log or echo ``token``."""
+
+    token: str
+    scopes: tuple[str, ...] = ()
+
+
+@dataclass
+class DatabricksMockConfig:
+    """Runtime options for a Databricks mock server."""
+
+    seed: int | None = 42
+    auth: AuthMode = "optional"
+    tokens: tuple[TokenGrant, ...] = ()
+    error_rate: float = 0.0
+    error_codes: list[int] | None = None
+    latency_ms: float = 0.0
+    jitter_ms: float = 0.0
+    rate_limit: float | None = None
+    default_page_size: int = 20
+    max_page_size: int = 100
+    clock: ClockMode = "virtual"
+    org_id: str | None = "1234567890"
+    fail_stage: FailStage = None
+    dbfs_temp_dir: str | None = None
+    dbfs_max_files: int = 100
+    dbfs_max_bytes: int = 1_048_576
+
+    def __post_init__(self) -> None:
+        if self.error_codes is None:
+            self.error_codes = [500]
+        if self.auth not in ("disabled", "optional", "strict"):
+            raise ValueError(
+                f"auth must be 'disabled', 'optional', or 'strict', got {self.auth!r}"
+            )
+        if self.clock not in ("virtual", "real"):
+            raise ValueError("clock must be 'virtual' or 'real'")
+        if not 0.0 <= self.error_rate <= 1.0:
+            raise ValueError("error_rate must be between 0 and 1")

semblance_databricks/errors.py

@@ -0,0 +1,72 @@
+"""Databricks error_code/message mapper."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import Request
+from fastapi.exceptions import RequestValidationError
+from fastapi.responses import JSONResponse
+from starlette.exceptions import HTTPException as StarletteHTTPException
+
+
+class DatabricksError(Exception):
+    """Raised to produce a Databricks-style error response."""
+
+    def __init__(self, status_code: int, error_code: str, message: str) -> None:
+        self.status_code = status_code
+        self.error_code = error_code
+        self.message = message
+        super().__init__(message)
+
+    def to_body(self) -> dict[str, str]:
+        return {"error_code": self.error_code, "message": self.message}
+
+
+def databricks_error_body(error_code: str, message: str) -> dict[str, str]:
+    return {"error_code": error_code, "message": message}
+
+
+def register_exception_handlers(app: Any) -> None:
+    @app.exception_handler(DatabricksError)
+    async def databricks_error_handler(
+        request: Request, exc: DatabricksError
+    ) -> JSONResponse:
+        return JSONResponse(status_code=exc.status_code, content=exc.to_body())
+
+    @app.exception_handler(RequestValidationError)
+    async def validation_handler(
+        request: Request, exc: RequestValidationError
+    ) -> JSONResponse:
+        return JSONResponse(
+            status_code=400,
+            content=databricks_error_body(
+                "INVALID_PARAMETER_VALUE",
+                "Invalid request",
+            ),
+        )
+
+    @app.exception_handler(StarletteHTTPException)
+    async def http_exception_handler(
+        request: Request, exc: StarletteHTTPException
+    ) -> JSONResponse:
+        if exc.status_code == 404:
+            return JSONResponse(
+                status_code=404,
+                content=databricks_error_body(
+                    "RESOURCE_DOES_NOT_EXIST",
+                    f"Path {request.url.path} does not exist",
+                ),
+            )
+        if exc.status_code == 405:
+            return JSONResponse(
+                status_code=405,
+                content=databricks_error_body(
+                    "INVALID_PARAMETER_VALUE",
+                    "Method not allowed",
+                ),
+            )
+        return JSONResponse(
+            status_code=exc.status_code,
+            content=databricks_error_body("UNKNOWN", "HTTP error"),
+        )

semblance_databricks/fixtures/__init__.py

@@ -0,0 +1,5 @@
+"""Fixture package."""
+
+from semblance_databricks.fixtures.loaders import bundled_acme_path, load_fixture_file
+
+__all__ = ["bundled_acme_path", "load_fixture_file"]

semblance_databricks/fixtures/defaults/acme.yaml

@@ -0,0 +1,47 @@
+version: 1
+workspace:
+  workspaceId: "1234567890"
+  name: acme
+clusters:
+  - cluster_id: "0101-acme-0001"
+    cluster_name: ingest-1
+    state: RUNNING
+  - cluster_id: "0101-acme-0002"
+    cluster_name: ingest-2
+    state: RUNNING
+  - cluster_id: "0101-acme-0003"
+    cluster_name: ingest-3
+    state: TERMINATED
+  - cluster_id: "0101-acme-0004"
+    cluster_name: ingest-4
+    state: PENDING
+    startup_delay_ticks: 1
+jobs:
+  - job_id: "1001"
+    name: nightly-etl
+  - job_id: "1002"
+    name: weekly-report
+runs:
+  - run_id: "2001"
+    job_id: "1001"
+    run_name: nightly-etl-run
+    life_cycle_state: TERMINATED
+    result_state: SUCCESS
+    output_notebook: "ok"
+warehouses:
+  - id: "wh-acme"
+    name: analytics
+    state: RUNNING
+secretScopes:
+  - name: acme-scope
+    keys:
+      - api-key
+dbfs:
+  - path: "/mnt/acme"
+    is_dir: true
+  - path: "/mnt/acme/readme.txt"
+    is_dir: false
+    content: hello
+workspaceObjects:
+  - path: "/Users/user@acme.example/notebook"
+    object_type: NOTEBOOK

semblance_databricks/fixtures/loaders.py

@@ -0,0 +1,196 @@
+"""Fixture v1 schema and loaders. YAML/JSON is data only — never evaluated."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+
+import yaml
+from pydantic import BaseModel, ConfigDict, Field
+
+from semblance_databricks.state import (
+    DatabricksState,
+    DbfsNode,
+    JobRecord,
+    RunRecord,
+    SecretScope,
+    WarehouseRecord,
+    WorkspaceObject,
+)
+
+
+class _Strict(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class FixtureWorkspace(_Strict):
+    workspace_id: str = Field(alias="workspaceId", default="1234567890")
+    name: str = "acme"
+
+
+class FixtureCluster(_Strict):
+    cluster_id: str | None = Field(default=None, alias="cluster_id")
+    cluster_name: str
+    spark_version: str = "13.3.x-scala2.12"
+    node_type_id: str = "i3.xlarge"
+    state: str = "RUNNING"
+    num_workers: int = 1
+    startup_delay_ticks: int = 0
+    startup_failure: bool = False
+
+
+class FixtureJob(_Strict):
+    job_id: str | None = None
+    name: str
+    creator_user_name: str = "user@acme.example"
+
+
+class FixtureRun(_Strict):
+    run_id: str | None = None
+    job_id: str | None = None
+    run_name: str = "nightly"
+    life_cycle_state: str = "TERMINATED"
+    result_state: str | None = "SUCCESS"
+    output_notebook: str | None = None
+
+
+class FixtureWarehouse(_Strict):
+    id: str | None = None
+    name: str
+    state: str = "RUNNING"
+
+
+class FixtureSecretScope(_Strict):
+    name: str
+    keys: list[str] = Field(default_factory=list)
+
+
+class FixtureDbfs(_Strict):
+    path: str
+    is_dir: bool = False
+    content: str = ""
+
+
+class FixtureWorkspaceObject(_Strict):
+    path: str
+    object_type: str = "NOTEBOOK"
+
+
+class FixtureDocument(_Strict):
+    version: int
+    workspace: FixtureWorkspace = Field(default_factory=FixtureWorkspace)
+    clusters: list[FixtureCluster] = Field(default_factory=list)
+    jobs: list[FixtureJob] = Field(default_factory=list)
+    runs: list[FixtureRun] = Field(default_factory=list)
+    warehouses: list[FixtureWarehouse] = Field(default_factory=list)
+    secret_scopes: list[FixtureSecretScope] = Field(
+        default_factory=list, alias="secretScopes"
+    )
+    dbfs: list[FixtureDbfs] = Field(default_factory=list)
+    workspace_objects: list[FixtureWorkspaceObject] = Field(
+        default_factory=list, alias="workspaceObjects"
+    )
+
+
+def parse_fixture(data: dict[str, Any]) -> FixtureDocument:
+    doc = FixtureDocument.model_validate(data)
+    if doc.version != 1:
+        raise ValueError("only fixture version 1 is supported")
+    return doc
+
+
+def load_fixture_file(path: str | Path) -> FixtureDocument:
+    raw = Path(path).read_text(encoding="utf-8")
+    if Path(path).suffix.lower() in {".yaml", ".yml"}:
+        loaded = yaml.safe_load(raw)
+    else:
+        import json
+
+        loaded = json.loads(raw)
+    if not isinstance(loaded, dict):
+        raise ValueError("fixture root must be a mapping")
+    return parse_fixture(loaded)
+
+
+def apply_fixture(doc: FixtureDocument, state: DatabricksState) -> None:
+    seen_jobs: set[str] = set()
+    for cluster in doc.clusters:
+        rec = state.add_cluster(
+            cluster.cluster_name,
+            cluster_id=cluster.cluster_id,
+            spark_version=cluster.spark_version,
+            node_type_id=cluster.node_type_id,
+            state=cluster.state,
+            ticks_remaining=cluster.startup_delay_ticks,
+            fail_after=cluster.startup_failure,
+            num_workers=cluster.num_workers,
+        )
+        _ = rec
+    for job in doc.jobs:
+        jid = job.job_id or str(1000 + len(state.jobs) + 1)
+        if jid in seen_jobs or jid in state.jobs:
+            raise ValueError(f"Duplicate job_id {jid}")
+        seen_jobs.add(jid)
+        state.jobs[jid] = JobRecord(
+            job_id=jid,
+            settings={"name": job.name},
+            creator_user_name=job.creator_user_name,
+        )
+        state.permissions[("jobs", jid)] = {
+            "object_id": jid,
+            "object_type": "job",
+            "access_control_list": [
+                {
+                    "user_name": job.creator_user_name,
+                    "all_permissions": [{"permission_level": "CAN_MANAGE"}],
+                }
+            ],
+        }
+    for run in doc.runs:
+        rid = run.run_id or str(2000 + len(state.runs) + 1)
+        if rid in state.runs:
+            raise ValueError(f"Duplicate run_id {rid}")
+        if run.job_id and run.job_id not in state.jobs:
+            raise ValueError(f"Dangling job_id {run.job_id} on run")
+        state.runs[rid] = RunRecord(
+            run_id=rid,
+            job_id=run.job_id,
+            run_name=run.run_name,
+            life_cycle_state=run.life_cycle_state,
+            result_state=run.result_state,
+            output={"notebook_output": {"result": run.output_notebook or "ok"}},
+        )
+    for wh in doc.warehouses:
+        wid = wh.id or f"wh-{len(state.warehouses) + 1}"
+        if wid in state.warehouses:
+            raise ValueError(f"Duplicate warehouse id {wid}")
+        state.warehouses[wid] = WarehouseRecord(id=wid, name=wh.name, state=wh.state)
+    for scope in doc.secret_scopes:
+        if scope.name in state.secret_scopes:
+            raise ValueError(f"Duplicate secret scope {scope.name}")
+        state.secret_scopes[scope.name] = SecretScope(
+            name=scope.name,
+            keys={k: "redacted" for k in scope.keys},
+        )
+    for node in doc.dbfs:
+        state.dbfs[node.path] = DbfsNode(
+            path=node.path,
+            is_dir=node.is_dir,
+            data=node.content.encode(),
+        )
+    for obj in doc.workspace_objects:
+        state.workspace_objects[obj.path] = WorkspaceObject(
+            path=obj.path, object_type=obj.object_type
+        )
+    state.reindex_seq()
+    state.bump()
+
+
+def bundled_acme_path() -> Path:
+    return Path(__file__).resolve().parent / "defaults" / "acme.yaml"
+
+
+def load_bundled_fixture(name: str = "acme") -> FixtureDocument:
+    if name != "acme":
+        raise ValueError(f"unknown bundled fixture {name!r}; expected 'acme'")
+    return load_fixture_file(bundled_acme_path())

semblance_databricks/http.py

@@ -0,0 +1,34 @@
+"""Shared request parsing for Databricks REST handlers."""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from fastapi import Request
+
+from semblance_databricks.errors import DatabricksError
+
+
+async def json_object(request: Request) -> dict[str, Any]:
+    raw = await request.body()
+    if not raw:
+        return {}
+    try:
+        body = json.loads(raw)
+    except json.JSONDecodeError as exc:
+        raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "Invalid JSON") from exc
+    if not isinstance(body, dict):
+        raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "JSON object required")
+    return body
+
+
+def require_int(value: Any, default: int, name: str) -> int:
+    if value is None:
+        return default
+    try:
+        return int(value)
+    except (TypeError, ValueError) as exc:
+        raise DatabricksError(
+            400, "INVALID_PARAMETER_VALUE", f"{name} must be an integer"
+        ) from exc

semblance_databricks/ids.py

@@ -0,0 +1,88 @@
+"""Deterministic resource IDs and opaque page tokens."""
+
+from __future__ import annotations
+
+import base64
+import hashlib
+import hmac
+import json
+from dataclasses import dataclass
+
+from semblance_databricks.errors import DatabricksError
+
+
+def token_secret(seed: int | None) -> bytes:
+    material = f"semblance-databricks:{seed if seed is not None else 'none'}"
+    return hashlib.sha256(material.encode()).digest()
+
+
+def make_id(kind: str, identity: str, seed: int | None = 42) -> str:
+    digest = hashlib.sha256(
+        f"semblance-databricks:{seed}:{kind}:{identity}".encode()
+    ).hexdigest()[:12]
+    if kind == "cluster":
+        return f"0101-{digest[:8]}-{digest[8:]}"
+    if kind == "job":
+        return str(int(digest[:8], 16) % 10_000_000 + 1000)
+    if kind == "run":
+        return str(int(digest[:8], 16) % 10_000_000 + 2000)
+    if kind == "warehouse":
+        return digest
+    return digest
+
+
+@dataclass(frozen=True)
+class PageCursor:
+    resource: str
+    offset: int
+    revision: int
+
+
+class PageTokenCodec:
+    """URL-safe checksummed page tokens. Opaque to clients."""
+
+    def __init__(self, secret: bytes) -> None:
+        self._secret = secret
+
+    def encode(self, resource: str, offset: int, revision: int) -> str:
+        payload = json.dumps(
+            {"r": resource, "o": offset, "v": revision},
+            separators=(",", ":"),
+        )
+        sig = hmac.new(self._secret, payload.encode(), hashlib.sha256).hexdigest()[:16]
+        raw = f"{payload}|{sig}".encode()
+        return base64.urlsafe_b64encode(raw).decode().rstrip("=")
+
+    def decode(self, token: str, expected_resource: str) -> PageCursor:
+        try:
+            padded = token + "=" * (-len(token) % 4)
+            decoded = base64.urlsafe_b64decode(padded.encode()).decode()
+            payload, sig = decoded.rsplit("|", 1)
+            expected = hmac.new(
+                self._secret, payload.encode(), hashlib.sha256
+            ).hexdigest()[:16]
+            if len(sig) != len(expected) or not hmac.compare_digest(sig, expected):
+                raise ValueError("checksum mismatch")
+            data = json.loads(payload)
+            resource = str(data["r"])
+            offset = int(data["o"])
+            revision = int(data["v"])
+        except (
+            ValueError,
+            KeyError,
+            TypeError,
+            json.JSONDecodeError,
+            UnicodeDecodeError,
+        ) as exc:
+            raise DatabricksError(
+                400,
+                "INVALID_PARAMETER_VALUE",
+                "Invalid page token",
+            ) from exc
+        if resource != expected_resource or offset < 0:
+            raise DatabricksError(
+                400,
+                "INVALID_PARAMETER_VALUE",
+                "Invalid page token",
+            )
+        return PageCursor(resource=resource, offset=offset, revision=revision)

semblance_databricks/models/__init__.py

@@ -0,0 +1 @@
+"""Shared models."""

semblance_databricks/pagination.py

@@ -0,0 +1,48 @@
+"""Pagination helpers using PageTokenCodec."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import TypeVar
+
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.ids import PageTokenCodec
+
+T = TypeVar("T")
+
+
+def paginate(
+    items: Sequence[T],
+    *,
+    page_size: int | None,
+    page_token: str | None,
+    resource: str,
+    codec: PageTokenCodec,
+    revision: int,
+    default_page_size: int = 20,
+    max_page_size: int = 100,
+) -> tuple[list[T], str | None]:
+    size = default_page_size if page_size is None else page_size
+    if size <= 0:
+        raise DatabricksError(
+            400,
+            "INVALID_PARAMETER_VALUE",
+            "page_size must be positive",
+        )
+    size = min(size, max_page_size)
+    offset = 0
+    if page_token:
+        cursor = codec.decode(page_token, resource)
+        if cursor.revision != revision:
+            raise DatabricksError(
+                400,
+                "INVALID_PARAMETER_VALUE",
+                "Invalid page token",
+            )
+        offset = cursor.offset
+    page = list(items[offset : offset + size])
+    next_offset = offset + len(page)
+    next_token = None
+    if next_offset < len(items):
+        next_token = codec.encode(resource, next_offset, revision)
+    return page, next_token

semblance_databricks/registry.py

@@ -0,0 +1,29 @@
+"""Named Databricks operations for CLI and compatibility lookup."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from semblance_databricks.compatibility import operations_table
+
+
+@dataclass(frozen=True)
+class DatabricksOperation:
+    operation_id: str
+    method: str
+    path: str
+    support_level: str
+
+
+def registered_operations() -> list[DatabricksOperation]:
+    ops: list[DatabricksOperation] = []
+    for row in operations_table():
+        ops.append(
+            DatabricksOperation(
+                operation_id=str(row.get("operationId", "")),
+                method=str(row.get("method", "")),
+                path=str(row.get("path", "")),
+                support_level=str(row.get("supportLevel", "")),
+            )
+        )
+    return ops

semblance_databricks/services/__init__.py

@@ -0,0 +1 @@
+"""Databricks service routers."""

semblance_databricks/services/clusters/__init__.py

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