semblance-databricks 0.1.0__py3-none-any.whl

semblance_databricks/services/clusters/routes.py

@@ -0,0 +1,191 @@
+"""Clusters list/get/create/edit/delete/restart/events/libraries."""
+
+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, require_int
+from semblance_databricks.pagination import paginate
+from semblance_databricks.services.deps import (
+    mock_from,
+    require_fail_stage,
+    stub_unimplemented,
+)
+
+
+def create_clusters_router() -> APIRouter:
+    router = APIRouter()
+
+    def _list(
+        request: Request, page_size: int | None, page_token: str | None
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        mock.state.maybe_tick_real(mock.config.clock)
+        items = [
+            mock.state.cluster_json(c)
+            for c in sorted(mock.state.clusters.values(), key=lambda r: r.cluster_id)
+            if not c.deleted
+        ]
+        page, nxt = paginate(
+            items,
+            page_size=page_size,
+            page_token=page_token,
+            resource="clusters",
+            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] = {"clusters": page}
+        if nxt:
+            body["next_page_token"] = nxt
+        return body
+
+    def _get(request: Request, cluster_id: str | None) -> dict[str, Any]:
+        mock = mock_from(request)
+        mock.state.maybe_tick_real(mock.config.clock)
+        if not cluster_id:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "cluster_id required")
+        rec = mock.state.clusters.get(cluster_id)
+        if rec is None or rec.deleted:
+            raise DatabricksError(
+                404, "RESOURCE_DOES_NOT_EXIST", f"Cluster {cluster_id} not found"
+            )
+        return mock.state.cluster_json(rec)
+
+    @router.get("/api/2.1/clusters/list")
+    def list_21(
+        request: Request,
+        page_size: int | None = None,
+        page_token: str | None = None,
+    ) -> dict[str, Any]:
+        return _list(request, page_size, page_token)
+
+    @router.get("/api/2.1/clusters/get")
+    def get_21(request: Request, cluster_id: str | None = None) -> dict[str, Any]:
+        return _get(request, cluster_id)
+
+    @router.get("/api/2.0/clusters/get")
+    def get_20(request: Request, cluster_id: str | None = None) -> dict[str, Any]:
+        return _get(request, cluster_id)
+
+    @router.get("/api/2.0/clusters/list")
+    def list_20(
+        request: Request,
+        page_size: int | None = None,
+        page_token: str | None = None,
+    ) -> dict[str, Any]:
+        return _list(request, page_size, page_token)
+
+    @router.post("/api/2.1/clusters/create")
+    @router.post("/api/2.0/clusters/create")
+    async def create(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")
+        rec = mock.state.add_cluster(
+            str(body.get("cluster_name", "cluster")),
+            spark_version=str(body.get("spark_version", "13.3.x-scala2.12")),
+            node_type_id=str(body.get("node_type_id", "i3.xlarge")),
+            state="PENDING",
+            ticks_remaining=require_int(
+                body.get("startup_delay_ticks"), 1, "startup_delay_ticks"
+            ),
+            fail_after=bool(body.get("startup_failure", False)),
+            num_workers=require_int(body.get("num_workers"), 1, "num_workers"),
+        )
+        require_fail_stage(request, "after_write")
+        return {"cluster_id": rec.cluster_id}
+
+    @router.post("/api/2.1/clusters/edit")
+    @router.post("/api/2.0/clusters/edit")
+    async def edit(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        if "cluster_name" in body:
+            rec.cluster_name = str(body["cluster_name"])
+        if "num_workers" in body:
+            rec.num_workers = require_int(body["num_workers"], 1, "num_workers")
+        mock.state.bump()
+        return {}
+
+    @router.post("/api/2.1/clusters/delete")
+    @router.post("/api/2.0/clusters/delete")
+    async def delete(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        rec.state = "TERMINATING"
+        mock.state.bump()
+        return {}
+
+    @router.post("/api/2.1/clusters/restart")
+    @router.post("/api/2.0/clusters/restart")
+    async def restart(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        rec.state = "RESTARTING"
+        rec.ticks_remaining = 1
+        rec.fail_after = False
+        mock.state.bump()
+        return {}
+
+    @router.post("/api/2.1/clusters/permanent-delete")
+    @router.post("/api/2.0/clusters/permanent-delete")
+    async def permanent_delete(request: Request) -> dict[str, Any]:
+        stub_unimplemented(request, "clusters/permanent-delete")
+        return {}
+
+    @router.post("/api/2.1/clusters/events")
+    @router.post("/api/2.0/clusters/events")
+    async def events(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        return {"events": rec.events, "total_count": len(rec.events)}
+
+    @router.post("/api/2.0/libraries/install")
+    async def install(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        rec.libraries.extend(list(body.get("libraries") or []))
+        return {}
+
+    @router.post("/api/2.0/libraries/uninstall")
+    async def uninstall(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        cid = str(body.get("cluster_id", ""))
+        rec = mock.state.clusters.get(cid)
+        if rec is None or rec.deleted:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Cluster not found")
+        requested = list(body.get("libraries") or [])
+        if requested:
+            rec.libraries = [lib for lib in rec.libraries if lib not in requested]
+        else:
+            rec.libraries = []
+        return {}
+
+    return router

semblance_databricks/services/dbfs/__init__.py

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

semblance_databricks/services/dbfs/routes.py

@@ -0,0 +1,126 @@
+"""DBFS list/read, put, and add-block stub."""
+
+from __future__ import annotations
+
+import base64
+from pathlib import Path
+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, stub_unimplemented
+from semblance_databricks.state import DbfsNode
+
+
+def _normalize(path: str) -> str:
+    raw = path.replace("\\", "/")
+    if not raw.startswith("/"):
+        raw = "/" + raw
+    parts: list[str] = []
+    for seg in raw.split("/"):
+        if seg in {"", "."}:
+            continue
+        if seg == ".." or ":" in seg:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "Invalid path")
+        parts.append(seg)
+    return "/" + "/".join(parts) if parts else "/"
+
+
+def _parent(path: str) -> str:
+    if path in {"", "/"}:
+        return "/"
+    stripped = path.rstrip("/")
+    idx = stripped.rfind("/")
+    return stripped[:idx] or "/"
+
+
+def _write_temp(root: str, dbfs_path: str, data: bytes) -> None:
+    base = Path(root).resolve()
+    rel = dbfs_path.lstrip("/")
+    dest = base.joinpath(*rel.split("/")).resolve()
+    try:
+        dest.relative_to(base)
+    except ValueError as exc:
+        raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "Invalid path") from exc
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    dest.write_bytes(data)
+
+
+def create_dbfs_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.get("/api/2.0/dbfs/list")
+    def list_dbfs(request: Request, path: str | None = None) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not path:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "path required")
+        prefix = _normalize(path)
+        files: list[dict[str, Any]] = []
+        for node in mock.state.dbfs.values():
+            parent = _parent(node.path)
+            if parent == prefix:
+                if node.path == prefix:
+                    continue
+                files.append(
+                    {
+                        "path": node.path,
+                        "is_dir": node.is_dir,
+                        "file_size": 0 if node.is_dir else len(node.data),
+                    }
+                )
+        if prefix not in mock.state.dbfs and not files:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Path not found")
+        return {"files": files}
+
+    @router.get("/api/2.0/dbfs/read")
+    def read_dbfs(
+        request: Request,
+        path: str | None = None,
+        offset: int = 0,
+        length: int = 1048576,
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        if not path:
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "path required")
+        node = mock.state.dbfs.get(_normalize(path))
+        if node is None or node.is_dir:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "File not found")
+        data = node.data[offset : offset + length]
+        return {
+            "bytes_read": len(data),
+            "data": base64.b64encode(data).decode(),
+        }
+
+    @router.post("/api/2.0/dbfs/add-block")
+    async def add_block(request: Request) -> dict[str, Any]:
+        stub_unimplemented(request, "dbfs/add-block")
+        return {}
+
+    @router.post("/api/2.0/dbfs/put")
+    async def put_file(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        path = _normalize(str(body.get("path", "")))
+        if not path or path == "/":
+            raise DatabricksError(400, "INVALID_PARAMETER_VALUE", "path required")
+        if len(mock.state.dbfs) >= mock.config.dbfs_max_files:
+            raise DatabricksError(400, "MAX_BLOCK_SIZE_EXCEEDED", "DBFS file cap")
+        raw = str(body.get("contents", ""))
+        try:
+            content = base64.b64decode(raw, validate=True)
+        except Exception as exc:
+            raise DatabricksError(
+                400, "INVALID_PARAMETER_VALUE", "contents must be base64"
+            ) from exc
+        total = sum(len(n.data) for n in mock.state.dbfs.values()) + len(content)
+        if total > mock.config.dbfs_max_bytes:
+            raise DatabricksError(400, "MAX_BLOCK_SIZE_EXCEEDED", "DBFS size cap")
+        mock.state.dbfs[path] = DbfsNode(path=path, is_dir=False, data=content)
+        if mock.config.dbfs_temp_dir:
+            _write_temp(mock.config.dbfs_temp_dir, path, content)
+        mock.state.bump()
+        return {}
+
+    return router

semblance_databricks/services/dbsql/__init__.py

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

semblance_databricks/services/dbsql/routes.py

@@ -0,0 +1,163 @@
+"""SQL warehouses and statement execution (no Spark)."""
+
+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.ids import make_id
+from semblance_databricks.pagination import paginate
+from semblance_databricks.services.deps import mock_from
+from semblance_databricks.state import WarehouseRecord
+
+
+def _statement_payload(statement_id: str, rec: dict[str, Any]) -> dict[str, Any]:
+    chunks: list[Any] = rec.get("chunks") or []
+    index = int(rec.get("chunk_index", 0))
+    chunk = chunks[index] if index < len(chunks) else []
+    next_chunk = None
+    if index + 1 < len(chunks):
+        next_chunk = index + 1
+    result: dict[str, Any] = {
+        "chunk_index": index,
+        "row_count": len(chunk) if isinstance(chunk, list) else 0,
+        "data_array": chunk,
+    }
+    if next_chunk is not None:
+        result["next_chunk_internal"] = next_chunk
+    return {
+        "statement_id": statement_id,
+        "status": {"state": rec.get("state", "SUCCEEDED")},
+        "manifest": {
+            "format": "JSON_ARRAY",
+            "total_chunk_count": len(chunks),
+            "total_row_count": sum(
+                len(c) if isinstance(c, list) else 0 for c in chunks
+            ),
+        },
+        "result": result,
+    }
+
+
+def create_dbsql_router() -> APIRouter:
+    router = APIRouter()
+
+    @router.get("/api/2.0/sql/warehouses")
+    def list_warehouses(
+        request: Request,
+        page_size: int | None = None,
+        page_token: str | None = None,
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        items = [
+            {
+                "id": w.id,
+                "name": w.name,
+                "state": w.state,
+                "cluster_size": w.cluster_size,
+            }
+            for w in mock.state.warehouses.values()
+        ]
+        page, nxt = paginate(
+            items,
+            page_size=page_size,
+            page_token=page_token,
+            resource="warehouses",
+            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] = {"warehouses": page}
+        if nxt:
+            body["next_page_token"] = nxt
+        return body
+
+    @router.post("/api/2.0/sql/warehouses")
+    async def create_warehouse(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        name = str(body.get("name", "warehouse"))
+        wid = make_id("warehouse", f"{name}:{mock.state.next_seq()}", mock.config.seed)
+        rec = WarehouseRecord(
+            id=wid,
+            name=name,
+            cluster_size=str(body.get("cluster_size", "X-Small")),
+        )
+        mock.state.warehouses[wid] = rec
+        mock.state.permissions[("warehouses", wid)] = {
+            "object_id": wid,
+            "object_type": "warehouses",
+            "access_control_list": [],
+        }
+        mock.state.bump()
+        return {"id": wid}
+
+    @router.get("/api/2.0/sql/warehouses/{warehouse_id}")
+    def get_warehouse(request: Request, warehouse_id: str) -> dict[str, Any]:
+        mock = mock_from(request)
+        rec = mock.state.warehouses.get(warehouse_id)
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Warehouse not found")
+        return {
+            "id": rec.id,
+            "name": rec.name,
+            "state": rec.state,
+            "cluster_size": rec.cluster_size,
+        }
+
+    @router.delete("/api/2.0/sql/warehouses/{warehouse_id}")
+    def delete_warehouse(request: Request, warehouse_id: str) -> dict[str, Any]:
+        mock = mock_from(request)
+        if warehouse_id not in mock.state.warehouses:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Warehouse not found")
+        del mock.state.warehouses[warehouse_id]
+        mock.state.bump()
+        return {}
+
+    @router.post("/api/2.0/sql/statements")
+    async def execute_statement(request: Request) -> dict[str, Any]:
+        mock = mock_from(request)
+        body = await json_object(request)
+        warehouse_id = str(body.get("warehouse_id") or "")
+        sql = str(body.get("statement", ""))
+        if not warehouse_id:
+            raise DatabricksError(
+                400, "INVALID_PARAMETER_VALUE", "warehouse_id required"
+            )
+        if warehouse_id not in mock.state.warehouses:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Warehouse not found")
+        sid = make_id("statement", f"{sql}:{mock.state.next_seq()}", mock.config.seed)
+        callback = mock.state.statement_callbacks.get(
+            sql
+        ) or mock.state.statement_callbacks.get(warehouse_id)
+        if callback is not None:
+            payload = callback(body, mock.state)
+            chunks = payload.get("chunks", payload.get("data_array", [["1"]]))
+            if chunks and not isinstance(chunks[0], list):
+                chunks = [chunks]
+        else:
+            chunks = [[["ok"]]]
+        rec = {"state": "SUCCEEDED", "chunks": chunks, "chunk_index": 0}
+        mock.state.statements[sid] = rec
+        mock.state.bump()
+        return _statement_payload(sid, rec)
+
+    @router.get("/api/2.0/sql/statements/{statement_id}")
+    def get_statement(
+        request: Request,
+        statement_id: str,
+        chunk_index: int | None = None,
+    ) -> dict[str, Any]:
+        mock = mock_from(request)
+        rec = mock.state.statements.get(statement_id)
+        if rec is None:
+            raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", "Statement not found")
+        if chunk_index is not None:
+            rec = {**rec, "chunk_index": chunk_index}
+        return _statement_payload(statement_id, rec)
+
+    return router

semblance_databricks/services/deps.py

@@ -0,0 +1,37 @@
+"""Request helpers."""
+
+from __future__ import annotations
+
+from typing import Protocol, cast
+
+from fastapi import Request
+
+from semblance_databricks.config import DatabricksMockConfig
+from semblance_databricks.errors import DatabricksError
+from semblance_databricks.ids import PageTokenCodec
+from semblance_databricks.state import DatabricksState
+
+
+class DatabricksMockLike(Protocol):
+    config: DatabricksMockConfig
+    state: DatabricksState
+    page_token_codec: PageTokenCodec
+
+    def tick(self) -> None: ...
+
+
+def mock_from(request: Request) -> DatabricksMockLike:
+    return cast(DatabricksMockLike, request.app.state.databricks_mock)
+
+
+def require_fail_stage(request: Request, stage: str) -> None:
+    mock = mock_from(request)
+    if mock.config.fail_stage == stage:
+        raise DatabricksError(500, "INTERNAL_ERROR", f"fail_stage={stage}")
+
+
+def stub_unimplemented(request: Request, name: str) -> None:
+    mock = mock_from(request)
+    if mock.config.auth == "strict":
+        raise DatabricksError(501, "FEATURE_DISABLED", f"{name} is not implemented")
+    raise DatabricksError(404, "RESOURCE_DOES_NOT_EXIST", f"{name} not found")

semblance_databricks/services/jobs/__init__.py

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