wehub-resource-sync
271 lines
11 KiB
Python
271 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from tracker.config import Config, clear_tokens, load_config, update_tokens
|
|
from tracker.models import Experiment, Run
|
|
|
|
|
|
class AuthError(Exception):
|
|
"""Raised when an operation requires valid credentials that are absent or expired."""
|
|
|
|
|
|
class APIError(Exception):
|
|
"""Raised when the InsForge API returns a non-2xx response."""
|
|
|
|
def __init__(self, status_code: int, code: str, message: str) -> None:
|
|
"""Store the HTTP status code, machine-readable error code, and human-readable message."""
|
|
super().__init__(message)
|
|
self.status_code = status_code
|
|
self.code = code
|
|
self.message = message
|
|
|
|
|
|
class InsForgeClient:
|
|
"""HTTP client for the InsForge BaaS API with automatic token refresh on 401."""
|
|
|
|
def __init__(self, config: Config | None = None) -> None:
|
|
"""Initialise the client, loading config from disk if not provided."""
|
|
self._config = config or load_config()
|
|
self._http = httpx.Client(
|
|
base_url=self._config.server_url,
|
|
timeout=30.0,
|
|
)
|
|
|
|
def close(self) -> None:
|
|
"""Close the underlying HTTP connection pool."""
|
|
self._http.close()
|
|
|
|
def __enter__(self) -> "InsForgeClient":
|
|
"""Support use as a context manager, returning self."""
|
|
return self
|
|
|
|
def __exit__(self, *args: Any) -> None:
|
|
"""Close the HTTP client on context-manager exit."""
|
|
self.close()
|
|
|
|
# ── Internal helpers ──────────────────────────────────────────────────
|
|
|
|
def _auth_headers(self) -> dict[str, str]:
|
|
"""Return the Authorization header dict, raising AuthError if not logged in."""
|
|
if not self._config.access_token:
|
|
raise AuthError("Not logged in. Run `tracker login` first.")
|
|
return {"Authorization": f"Bearer {self._config.access_token}"}
|
|
|
|
def _raise_for_error(self, response: httpx.Response) -> None:
|
|
"""Raise APIError if the response is not a 2xx success."""
|
|
if response.is_success:
|
|
return
|
|
try:
|
|
body = response.json()
|
|
code = body.get("error", "UNKNOWN")
|
|
message = body.get("message", response.text)
|
|
except Exception:
|
|
code = "UNKNOWN"
|
|
message = response.text or f"HTTP {response.status_code}"
|
|
raise APIError(response.status_code, code, message)
|
|
|
|
def _refresh_token(self) -> None:
|
|
"""Exchange the stored refresh token for a new access token, updating config on disk."""
|
|
if not self._config.refresh_token:
|
|
raise AuthError("Session expired. Run `tracker login` again.")
|
|
resp = self._http.post(
|
|
"/api/auth/sessions",
|
|
json={"refreshToken": self._config.refresh_token},
|
|
)
|
|
if not resp.is_success:
|
|
clear_tokens()
|
|
raise AuthError("Session expired. Run `tracker login` again.")
|
|
data = resp.json()
|
|
new_access = data.get("accessToken") or data.get("access_token", "")
|
|
if not new_access:
|
|
clear_tokens()
|
|
raise AuthError("Session expired. Run `tracker login` again.")
|
|
new_refresh = data.get("refreshToken") or data.get("refresh_token", self._config.refresh_token)
|
|
self._config.access_token = new_access
|
|
self._config.refresh_token = new_refresh
|
|
update_tokens(new_access, new_refresh)
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
path: str,
|
|
*,
|
|
json: object = None,
|
|
params: dict[str, Any] | None = None,
|
|
extra_headers: dict[str, str] | None = None,
|
|
retry_on_401: bool = True,
|
|
) -> httpx.Response:
|
|
"""Send an authenticated request, transparently retrying once after a 401."""
|
|
headers = {**self._auth_headers(), **(extra_headers or {})}
|
|
resp = self._http.request(method, path, json=json, params=params, headers=headers)
|
|
|
|
if resp.status_code == 401 and retry_on_401:
|
|
self._refresh_token()
|
|
return self._request(
|
|
method, path,
|
|
json=json, params=params, extra_headers=extra_headers,
|
|
retry_on_401=False,
|
|
)
|
|
|
|
self._raise_for_error(resp)
|
|
return resp
|
|
|
|
# ── Auth ──────────────────────────────────────────────────────────────
|
|
|
|
def register(self, email: str, password: str) -> dict[str, Any]:
|
|
"""Create a new InsForge user account and return the raw response body."""
|
|
resp = self._http.post("/api/auth/users", json={"email": email, "password": password})
|
|
self._raise_for_error(resp)
|
|
return resp.json()
|
|
|
|
def login(self, email: str, password: str) -> tuple[str, str]:
|
|
"""Authenticate with email/password and return (access_token, refresh_token)."""
|
|
resp = self._http.post("/api/auth/sessions", json={"email": email, "password": password})
|
|
self._raise_for_error(resp)
|
|
data = resp.json()
|
|
access = data.get("accessToken") or data.get("access_token", "")
|
|
refresh = data.get("refreshToken") or data.get("refresh_token", "")
|
|
return access, refresh
|
|
|
|
# ── Table management ─────────────────────────────────────────────────
|
|
|
|
def table_exists(self, table: str) -> bool:
|
|
"""Return True if the table is accessible, False if it doesn't exist (404)."""
|
|
try:
|
|
self._request("GET", f"/api/database/records/{table}", params={"limit": 0})
|
|
return True
|
|
except APIError as e:
|
|
if e.status_code == 404:
|
|
return False
|
|
raise
|
|
|
|
# ── Generic CRUD helpers ─────────────────────────────────────────────
|
|
|
|
def _list_records(
|
|
self,
|
|
table: str,
|
|
filters: dict[str, str] | None = None,
|
|
limit: int | None = None,
|
|
offset: int = 0,
|
|
order: str | None = None,
|
|
select: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
"""Fetch rows from a database table with optional filtering, ordering, and pagination."""
|
|
params: dict[str, Any] = {"offset": offset}
|
|
if limit is not None:
|
|
params["limit"] = limit
|
|
if order:
|
|
params["order"] = order
|
|
if select:
|
|
params["select"] = select
|
|
if filters:
|
|
params.update(filters)
|
|
resp = self._request("GET", f"/api/database/records/{table}", params=params)
|
|
return resp.json()
|
|
|
|
def _create_record(self, table: str, record: dict[str, Any]) -> dict[str, Any]:
|
|
"""Insert a single record into a table and return the created row."""
|
|
resp = self._request(
|
|
"POST",
|
|
f"/api/database/records/{table}",
|
|
json=[record],
|
|
extra_headers={"Prefer": "return=representation"},
|
|
)
|
|
rows = resp.json()
|
|
return rows[0] if isinstance(rows, list) else rows
|
|
|
|
# ── Experiments ──────────────────────────────────────────────────────
|
|
|
|
def get_or_create_experiment(
|
|
self, name: str, description: str | None = None
|
|
) -> Experiment:
|
|
"""Return the named experiment, creating it if it does not yet exist."""
|
|
rows = self._list_records(
|
|
"experiments",
|
|
filters={"name": f"eq.{name}"},
|
|
limit=1,
|
|
)
|
|
if rows:
|
|
return Experiment.from_dict(rows[0])
|
|
|
|
record: dict[str, Any] = {"name": name}
|
|
if description:
|
|
record["description"] = description
|
|
try:
|
|
created = self._create_record("experiments", record)
|
|
return Experiment.from_dict(created)
|
|
except APIError as e:
|
|
# UNIQUE violation: another caller created it; retry GET
|
|
if e.status_code in (409, 422) or "unique" in e.message.lower():
|
|
rows = self._list_records(
|
|
"experiments",
|
|
filters={"name": f"eq.{name}"},
|
|
limit=1,
|
|
)
|
|
if rows:
|
|
return Experiment.from_dict(rows[0])
|
|
raise
|
|
|
|
def list_experiments(self, limit: int = 50) -> list[Experiment]:
|
|
"""Return up to *limit* experiments ordered by creation time, newest first."""
|
|
rows = self._list_records("experiments", limit=limit, order="created_at.desc")
|
|
return [Experiment.from_dict(r) for r in rows]
|
|
|
|
# ── Runs ─────────────────────────────────────────────────────────────
|
|
|
|
def create_run(
|
|
self,
|
|
experiment_id: str,
|
|
experiment_name: str,
|
|
run_name: str | None,
|
|
params: dict[str, Any],
|
|
metrics: dict[str, Any],
|
|
notes: str | None,
|
|
started_at: str,
|
|
finished_at: str,
|
|
) -> Run:
|
|
"""Insert a run record and return the persisted Run model."""
|
|
record: dict[str, Any] = {
|
|
"experiment_id": experiment_id,
|
|
"experiment_name": experiment_name,
|
|
"params": params,
|
|
"metrics": metrics,
|
|
"started_at": started_at,
|
|
"finished_at": finished_at,
|
|
}
|
|
if run_name:
|
|
record["run_name"] = run_name
|
|
if notes:
|
|
record["notes"] = notes
|
|
created = self._create_record("runs", record)
|
|
return Run.from_dict(created)
|
|
|
|
def list_runs(
|
|
self,
|
|
experiment_name: str | None = None,
|
|
limit: int = 20,
|
|
offset: int = 0,
|
|
) -> list[Run]:
|
|
"""Return recent runs, optionally filtered by experiment name."""
|
|
filters: dict[str, str] = {}
|
|
if experiment_name:
|
|
filters["experiment_name"] = f"eq.{experiment_name}"
|
|
rows = self._list_records(
|
|
"runs",
|
|
filters=filters,
|
|
limit=limit,
|
|
offset=offset,
|
|
order="created_at.desc",
|
|
)
|
|
return [Run.from_dict(r) for r in rows]
|
|
|
|
def get_run(self, run_id: str) -> Run:
|
|
"""Fetch a single run by its full UUID, raising APIError(404) if missing."""
|
|
rows = self._list_records("runs", filters={"id": f"eq.{run_id}"}, limit=1)
|
|
if not rows:
|
|
raise APIError(404, "NOT_FOUND", f"Run not found: {run_id}")
|
|
return Run.from_dict(rows[0])
|