| 1 | """Locate and talk to `witdb-bridge` over NDJSON stdin/stdout.""" |
| 2 | |
| 3 | from __future__ import annotations |
| 4 | |
| 5 | import base64 |
| 6 | import json |
| 7 | import os |
| 8 | import subprocess |
| 9 | import sys |
| 10 | from functools import lru_cache |
| 11 | from pathlib import Path |
| 12 | from typing import Any |
| 13 | |
| 14 | from pywitdb.exceptions import BridgeNotAvailableError, WitDbProtocolError, WitDbStoreError |
| 15 | |
| 16 | _BRIDGE_EXE = "witdb-bridge.exe" if sys.platform == "win32" else "witdb-bridge" |
| 17 | |
| 18 | |
| 19 | def _pywitdb_root() -> Path: |
| 20 | return Path(__file__).resolve().parents[3] |
| 21 | |
| 22 | |
| 23 | def _search_paths() -> list[Path]: |
| 24 | root = _pywitdb_root() |
| 25 | candidates: list[Path] = [] |
| 26 | env = os.environ.get("WITDB_BRIDGE_PATH") |
| 27 | if env: |
| 28 | candidates.append(Path(env)) |
| 29 | for tfm in ("net10.0", "net9.0"): |
| 30 | candidates.append( |
| 31 | root / "bridge" / "WitDbBridge" / "bin" / "Release" / tfm / _BRIDGE_EXE |
| 32 | ) |
| 33 | return candidates |
| 34 | |
| 35 | |
| 36 | @lru_cache(maxsize=1) |
| 37 | def bridge_available() -> bool: |
| 38 | return resolve_bridge_executable() is not None |
| 39 | |
| 40 | |
| 41 | @lru_cache(maxsize=1) |
| 42 | def resolve_bridge_executable() -> Path | None: |
| 43 | for path in _search_paths(): |
| 44 | if path.is_file(): |
| 45 | return path |
| 46 | return None |
| 47 | |
| 48 | |
| 49 | class BridgeClient: |
| 50 | def __init__(self, proc: subprocess.Popen[bytes]) -> None: |
| 51 | self._proc = proc |
| 52 | self._next_id = 0 |
| 53 | |
| 54 | @classmethod |
| 55 | def start(cls) -> BridgeClient: |
| 56 | exe = resolve_bridge_executable() |
| 57 | if exe is None: |
| 58 | raise BridgeNotAvailableError( |
| 59 | "witdb-bridge not found; build bridge/WitDbBridge or set WITDB_BRIDGE_PATH" |
| 60 | ) |
| 61 | proc = subprocess.Popen( |
| 62 | [str(exe)], |
| 63 | stdin=subprocess.PIPE, |
| 64 | stdout=subprocess.PIPE, |
| 65 | stderr=subprocess.PIPE, |
| 66 | ) |
| 67 | if proc.stdin is None or proc.stdout is None: |
| 68 | proc.kill() |
| 69 | raise BridgeNotAvailableError("failed to start witdb-bridge (no pipes)") |
| 70 | return cls(proc) |
| 71 | |
| 72 | def request(self, op: str, **payload: Any) -> dict[str, Any]: |
| 73 | if self._proc.poll() is not None: |
| 74 | raise BridgeNotAvailableError("witdb-bridge process exited") |
| 75 | |
| 76 | self._next_id += 1 |
| 77 | msg_id = self._next_id |
| 78 | body = {"id": msg_id, "op": op, **payload} |
| 79 | line = (json.dumps(body, separators=(",", ":")) + "\n").encode("utf-8") |
| 80 | assert self._proc.stdin is not None |
| 81 | assert self._proc.stdout is not None |
| 82 | self._proc.stdin.write(line) |
| 83 | self._proc.stdin.flush() |
| 84 | |
| 85 | raw = self._proc.stdout.readline() |
| 86 | if not raw: |
| 87 | raise WitDbProtocolError("witdb-bridge closed stdout") |
| 88 | try: |
| 89 | resp = json.loads(raw.decode("utf-8")) |
| 90 | except json.JSONDecodeError as exc: |
| 91 | raise WitDbProtocolError(f"invalid JSON from bridge: {raw!r}") from exc |
| 92 | |
| 93 | if resp.get("id") != msg_id: |
| 94 | raise WitDbProtocolError(f"response id mismatch: expected {msg_id}, got {resp.get('id')}") |
| 95 | return resp |
| 96 | |
| 97 | def ensure_ok(self, resp: dict[str, Any]) -> dict[str, Any]: |
| 98 | if resp.get("ok"): |
| 99 | return resp |
| 100 | err = resp.get("error") or {} |
| 101 | code = err.get("code", "store_error") |
| 102 | message = err.get("message", "bridge error") |
| 103 | raise WitDbStoreError(f"{code}: {message}") |
| 104 | |
| 105 | def open(self, path: str, *, password: str | None = None, create: bool = True) -> dict[str, Any]: |
| 106 | payload: dict[str, Any] = {"path": path, "create": create} |
| 107 | if password is not None: |
| 108 | payload["password"] = password |
| 109 | return self.ensure_ok(self.request("open", **payload)) |
| 110 | |
| 111 | def close(self) -> None: |
| 112 | try: |
| 113 | self.ensure_ok(self.request("close")) |
| 114 | finally: |
| 115 | if self._proc.stdin: |
| 116 | self._proc.stdin.close() |
| 117 | self._proc.wait(timeout=10) |
| 118 | |
| 119 | @staticmethod |
| 120 | def b64(data: bytes) -> str: |
| 121 | return base64.b64encode(data).decode("ascii") |
| 122 | |
| 123 | @staticmethod |
| 124 | def unb64(text: str) -> bytes: |
| 125 | return base64.b64decode(text.encode("ascii")) |
| 126 | |