databricks-explore 0.5.0__py3-none-any.whl

databricks_explore/canvases.py

@@ -0,0 +1,333 @@
+"""Canvas files under data/canvases/<name>.json, format versions 1 and 2.
+
+The pydantic models mirror contracts/canvas.schema.json. Checks the schema cannot express
+(unique window titles and ids, links and connectors that point at existing ids) run in check_references().
+Validated bodies are stored exactly as received, so save followed by load returns the same object.
+Version 2 adds the sketch node types and a top-level `edges` list of user-drawn connectors.
+Each canvas also has a revision number (a counter in data/canvas_history/<name>/meta.json, never inside the canvas file)
+that every save bumps, and agent writes keep the replaced file as a snapshot so they can be undone.
+"""
+import json
+import os
+import threading
+from typing import Annotated, Any, Literal, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, StrictBool, StrictFloat, StrictInt, ValidationError, model_validator
+from pydantic_core import PydanticCustomError
+
+from . import dbx
+
+CANVASES = dbx.DATA / "canvases"
+CANVASES.mkdir(parents=True, exist_ok=True)
+
+NAME_PATTERN = r"^[A-Za-z0-9_.-]+$"
+KEEP_SNAPSHOTS = 20  # per canvas
+LOCK = threading.RLock()  # one writer at a time, so an agent's read-modify-write cannot interleave with a save
+
+
+class CanvasInvalid(ValueError):
+    """Validation failed. str(e) is '<path>: <reason>' (route returns it as HTTP 400)."""
+
+
+class CanvasNotFound(LookupError):
+    """No stored canvas with that name (route returns HTTP 404)."""
+
+
+class CanvasConflict(RuntimeError):
+    """The canvas was saved by someone else since the revision the writer started from (route returns HTTP 409).
+
+    `info` is what revision_info() says now, so the page can offer Reload or Keep mine."""
+
+    def __init__(self, name: str, expected: int, info: dict):
+        super().__init__(f"canvas {name} was saved by {info['by'] or 'someone else'} since revision {expected} (now revision {info['revision']})")
+        self.info = info
+
+
+class Point(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    x: StrictFloat
+    y: StrictFloat
+
+
+class Size(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    w: Annotated[StrictFloat, Field(gt=0)]
+    h: Annotated[StrictFloat, Field(gt=0)]
+
+
+class Viewport(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    x: StrictFloat
+    y: StrictFloat
+    zoom: Annotated[StrictFloat, Field(gt=0)]
+
+
+class Chart(BaseModel):
+    vega_lite: dict[str, Any]
+
+
+WindowType = Literal["query", "chart", "note", "table", "shape", "text", "sticky", "frame", "placeholder"]
+REQUIRED_BY_TYPE = {"query": "sql", "chart": "chart", "note": "text", "table": "table",
+                    "shape": "shape", "placeholder": "placeholder"}
+SKETCH_TYPES = ("shape", "text", "sticky", "frame", "placeholder")  # version 2 only; titled by `label`, not `title`
+ShapeKind = Literal["rect", "rounded", "ellipse", "diamond", "cylinder"]
+LineKind = Literal["straight", "step", "curve"]
+
+
+class PlaceholderColumn(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    name: str
+    type: str = ""
+    key: Optional[Literal["pk", "fk"]] = None
+
+
+class Placeholder(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    id: Annotated[str, Field(pattern=r"^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$")]
+    name: str
+    columns: list[PlaceholderColumn] = Field(default_factory=list)
+    linked: Optional[Annotated[str, Field(min_length=1)]] = None
+
+
+class Window(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    id: Annotated[str, Field(min_length=1)]
+    type: WindowType
+    title: Optional[Annotated[str, Field(min_length=1, pattern=r"^[^{}]+$")]] = None
+    position: Point
+    size: Size
+    sql: Optional[str] = None
+    chart: Optional[Chart] = None
+    text: Optional[str] = None
+    table: Optional[Annotated[str, Field(min_length=1)]] = None
+    collapsed: StrictBool = False
+    label: Optional[str] = None
+    shape: Optional[ShapeKind] = None
+    fill: Optional[str] = None
+    stroke: Optional[str] = None
+    font_size: Optional[Annotated[StrictFloat, Field(gt=0)]] = None
+    locked: Optional[StrictBool] = None
+    z: Optional[StrictFloat] = None
+    placeholder: Optional[Placeholder] = None
+
+    @model_validator(mode="after")
+    def _type_needs_its_field(self):
+        if self.type not in SKETCH_TYPES and self.title is None:
+            raise PydanticCustomError("type_requires_field", 'type "{type}" requires "{field}"',
+                                      {"type": self.type, "field": "title"})
+        field = REQUIRED_BY_TYPE.get(self.type)
+        if field is not None and getattr(self, field) is None:
+            raise PydanticCustomError("type_requires_field", 'type "{type}" requires "{field}"',
+                                      {"type": self.type, "field": field})
+        return self
+
+
+class Link(BaseModel):
+    model_config = ConfigDict(extra="forbid", populate_by_name=True)
+    from_: Annotated[str, Field(alias="from", min_length=1)]
+    to: Annotated[str, Field(min_length=1)]
+
+
+class RelationshipRef(BaseModel):
+    """What a connector remembers about the `sketch` edge saved from it (all optional)."""
+    model_config = ConfigDict(extra="forbid")
+    id: Optional[Annotated[str, Field(min_length=1)]] = None
+    source_column: Optional[Annotated[str, Field(min_length=1)]] = None
+    target_column: Optional[Annotated[str, Field(min_length=1)]] = None
+
+
+class CanvasEdge(BaseModel):
+    """A user-drawn connector between two nodes. Version 2 only."""
+    model_config = ConfigDict(extra="forbid")
+    id: Annotated[str, Field(min_length=1)]
+    source: Annotated[str, Field(min_length=1)]
+    target: Annotated[str, Field(min_length=1)]
+    source_handle: Optional[Literal["t", "r", "b", "l"]] = None
+    target_handle: Optional[Literal["t", "r", "b", "l"]] = None
+    line: LineKind = "step"
+    arrow_start: StrictBool = False
+    arrow_end: StrictBool = True
+    dashed: StrictBool = False
+    label: Optional[str] = None
+    label_offset: StrictFloat = 0.0
+    color: Optional[str] = None
+    relationship: Optional[RelationshipRef] = None
+
+
+class Canvas(BaseModel):
+    model_config = ConfigDict(extra="allow")
+    version: Annotated[StrictInt, Field(ge=1, le=2)]
+    name: Annotated[str, Field(min_length=1, pattern=NAME_PATTERN)]
+    profile: Optional[Annotated[str, Field(min_length=1)]] = None
+    warehouse_id: Optional[str] = None
+    viewport: Optional[Viewport] = None
+    windows: list[Window]
+    links: list[Link] = Field(default_factory=list)
+    edges: list[CanvasEdge] = Field(default_factory=list)
+
+
+def _path(loc):
+    out = ""
+    for part in loc:
+        out += f"[{part}]" if isinstance(part, int) else (f".{part}" if out else str(part))
+    return out or "(root)"
+
+
+def check_references(canvas: Canvas) -> None:
+    """Checks the schema cannot express: unique titles and ids, and links and connectors that name existing ids."""
+    if canvas.version == 1:
+        if canvas.profile is None:
+            raise CanvasInvalid('profile: required in version 1')
+        if "edges" in canvas.model_fields_set:
+            raise CanvasInvalid('edges: needs version 2')
+    ids, titles = {}, {}
+    for i, w in enumerate(canvas.windows):
+        if canvas.version == 1 and w.type in SKETCH_TYPES:
+            raise CanvasInvalid(f'windows[{i}].type: "{w.type}" needs version 2')
+        if w.id in ids:
+            raise CanvasInvalid(f'windows[{i}].id: duplicate id "{w.id}"')
+        ids[w.id] = i
+        if w.type in SKETCH_TYPES:
+            continue  # sketch nodes are not referenced from SQL, so their labels need not be unique
+        if w.title in titles:
+            raise CanvasInvalid(f'windows[{i}].title: duplicate title "{w.title}"')
+        titles[w.title] = i
+    for i, link in enumerate(canvas.links):
+        for key, value in (("from", link.from_), ("to", link.to)):
+            if value not in ids:
+                raise CanvasInvalid(f'links[{i}].{key}: no window with id "{value}"')
+    edge_ids = set()
+    for i, edge in enumerate(canvas.edges):
+        if edge.id in edge_ids:
+            raise CanvasInvalid(f'edges[{i}].id: duplicate id "{edge.id}"')
+        edge_ids.add(edge.id)
+        for key, value in (("source", edge.source), ("target", edge.target)):
+            if value not in ids:
+                raise CanvasInvalid(f'edges[{i}].{key}: no window with id "{value}"')
+
+
+def parse_canvas(data) -> Canvas:
+    """Validate a canvas body. Raises CanvasInvalid with a path-prefixed message."""
+    if not isinstance(data, dict):
+        raise CanvasInvalid("(root): body must be a JSON object")
+    try:
+        canvas = Canvas.model_validate(data)
+    except ValidationError as e:
+        first = e.errors()[0]
+        raise CanvasInvalid(f"{_path(first['loc'])}: {first['msg']}") from e
+    check_references(canvas)
+    return canvas
+
+
+def path_for(name: str):
+    return CANVASES / f"{dbx.safe(name)}.json"
+
+
+def list_canvases() -> list[str]:
+    return sorted(p.stem for p in CANVASES.glob("*.json"))
+
+
+def history_dir(name: str):
+    """data/canvas_history/<name>/: the revision counter and the snapshots of one canvas."""
+    return CANVASES.parent / "canvas_history" / dbx.safe(name)
+
+
+def _meta(name: str) -> dict:
+    try:
+        meta = json.loads((history_dir(name) / "meta.json").read_text(encoding="utf-8"))
+    except (OSError, ValueError):
+        return {}
+    return meta if isinstance(meta, dict) else {}
+
+
+def _snapshots(name: str) -> list:
+    d = history_dir(name)
+    return sorted(d.glob("[0-9]*.json")) if d.is_dir() else []
+
+
+def _atomic_write(path, text: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.parent / (path.name + ".tmp")
+    tmp.write_text(text, encoding="utf-8")
+    os.replace(tmp, path)  # overwrite; never leaves a half-written file behind
+
+
+def revision_info(name: str) -> dict:
+    """{"name", "revision", "by", "snapshots"} for a saved canvas. `by` is ui, agent or undo (None before any save)."""
+    with LOCK:
+        if dbx.safe(name) != name or not path_for(name).is_file():
+            raise CanvasNotFound(f"canvas not found: {name}")
+        meta = _meta(name)
+        rev = meta.get("revision")
+        return {"name": name, "revision": rev if isinstance(rev, int) else 0, "by": meta.get("by"),
+                "snapshots": len(_snapshots(name))}
+
+
+def write_canvas(data, by: str = "ui", snapshot: bool = False, expect: Optional[int] = None) -> tuple[str, int]:
+    """Validate and write data/canvases/<name>.json, storing the body exactly as received. Returns (name, revision).
+
+    Every write bumps the revision and records who made it (`by`). With `snapshot`, the file being replaced is kept
+    first (the newest KEEP_SNAPSHOTS are kept), which is what undo_canvas restores. With `expect`, the write is refused
+    (CanvasConflict) when the canvas exists and is not at that revision, so a page cannot overwrite an agent's newer save unseen.
+    """
+    name = parse_canvas(data).name
+    # On Windows (and macOS by default) "Hours" and "hours" are the same file. Refuse
+    # rather than silently overwrite a different canvas.
+    clash = [s for s in list_canvases() if s.lower() == name.lower() and s != name]
+    if clash:
+        raise CanvasInvalid(f"name: '{name}' differs from existing canvas '{clash[0]}' only by case; pick another name")
+    try:
+        text = json.dumps(data, indent=2, allow_nan=False)
+    except ValueError as e:  # NaN or Infinity cannot be stored as JSON
+        raise CanvasInvalid(f"(root): {e}") from e
+    path = path_for(name)
+    with LOCK:
+        meta = _meta(name)
+        old = meta.get("revision") if isinstance(meta.get("revision"), int) else 0
+        if expect is not None and path.is_file() and old != expect:
+            raise CanvasConflict(name, expect, revision_info(name))
+        if snapshot and path.is_file():
+            _atomic_write(history_dir(name) / f"{old:08d}.json", path.read_text(encoding="utf-8"))
+            for stale in _snapshots(name)[:-KEEP_SNAPSHOTS]:
+                stale.unlink(missing_ok=True)
+        _atomic_write(path, text)
+        _atomic_write(history_dir(name) / "meta.json", json.dumps({"revision": old + 1, "by": by}))
+        return name, old + 1
+
+
+def save_canvas(data) -> str:
+    """Validate and write data/canvases/<name>.json, storing the body exactly as received."""
+    return write_canvas(data)[0]
+
+
+def undo_canvas(name: str) -> dict:
+    """Put back the newest snapshot of a canvas (the state before an agent's last write) and drop that snapshot.
+
+    Returns {"name", "revision", "snapshots"} (snapshots left). CanvasInvalid when there is nothing to undo.
+    """
+    with LOCK:
+        revision_info(name)  # CanvasNotFound for an unknown canvas
+        snaps = _snapshots(name)
+        if not snaps:
+            raise CanvasInvalid(f"nothing to undo for canvas {name}")
+        data = json.loads(snaps[-1].read_text(encoding="utf-8"))
+        _, rev = write_canvas(data, by="undo")
+        snaps[-1].unlink(missing_ok=True)
+        return {"name": name, "revision": rev, "snapshots": len(_snapshots(name))}
+
+
+def load_canvas(name: str) -> dict:
+    """Return the stored canvas object as saved. Validates it on the way out."""
+    if dbx.safe(name) != name:  # names are validated on save, so this cannot exist
+        raise CanvasNotFound(f"canvas not found: {name}")
+    path = path_for(name)
+    if not path.is_file():
+        raise CanvasNotFound(f"canvas not found: {name}")
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+    except (OSError, ValueError) as e:
+        raise CanvasInvalid(f"(root): stored file is not valid JSON: {e}") from e
+    canvas = parse_canvas(data)
+    if canvas.name != name:
+        raise CanvasInvalid(f'name: "{canvas.name}" does not match file name "{name}"')
+    return data

databricks_explore/cli.py

@@ -0,0 +1,80 @@
+"""The `databricks-explore` command: parse options, pick the data folder, start the server and open the browser."""
+import argparse
+import os
+import sys
+import threading
+import webbrowser
+from pathlib import Path
+from typing import Optional, Sequence
+
+from . import __version__
+
+HOST = "127.0.0.1"  # the app never listens on anything else
+DEFAULT_PORT, DEMO_PORT = 8765, 8767  # demo mode has its own port so the browser keeps its saved profile apart
+
+
+def build_parser() -> argparse.ArgumentParser:
+    """The option parser. Defaults for --port and --data-dir are filled in by main() because they depend on --demo."""
+    p = argparse.ArgumentParser(
+        prog="databricks-explore",
+        description="Local Unity Catalog explorer: query canvas, relationship graph and sketching. Serves on 127.0.0.1 only.",
+        epilog="Environment: EXPLORER_DATA_DIR (same as --data-dir), PORT (same as --port), EXPLORER_DEMO=1 (same as --demo).",
+    )
+    p.add_argument("--version", action="version", version=f"databricks-explore {__version__}")
+    p.add_argument("--demo", action="store_true", help="serve an invented workspace (no Databricks needed) on its own port and data folder")
+    p.add_argument("--redacted", action="store_true", help="with --demo: make query history text read <REDACTED>, to try that path")
+    p.add_argument("--port", type=int, help=f"port to listen on (default {DEFAULT_PORT}, or {DEMO_PORT} with --demo)")
+    p.add_argument("--data-dir", metavar="DIR", help="where the app keeps its cache, canvases and relationships "
+                   "(default: your per-user data folder, with a demo/ folder inside it for --demo)")
+    p.add_argument("--debug", action="store_true", help="Werkzeug debugger and reloader (the debugger can run code from an error page)")
+    p.add_argument("--no-browser", action="store_true", help="do not open a browser tab")
+    return p
+
+
+def source_checkout_hint(data_dir_given: bool) -> Optional[str]:
+    """One line of advice when running from a source checkout that still has its old data/ folder, else None."""
+    if data_dir_given:
+        return None
+    root = Path(__file__).resolve().parents[2]
+    if (root / "pyproject.toml").is_file() and (root / "data").is_dir():
+        return f"Hint: {root / 'data'} is from an earlier version and is not used now. To keep using it: --data-dir {root / 'data'}"
+    return None
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    """Run the app. Returns the exit status."""
+    args = build_parser().parse_args(argv)
+    # dbx reads these when it is first imported, so they are set before the app is.
+    demo = args.demo or os.environ.get("EXPLORER_DEMO") == "1"
+    if demo:
+        os.environ["EXPLORER_DEMO"] = "1"
+    if args.redacted:
+        if not demo:
+            print("--redacted only applies with --demo", file=sys.stderr)
+            return 2
+        os.environ["EXPLORER_DEMO_REDACTED"] = "1"
+    data_dir_given = bool(args.data_dir or os.environ.get("EXPLORER_DATA_DIR"))
+    if args.data_dir:
+        os.environ["EXPLORER_DATA_DIR"] = str(Path(args.data_dir).expanduser().resolve())
+    port = args.port if args.port is not None else int(os.environ.get("PORT") or (DEMO_PORT if demo else DEFAULT_PORT))
+    debug = args.debug or os.environ.get("EXPLORER_DEBUG") == "1"
+
+    from . import dbx
+    from .app import app
+
+    url = f"http://{HOST}:{port}"
+    # With the reloader the first process only supervises; the server (and the message) belong to its child.
+    serving = not debug or os.environ.get("WERKZEUG_RUN_MAIN") == "true"
+    if serving:
+        print(f"Databricks Explorer {__version__}{' (demo mode)' if demo else ''} at {url}", flush=True)
+        print(f"Data folder: {dbx.DATA}", flush=True)
+        if hint := source_checkout_hint(data_dir_given):
+            print(hint, flush=True)
+        if not args.no_browser:
+            threading.Timer(1.0, webbrowser.open, args=(url,)).start()
+    try:
+        app.run(host=HOST, port=port, debug=debug)
+    except OSError as e:
+        print(f"Could not listen on {HOST}:{port} ({e.strerror or e}). Pick another port with --port.", file=sys.stderr)
+        return 1
+    return 0