gatewaysdk 0.3.2__py3-none-any.whl

gatewaysdk/world_data.py

@@ -0,0 +1,453 @@
+"""Rows into a platform world, in bulk: the client side of world data batches.
+
+The wire is the platform's ``/api/public/worlds/{slug}/data`` routes (the contract,
+section 7): open a batch, register each chunk by digest, PUT the chunk's bytes to
+the presigned URL the platform hands back (or skip it when the platform already
+holds that digest), mark it uploaded, complete the batch, poll it to a terminal
+state. Auth and host resolution match the other clients: the project's public +
+secret key as HTTP Basic (``GATEWAY_PUBLIC_KEY`` / ``GATEWAY_SECRET_KEY`` /
+``GATEWAY_HOST``). Standard library only.
+
+:func:`import_rows` is the whole flow the CLI runs: rows streamed from files, cut
+into gzipped chunks one at a time, uploaded ``parallel`` at a time, with a manifest
+at ``.gateway/imports/<batchId>.json`` (section 10) written as each chunk lands so
+``--resume`` can pick the upload up where it stopped.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+import sys
+import tempfile
+import threading
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable, Dict, Iterable, List, Optional
+
+from gatewaysdk.environment.schema.batch import CHUNK_BYTES, Chunk, chunk_rows, iter_rows, write_report
+
+MANIFEST_VERSION = 1
+MANIFEST_DIR = Path(".gateway") / "imports"
+TERMINAL = ("applied", "refused", "published", "failed")
+RESUMABLE = ("open", "uploading")
+POLL_SECONDS = 2.0
+UPLOAD_TIMEOUT = 600
+
+
+class WorldDataError(RuntimeError):
+    """Raised when any world data API call fails."""
+
+
+# ---------------------------------------------------------------- the client
+
+
+class WorldDataClient:
+    """Talks to a Gateway host's world data batch API."""
+
+    def __init__(self, host: str, public_key: str, secret_key: str) -> None:
+        if not host:
+            raise ValueError("host is required (e.g. https://withgateway.ai)")
+        if not public_key or not secret_key:
+            raise ValueError("both public_key and secret_key are required")
+        self.host = host.rstrip("/")
+        self._auth = base64.b64encode(f"{public_key}:{secret_key}".encode("utf-8")).decode("ascii")
+
+    @classmethod
+    def from_env(cls, host: Optional[str] = None) -> "WorldDataClient":
+        return cls(
+            host=host or os.environ.get("GATEWAY_HOST", ""),
+            public_key=os.environ.get("GATEWAY_PUBLIC_KEY", ""),
+            secret_key=os.environ.get("GATEWAY_SECRET_KEY", ""),
+        )
+
+    def request(self, method: str, path: str, body: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
+        headers = {"Authorization": f"Basic {self._auth}"}
+        data = None
+        if body is not None:
+            headers["Content-Type"] = "application/json"
+            data = json.dumps(body).encode("utf-8")
+        req = urllib.request.Request(f"{self.host}{path}", data=data, method=method, headers=headers)
+        try:
+            with urllib.request.urlopen(req, timeout=120) as resp:
+                text = resp.read().decode("utf-8")
+                return json.loads(text) if text else {}
+        except urllib.error.HTTPError as e:
+            detail = e.read().decode("utf-8", "replace")[:500]
+            raise WorldDataError(f"{method} {path} -> {e.code}: {detail}") from None
+        except urllib.error.URLError as e:
+            raise WorldDataError(f"{method} {path} -> {e.reason}") from None
+
+    @staticmethod
+    def route(slug: str) -> str:
+        return f"/api/public/worlds/{urllib.parse.quote(slug.split('@', 1)[0], safe='')}/data"
+
+    def open_batch(self, slug: str, **flags: Any) -> Dict[str, Any]:
+        return self.request("POST", f"{self.route(slug)}/batches", {k: v for k, v in flags.items() if v is not None})
+
+    def register_chunk(self, slug: str, batch_id: str, chunk: Dict[str, Any]) -> Dict[str, Any]:
+        return self.request("POST", f"{self.route(slug)}/batches/{batch_id}/chunks", chunk)
+
+    def mark_uploaded(self, slug: str, batch_id: str, chunk_id: str) -> Dict[str, Any]:
+        return self.request("POST", f"{self.route(slug)}/batches/{batch_id}/chunks/{chunk_id}/uploaded")
+
+    def complete(self, slug: str, batch_id: str) -> Dict[str, Any]:
+        return self.request("POST", f"{self.route(slug)}/batches/{batch_id}/complete")
+
+    def status(self, slug: str, batch_id: str) -> Dict[str, Any]:
+        return self.request("GET", f"{self.route(slug)}/batches/{batch_id}")
+
+    def batches(self, slug: str, *, limit: Optional[int] = None, cursor: Optional[str] = None) -> Dict[str, Any]:
+        query = urllib.parse.urlencode({k: v for k, v in (("limit", limit), ("cursor", cursor)) if v is not None})
+        return self.request("GET", f"{self.route(slug)}/batches" + (f"?{query}" if query else ""))
+
+    def publish(self, slug: str, message: Optional[str] = None) -> Dict[str, Any]:
+        return self.request("POST", f"{self.route(slug)}/publish", {"message": message} if message else {})
+
+    def counts(self, slug: str, entity: Optional[str] = None) -> Dict[str, Any]:
+        query = f"?entity={urllib.parse.quote(entity, safe='')}" if entity else ""
+        return self.request("GET", f"{self.route(slug)}{query}")
+
+    @staticmethod
+    def put_object(url: str, path: Path, size: int, headers: Optional[Dict[str, str]] = None) -> None:
+        """PUT a chunk's bytes to its presigned URL; no auth header, the URL is the grant.
+        ``headers`` are exactly the ones the platform signed the URL for (content-type,
+        content-length, x-amz-checksum-sha256); without them the plain pair is sent."""
+        signed = dict(headers) if headers else {"Content-Type": "application/gzip", "Content-Length": str(size)}
+        with open(path, "rb") as stream:
+            req = urllib.request.Request(url, data=stream, method="PUT", headers=signed)
+            try:
+                with urllib.request.urlopen(req, timeout=UPLOAD_TIMEOUT) as resp:
+                    resp.read()
+            except urllib.error.HTTPError as e:
+                detail = e.read().decode("utf-8", "replace")[:300]
+                raise WorldDataError(f"chunk PUT -> {e.code}: {detail}") from None
+            except urllib.error.URLError as e:
+                raise WorldDataError(f"chunk PUT -> {e.reason}") from None
+
+    @staticmethod
+    def fetch_report(url: str) -> str:
+        try:
+            with urllib.request.urlopen(urllib.request.Request(url), timeout=120) as resp:
+                return resp.read().decode("utf-8")
+        except (urllib.error.HTTPError, urllib.error.URLError) as e:
+            raise WorldDataError(f"report download failed: {e}") from None
+
+
+# ---------------------------------------------------------------- the manifest
+
+
+def _now() -> str:
+    return datetime.now(timezone.utc).isoformat(timespec="seconds").replace("+00:00", "Z")
+
+
+@dataclass
+class Flags:
+    mode: str = "append"
+    merge: str = "upsert"
+    atomic: bool = True
+    publish: str = "now"
+    dryRun: bool = False
+    chunkBytes: int = CHUNK_BYTES
+    chunkRows: Optional[int] = None
+    entity: Optional[str] = None
+
+    def to_json(self) -> Dict[str, Any]:
+        return dict(self.__dict__)
+
+
+class Manifest:
+    """``.gateway/imports/<batchId>.json``: what was uploaded, so a run can resume."""
+
+    def __init__(self, path: Path, data: Dict[str, Any]):
+        self.path = path
+        self.data = data
+        self._lock = threading.Lock()
+
+    @classmethod
+    def create(cls, base: Path, batch_id: str, slug: str, host: str, flags: Flags, inputs: List[str]) -> "Manifest":
+        directory = base / MANIFEST_DIR
+        directory.mkdir(parents=True, exist_ok=True)
+        ignore = base / ".gateway" / ".gitignore"
+        if not ignore.exists():
+            ignore.write_text("*\n")
+        manifest = cls(
+            directory / f"{batch_id}.json",
+            {
+                "version": MANIFEST_VERSION,
+                "batchId": batch_id,
+                "slug": slug,
+                "host": host,
+                "flags": flags.to_json(),
+                "inputs": list(inputs),
+                "chunks": [],
+                "completed": False,
+                "createdAt": _now(),
+                "updatedAt": _now(),
+            },
+        )
+        manifest.save()
+        return manifest
+
+    @classmethod
+    def load(cls, path: Path) -> "Manifest":
+        data = json.loads(path.read_text(encoding="utf-8"))
+        if data.get("version") != MANIFEST_VERSION:
+            raise WorldDataError(f"{path} is a manifest of version {data.get('version')}, not {MANIFEST_VERSION}")
+        return cls(path, data)
+
+    @classmethod
+    def find(cls, base: Path, slug: str, inputs: Optional[List[str]] = None) -> Optional["Manifest"]:
+        """The newest incomplete manifest for ``slug`` (and these inputs when given)."""
+        directory = base / MANIFEST_DIR
+        if not directory.is_dir():
+            return None
+        candidates = []
+        for path in directory.glob("*.json"):
+            try:
+                manifest = cls.load(path)
+            except (OSError, ValueError, WorldDataError):
+                continue
+            data = manifest.data
+            if data.get("completed") or data.get("slug") != slug.split("@", 1)[0]:
+                continue
+            if inputs is not None and data.get("inputs") != list(inputs):
+                continue
+            candidates.append(manifest)
+        if not candidates:
+            return None
+        return max(candidates, key=lambda m: m.data.get("updatedAt", ""))
+
+    @property
+    def batch_id(self) -> str:
+        return self.data["batchId"]
+
+    @property
+    def flags(self) -> Flags:
+        return Flags(**self.data["flags"])
+
+    def save(self) -> None:
+        with self._lock:
+            self.data["updatedAt"] = _now()
+            tmp = self.path.with_suffix(".json.tmp")
+            tmp.write_text(json.dumps(self.data, indent=1) + "\n", encoding="utf-8")
+            tmp.replace(self.path)
+
+    def record(self, entry: Dict[str, Any]) -> None:
+        with self._lock:
+            chunks = [c for c in self.data["chunks"] if c["ordinal"] != entry["ordinal"]]
+            chunks.append(entry)
+            self.data["chunks"] = sorted(chunks, key=lambda c: c["ordinal"])
+        self.save()
+
+    def uploaded(self, ordinal: int, digest: str) -> Optional[Dict[str, Any]]:
+        for entry in self.data["chunks"]:
+            if entry["ordinal"] == ordinal and entry["digest"] == digest and entry.get("uploaded"):
+                return entry
+        return None
+
+
+# ---------------------------------------------------------------- the flow
+
+
+Progress = Callable[[str], None]
+
+
+def _quiet(_: str) -> None:
+    pass
+
+
+def _upload(client: WorldDataClient, slug: str, batch_id: str, chunk: Chunk, manifest: Manifest) -> Dict[str, Any]:
+    entry = chunk.manifest_entry()
+    registered = client.register_chunk(
+        slug,
+        batch_id,
+        {
+            k: v
+            for k, v in (
+                ("ordinal", chunk.ordinal),
+                ("digest", chunk.digest),
+                ("bytes", chunk.bytes),
+                ("rows", chunk.rows),
+                ("entity", chunk.entity),
+            )
+            if v is not None
+        },
+    )
+    chunk_id = registered.get("chunkId")
+    entry["chunkId"] = chunk_id
+    if not registered.get("uploaded"):
+        url = registered.get("uploadUrl")
+        if not url:
+            raise WorldDataError(f"chunk {chunk.ordinal}: the platform gave neither uploaded nor an uploadUrl")
+        headers = registered.get("uploadHeaders")
+        if headers is not None and not (
+            isinstance(headers, dict) and all(isinstance(k, str) and isinstance(v, str) for k, v in headers.items())
+        ):
+            raise WorldDataError(f"chunk {chunk.ordinal}: uploadHeaders must map header names to strings")
+        client.put_object(url, chunk.path, chunk.bytes, headers)
+        client.mark_uploaded(slug, batch_id, chunk_id)
+    entry["uploaded"] = True
+    manifest.record(entry)
+    chunk.path.unlink(missing_ok=True)
+    return entry
+
+
+def import_rows(
+    client: WorldDataClient,
+    slug: str,
+    inputs: List[str],
+    *,
+    flags: Flags,
+    parallel: int = 4,
+    resume: bool = False,
+    message: Optional[str] = None,
+    base: Optional[Path] = None,
+    progress: Progress = _quiet,
+) -> Manifest:
+    """Stream the rows files into chunks, upload them, complete the batch. Returns the
+    manifest; the caller polls the batch with :func:`wait`."""
+    base = base or Path.cwd()
+    manifest: Optional[Manifest] = None
+    if resume:
+        manifest = Manifest.find(base, slug, inputs)
+        if manifest is None:
+            raise WorldDataError(f"nothing to resume: no incomplete manifest for {slug} under {base / MANIFEST_DIR}")
+        state = client.status(slug, manifest.batch_id).get("state")
+        if state not in RESUMABLE:
+            raise WorldDataError(f"batch {manifest.batch_id} is {state}; only an open or uploading batch resumes")
+        flags = manifest.flags
+        progress(f"resuming batch {manifest.batch_id} ({len(manifest.data['chunks'])} chunks recorded)")
+    else:
+        opened = client.open_batch(
+            slug,
+            mode=flags.mode,
+            merge=flags.merge,
+            atomic=flags.atomic,
+            dryRun=flags.dryRun,
+            publish=flags.publish,
+            message=message,
+        )
+        batch_id = opened.get("batchId")
+        if not isinstance(batch_id, str) or not batch_id:
+            raise WorldDataError("the platform opened no batch (no batchId in the answer)")
+        limits = opened.get("limits") or {}
+        if isinstance(limits.get("chunkBytes"), int) and flags.chunkBytes > limits["chunkBytes"]:
+            raise WorldDataError(
+                f"--chunk-bytes {flags.chunkBytes} exceeds the platform's chunk limit {limits['chunkBytes']}"
+            )
+        manifest = Manifest.create(base, batch_id, slug.split("@", 1)[0], client.host, flags, inputs)
+        progress(f"batch {batch_id} open")
+    assert manifest is not None
+    batch_id = manifest.batch_id
+    errors: List[BaseException] = []
+    with tempfile.TemporaryDirectory(prefix="gateway-chunks-") as workdir:
+        with ThreadPoolExecutor(max_workers=max(1, parallel)) as pool:
+            pending = []
+            rows = iter_rows(inputs, entity=flags.entity)
+            for chunk in chunk_rows(
+                rows, workdir, chunk_bytes=flags.chunkBytes, chunk_rows=flags.chunkRows, entity=flags.entity
+            ):
+                if errors:
+                    break
+                known = manifest.uploaded(chunk.ordinal, chunk.digest)
+                if known is not None:
+                    progress(f"chunk {chunk.ordinal}: already uploaded ({chunk.rows} rows)")
+                    chunk.path.unlink(missing_ok=True)
+                    continue
+
+                def job(chunk: Chunk = chunk) -> None:
+                    try:
+                        entry = _upload(client, slug, batch_id, chunk, manifest)
+                        progress(
+                            f"chunk {chunk.ordinal}: uploaded {chunk.rows} rows, {chunk.bytes} bytes ({entry['chunkId']})"
+                        )
+                    except BaseException as error:  # surfaced after the pool drains
+                        errors.append(error)
+
+                pending.append(pool.submit(job))
+                # Bound the chunks on disk to the workers: the producer waits for a slot.
+                while sum(1 for f in pending if not f.done()) >= max(1, parallel):
+                    time.sleep(0.05)
+            for future in pending:
+                future.result()
+    if errors:
+        raise errors[0]
+    client.complete(slug, batch_id)
+    manifest.data["completed"] = True
+    manifest.save()
+    progress(f"batch {batch_id} complete: {len(manifest.data['chunks'])} chunks")
+    return manifest
+
+
+def wait(
+    client: WorldDataClient,
+    slug: str,
+    batch_id: str,
+    *,
+    progress: Progress = _quiet,
+    poll: float = POLL_SECONDS,
+    timeout: Optional[float] = None,
+) -> Dict[str, Any]:
+    """Poll the batch until a terminal state, one progress line per phase change."""
+    started = time.monotonic()
+    last = None
+    while True:
+        status = client.status(slug, batch_id)
+        phase = (status.get("state"), (status.get("progress") or {}).get("phase"))
+        if phase != last:
+            last = phase
+            detail = status.get("progress") or {}
+            progress(f"{status.get('state')}" + (f" ({detail.get('phase')})" if detail.get("phase") else ""))
+        if status.get("state") in TERMINAL:
+            return status
+        if timeout is not None and time.monotonic() - started > timeout:
+            return status
+        time.sleep(poll)
+
+
+def write_status_report(client: WorldDataClient, status: Dict[str, Any], destination: Path) -> int:
+    """The batch's refusals as a JSONL report: the full report when the platform
+    offers one, else the records the status carries."""
+    url = status.get("reportUrl")
+    if url:
+        destination.write_text(client.fetch_report(url), encoding="utf-8")
+        return int(status.get("refusalCount") or 0)
+    records: Iterable[dict] = status.get("refusals") or []
+    with open(destination, "w", encoding="utf-8") as out:
+        return write_report(records, out)
+
+
+def tty_progress(stream: Any = None) -> Progress:
+    """Progress lines on a TTY; silence otherwise."""
+    stream = stream or sys.stderr
+    if not hasattr(stream, "isatty") or not stream.isatty():
+        return _quiet
+
+    def show(line: str) -> None:
+        print(line, file=stream, flush=True)
+
+    return show
+
+
+__all__ = [
+    "Flags",
+    "MANIFEST_DIR",
+    "MANIFEST_VERSION",
+    "Manifest",
+    "RESUMABLE",
+    "TERMINAL",
+    "WorldDataClient",
+    "WorldDataError",
+    "import_rows",
+    "tty_progress",
+    "wait",
+    "write_status_report",
+]