diff --git a/.env.example b/.env.example index 5dc20dc..b70a659 100644 --- a/.env.example +++ b/.env.example @@ -56,5 +56,7 @@ SOURCERYKIT_ORG_ID=00000000-0000-0000-0000-000000000000 # SOURCERYKIT_TOKEN_STORE=/path/to/app/sourcerykit.env # SOURCERYKIT_PROVABLY_APP_URL=https://app.provably.ai +# Consent page the browser sign-in opens (e.g. a local web app in dev). +# SOURCERYKIT_CONSENT_URL=https://switchboard.provably.ai/consent # SOURCERYKIT_PROVABLY_API_URL=https://api.provably.ai # SOURCERYKIT_PROVABLY_MCP_URL=https://mcp.provably.ai diff --git a/CHANGELOG.md b/CHANGELOG.md index 5afee59..7909cbc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,7 @@ ### Breaking changes - **User API key removed in favor of OAuth tokens** — authentication now uses `PROVABLY_ACCESS_TOKEN` (auto-refreshed via `PROVABLY_REFRESH_TOKEN`); the verify path uses the collection integration key. Existing users must re-run `sourcerykit init`. See the [migration guide](docs/migrations/unreleased/unreleased.md). - **OAuth browser-only login** — `--register`, `--email`, and `--password` are removed; login is browser OAuth (PKCE) only. New accounts are created on the Provably web app during login. +- **Provably API client moved to [`provably-sdk`](https://pypi.org/project/provably-sdk/)** — `sourcerykit.provably` is gone; import from `provably` instead (e.g. `from provably import ProvablyError`, `from provably.service import service`). `ProvablyError` no longer subclasses `SourceryKitError`; catch it explicitly. sourcerykit configures the SDK on import. ### Features - **Idempotent integration bootstrap** — re-running `init`/`doctor --fix` reuses the existing integration (exact collection match) instead of minting a duplicate key and shadow user. diff --git a/pyproject.toml b/pyproject.toml index 9e1020e..c9375dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,7 @@ dependencies = [ "requests >=2.34", "jsonschema>=4.0", "psycopg[binary]>=3.1", + "provably-sdk>=0.3.1,<1", "pydantic>=2.6", "python-dotenv>=1.2", "sqlalchemy>=2.0", diff --git a/src/sourcerykit/__init__.py b/src/sourcerykit/__init__.py index 19a453c..fab8d61 100644 --- a/src/sourcerykit/__init__.py +++ b/src/sourcerykit/__init__.py @@ -1,3 +1,4 @@ +from sourcerykit import _provably # noqa: F401 (configures the Provably SDK first) from sourcerykit.bootstrap import bootstrap_system from sourcerykit.errors import ( SourceryKitBootstrapError, diff --git a/src/sourcerykit/_provably.py b/src/sourcerykit/_provably.py new file mode 100644 index 0000000..500be1f --- /dev/null +++ b/src/sourcerykit/_provably.py @@ -0,0 +1,100 @@ +"""Glue between sourcerykit and the Provably SDK (``provably-sdk``). + +The SDK reads nothing on its own; this module tells it where sourcerykit keeps +its settings and session, and exempts the SDK's traffic from the interceptor. +It configures the SDK on import, which ``sourcerykit/__init__`` does first. +""" + +import json +import os + +import provably +from dotenv import set_key +from provably import OAuthTokens, ProvablyConfig + +from sourcerykit.config import ( + CONFIG_FILE, + get_bootstrap_app_url, + get_bootstrap_settings, + get_settings, + load_app_dir_config, + load_local_env, + save_app_dir_config, +) +from sourcerykit.errors import SourceryKitConfigError +from sourcerykit.intercept._self_egress import provably_self_egress + +# The OAuth client sourcerykit is registered as, and the loopback port it listens on. +OAUTH_CLIENT_ID = "sourcerykit-cli" +OAUTH_LOOPBACK_PORT = 8910 +# The web app's consent page +DEFAULT_CONSENT_URL = "https://switchboard.provably.ai/consent" + + +def _token_store() -> str | None: + """Return the app's token-store path (SOURCERYKIT_TOKEN_STORE) or None for global JSON. + + Resolved per-call so the app can set it at startup regardless of import order. + """ + return os.getenv("SOURCERYKIT_TOKEN_STORE") or None + + +class SourceryKitTokenStore: + """The session in sourcerykit's global JSON, or in an embedding app's ``.env``.""" + + def load(self) -> OAuthTokens | None: + # Raises SourceryKitConfigError when sourcerykit is not set up, as before the SDK split. + settings = get_settings() + # sourcerykit only ever signs in as its own client, so its tokens are that client's. + return OAuthTokens( + access_token=settings.access_token, + refresh_token=settings.refresh_token or None, + client_id=OAUTH_CLIENT_ID, + ) + + def save(self, tokens: OAuthTokens) -> None: + store = _token_store() + if store: + set_key(store, "PROVABLY_ACCESS_TOKEN", tokens.access_token) + set_key(store, "PROVABLY_REFRESH_TOKEN", tokens.refresh_token or "") + os.environ["PROVABLY_ACCESS_TOKEN"] = tokens.access_token + if tokens.refresh_token: + os.environ["PROVABLY_REFRESH_TOKEN"] = tokens.refresh_token + load_local_env.cache_clear() + else: + save_app_dir_config(token=tokens.access_token, refresh_token=tokens.refresh_token) + get_settings.cache_clear() + + def clear_refresh_token(self) -> None: + store = _token_store() + if store: + set_key(store, "PROVABLY_REFRESH_TOKEN", "") + os.environ.pop("PROVABLY_REFRESH_TOKEN", None) + load_local_env.cache_clear() + else: + payload = load_app_dir_config() + payload.pop("refresh_token", None) + CONFIG_FILE.write_text(json.dumps(payload)) + load_app_dir_config.cache_clear() + get_settings.cache_clear() + + +def sdk_config() -> ProvablyConfig: + """sourcerykit's settings as the SDK's; before setup, the API and app URLs alone.""" + try: + settings = get_settings() + except SourceryKitConfigError: + return ProvablyConfig(api_url=get_bootstrap_settings(), app_url=get_bootstrap_app_url()) + return ProvablyConfig(api_url=settings.provably_api, app_url=settings.provably_app, org_id=settings.org_id) + + +def consent_page_url() -> str: + """URL of the web app consent page: ``SOURCERYKIT_CONSENT_URL``, else production. + + A setting of its own, not ``SOURCERYKIT_PROVABLY_APP_URL``: that one is the + base of the app's query-record links, which a consent page URL would break. + """ + return (os.environ.get("SOURCERYKIT_CONSENT_URL") or "").strip() or DEFAULT_CONSENT_URL + + +provably.configure(config=sdk_config, tokens=SourceryKitTokenStore(), egress=provably_self_egress) diff --git a/src/sourcerykit/bootstrap/_cache.py b/src/sourcerykit/bootstrap/_cache.py index c709964..02edb39 100644 --- a/src/sourcerykit/bootstrap/_cache.py +++ b/src/sourcerykit/bootstrap/_cache.py @@ -5,12 +5,14 @@ from dataclasses import dataclass, field from uuid import UUID +from provably import ConnectionInfo, ProvablyError +from provably.service import service + from sourcerykit.config import Settings -from sourcerykit.db._engine import ConnectionInfo, get_connection_info +from sourcerykit.db._engine import get_connection_info +from sourcerykit.db._schema import INTERCEPTS_TABLE from sourcerykit.errors import SourceryKitBootstrapError, SourceryKitError from sourcerykit.logger import get_logger -from sourcerykit.provably._errors import ProvablyError -from sourcerykit.provably.service import service _log = get_logger(__name__) @@ -45,14 +47,17 @@ async def run_handshake(self, project_name: str) -> None: connection_info = get_connection_info() self.database_id = await self._resolve_database(connection_info) - ids = await service.get_database_schema_id_and_table_id(self.middleware_id, connection_info) + ids = await service.get_database_schema_id_and_table_id( + self.middleware_id, connection_info, table_name=INTERCEPTS_TABLE + ) self.schema_id = ids["schema_id"] self.table_id = ids["table_id"] self.collection_name = project_name self.collection_id = await self._resolve_collection(project_name) self.integration_key = await self._resolve_integration_key() - except SourceryKitError: + # ProvablyError comes from provably-sdk, so it is not a SourceryKitError. + except (SourceryKitError, ProvablyError): raise except Exception as e: _log.error("handshake_failed_unexpected", error=str(e)) @@ -110,7 +115,7 @@ async def _resolve_collection(self, project_name: str) -> UUID: async def _resolve_integration_key(self) -> str: if self.collection_id is None: raise SourceryKitBootstrapError("collection_id is not set; _resolve_collection() must succeed first") - _, key = await service.ensure_integration(self.collection_id) + _, key = await service.ensure_integration(self.collection_id, name=INTERCEPTS_TABLE) return key diff --git a/src/sourcerykit/bootstrap/bootstrap.py b/src/sourcerykit/bootstrap/bootstrap.py index a57b1eb..972f91b 100644 --- a/src/sourcerykit/bootstrap/bootstrap.py +++ b/src/sourcerykit/bootstrap/bootstrap.py @@ -1,3 +1,5 @@ +from provably.service import service as provably_service + from sourcerykit.bootstrap._cache import _BOOTSTRAP_INSTANCE, ProvablyBootstrapCache from sourcerykit.config import get_settings, save_local_env from sourcerykit.db._engine import get_engine @@ -8,7 +10,6 @@ ) from sourcerykit.intercept.interceptor import init_interceptor from sourcerykit.logger import get_logger -from sourcerykit.provably.service import service as provably_service _log = get_logger(__name__) diff --git a/src/sourcerykit/cli/doctor.py b/src/sourcerykit/cli/doctor.py index 0b4c18c..bdfdecc 100644 --- a/src/sourcerykit/cli/doctor.py +++ b/src/sourcerykit/cli/doctor.py @@ -4,13 +4,15 @@ import dataclasses from collections.abc import Callable +from provably import ProvablyConnectionError, ProvablyUnauthorizedError +from provably.auth_service import auth_service +from provably.service import service + from sourcerykit.cli.init import run_full_bootstrap from sourcerykit.cli.utils import console, mask_postgres_url, run_connectivity_check from sourcerykit.config import Settings, get_settings from sourcerykit.db._engine import get_connection_info -from sourcerykit.provably._errors import ProvablyConnectionError, ProvablyUnauthorizedError -from sourcerykit.provably.auth_service import auth_service -from sourcerykit.provably.service import service +from sourcerykit.db._schema import INTERCEPTS_TABLE def _check_token_and_org(settings: Settings) -> tuple[bool, str]: @@ -95,7 +97,9 @@ async def _deep_check_collection_and_ids(settings: Settings) -> tuple[bool, str] if remote_db != settings.database_id: return False, f"Database mismatch (local={settings.database_id}, remote={remote_db})" - ids = await service.get_database_schema_id_and_table_id(settings.middleware_id, connection_info) + ids = await service.get_database_schema_id_and_table_id( + settings.middleware_id, connection_info, table_name=INTERCEPTS_TABLE + ) if ids["schema_id"] != settings.schema_id: return False, f"Schema mismatch (local={settings.schema_id}, remote={ids['schema_id']})" if ids["table_id"] != settings.table_id: diff --git a/src/sourcerykit/cli/feedback.py b/src/sourcerykit/cli/feedback.py index d000c74..f0cf10b 100644 --- a/src/sourcerykit/cli/feedback.py +++ b/src/sourcerykit/cli/feedback.py @@ -3,9 +3,9 @@ import questionary import typer +from provably.service import ProvablyService from sourcerykit.cli.utils import console, require_settings -from sourcerykit.provably.service import ProvablyService service = ProvablyService() diff --git a/src/sourcerykit/cli/init.py b/src/sourcerykit/cli/init.py index 7d47ee5..44fcd76 100644 --- a/src/sourcerykit/cli/init.py +++ b/src/sourcerykit/cli/init.py @@ -6,7 +6,19 @@ import questionary import typer +from provably import ( + Organization, + OrganizationType, + ProvablyConnectionError, + ProvablyUnauthorizedError, + browser_login, +) +from provably._api import get_api as get_main_api +from provably._http import get_http +from provably.auth_service import ProvablyAuthService +from provably.service import service as provably_service +from sourcerykit._provably import OAUTH_CLIENT_ID, OAUTH_LOOPBACK_PORT, consent_page_url from sourcerykit.bootstrap._cache import _BOOTSTRAP_INSTANCE from sourcerykit.cli import logo from sourcerykit.cli.utils import ( @@ -20,16 +32,6 @@ from sourcerykit.config import load_app_dir_config, save_app_dir_config, save_local_env from sourcerykit.db._engine import get_engine from sourcerykit.db._schema import ensure_schema -from sourcerykit.provably._api import get_api as get_main_api -from sourcerykit.provably._auth_api import Organization, OrganizationType -from sourcerykit.provably._errors import ( - ProvablyConnectionError, - ProvablyUnauthorizedError, -) -from sourcerykit.provably._http import get_http -from sourcerykit.provably.auth_service import ProvablyAuthService -from sourcerykit.provably.oauth_login import browser_login -from sourcerykit.provably.service import service as provably_service service = ProvablyAuthService() @@ -48,7 +50,7 @@ def _run_oauth_browser( try: console.print("\n[bold]🔐 Browser authentication[/bold]") console.print("Opening browser for authentication...") - tokens = asyncio.run(browser_login()) + tokens = asyncio.run(browser_login(consent_page_url(), client_id=OAUTH_CLIENT_ID, port=OAUTH_LOOPBACK_PORT)) email = asyncio.run(service.get_user_email(tokens.access_token)) except ProvablyConnectionError as e: console.print(f"[red]❌ Network error: {e}[/red]") diff --git a/src/sourcerykit/cli/sandbox.py b/src/sourcerykit/cli/sandbox.py index 28c00ad..965f87d 100644 --- a/src/sourcerykit/cli/sandbox.py +++ b/src/sourcerykit/cli/sandbox.py @@ -5,10 +5,10 @@ import questionary import typer from dotenv import unset_key +from provably.service import service from sourcerykit.cli.utils import console, mask_postgres_url, require_settings from sourcerykit.config import LOCAL_ENV_FILE, save_local_env -from sourcerykit.provably.service import service sandbox = typer.Typer(no_args_is_help=True) diff --git a/src/sourcerykit/cli/trace.py b/src/sourcerykit/cli/trace.py index e58d448..25f821b 100644 --- a/src/sourcerykit/cli/trace.py +++ b/src/sourcerykit/cli/trace.py @@ -6,6 +6,8 @@ from uuid import UUID import typer +from provably import QueryAnswer +from provably.service import service from rich.markup import escape as rich_escape from rich.panel import Panel from rich.table import Table @@ -19,8 +21,6 @@ select_trace_intercepts_by_trace_id, select_traces_with_intercept_count, ) -from sourcerykit.provably._answer_model import QueryAnswer -from sourcerykit.provably.service import service from sourcerykit.utils import extract_actual trace = typer.Typer(no_args_is_help=True) diff --git a/src/sourcerykit/cli/utils.py b/src/sourcerykit/cli/utils.py index b9df64b..a67f3d8 100644 --- a/src/sourcerykit/cli/utils.py +++ b/src/sourcerykit/cli/utils.py @@ -7,11 +7,11 @@ import psycopg import questionary import typer +from provably._api import get_api as get_main_api +from provably._http import get_http from rich.console import Console from sourcerykit.config import CONFIG_FILE, Settings, get_settings, load_app_dir_config -from sourcerykit.provably._api import get_api as get_main_api -from sourcerykit.provably._http import get_http console = Console() diff --git a/src/sourcerykit/config.py b/src/sourcerykit/config.py index 7a7c7b4..9d273bd 100644 --- a/src/sourcerykit/config.py +++ b/src/sourcerykit/config.py @@ -218,3 +218,11 @@ def get_bootstrap_settings() -> str: if not raw: raw = load_app_dir_config().get("provably_api", "").strip().rstrip("/") return raw or DEFAULT_PROVABLY_API_URL + + +def get_bootstrap_app_url() -> str: + """Return the Provably app URL without requiring full settings validation.""" + raw = (os.getenv("SOURCERYKIT_PROVABLY_APP_URL") or "").strip().rstrip("/") + if not raw: + raw = str(load_app_dir_config().get("provably_app", "")).strip().rstrip("/") + return raw or DEFAULT_PROVABLY_APP_URL diff --git a/src/sourcerykit/db/_engine.py b/src/sourcerykit/db/_engine.py index 2c9c52e..82a7705 100644 --- a/src/sourcerykit/db/_engine.py +++ b/src/sourcerykit/db/_engine.py @@ -1,9 +1,6 @@ """SQLAlchemy engine""" -from dataclasses import dataclass -from typing import Any -from urllib.parse import unquote, urlparse - +from provably import ConnectionInfo from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine from sourcerykit.config import get_settings @@ -13,44 +10,6 @@ _log = get_logger(__name__) -@dataclass -class ConnectionInfo: - name: str - username: str - password: str - provider: str - uri: str - - @classmethod - def from_url(cls, url: str) -> "ConnectionInfo": - """Parse a PostgreSQL URL into a ConnectionInfo.""" - parsed = urlparse(url) - provider = parsed.scheme.split("+", 1)[0] - host = parsed.hostname or "" - port = parsed.port - uri = f"{host}:{port}" if port else host - return cls( - name=parsed.path.lstrip("/"), - username=unquote(parsed.username or ""), - password=unquote(parsed.password or ""), - provider=provider, - uri=uri, - ) - - def same_server(self, other: "ConnectionInfo") -> bool: - """True if both point to the same database server and name (ignores credentials and query params).""" - return (self.provider, self.uri, self.name) == (other.provider, other.uri, other.name) - - def to_dict(self) -> dict[str, Any]: - return { - "name": self.name, - "username": self.username, - "password": self.password, - "provider": self.provider, - "uri": self.uri, - } - - # Internal singleton to ensure we only ever create one engine per process _ENGINE: AsyncEngine | None = None diff --git a/src/sourcerykit/evaluator/evaluator.py b/src/sourcerykit/evaluator/evaluator.py index 42c2778..552ddd9 100644 --- a/src/sourcerykit/evaluator/evaluator.py +++ b/src/sourcerykit/evaluator/evaluator.py @@ -2,14 +2,15 @@ import uuid from typing import Any +from provably import ProvablyError, QueryAnswer +from provably.service import service + from sourcerykit.db._engine import get_engine from sourcerykit.db._traces import update_trace_intercept_outcome from sourcerykit.errors import SourceryKitError, SourceryKitStorageError from sourcerykit.evaluator._eval_modes import evaluate_claim from sourcerykit.intercept._self_egress import provably_self_egress from sourcerykit.logger import get_logger -from sourcerykit.provably._answer_model import QueryAnswer -from sourcerykit.provably.service import service from sourcerykit.schemas import HandoffPayload, Outcome from sourcerykit.trusted_endpoints.service import verify_claim_endpoints @@ -29,7 +30,7 @@ async def evaluate_handoff(*, payload: HandoffPayload) -> dict[str, Any]: verify_claim_endpoints(payload), asyncio.gather(*(service.verify_proof(qid, integration_api_key) for qid in query_ids)), ) - except (ValueError, SourceryKitError) as e: + except (ValueError, SourceryKitError, ProvablyError) as e: return {"outcome": Outcome.CAUGHT, "per_claim": [], "errors": [f"trust gate: {e}"]} per_claim: list[dict[str, Any]] = [] diff --git a/src/sourcerykit/handoff/_preprocess.py b/src/sourcerykit/handoff/_preprocess.py index bc573fe..7fa0943 100644 --- a/src/sourcerykit/handoff/_preprocess.py +++ b/src/sourcerykit/handoff/_preprocess.py @@ -1,10 +1,11 @@ import asyncio +from provably import ProvablyNotFoundError +from provably.service import service + from sourcerykit.bootstrap.bootstrap import get_bootstrap from sourcerykit.errors import SourceryKitBootstrapError from sourcerykit.logger import get_logger -from sourcerykit.provably._errors import ProvablyNotFoundError -from sourcerykit.provably.service import service _log = get_logger(__name__) _preprocess_lock = asyncio.Lock() diff --git a/src/sourcerykit/handoff/_query_records.py b/src/sourcerykit/handoff/_query_records.py index 47a3e27..fe02db5 100644 --- a/src/sourcerykit/handoff/_query_records.py +++ b/src/sourcerykit/handoff/_query_records.py @@ -1,5 +1,7 @@ from uuid import UUID +from provably.service import service + from sourcerykit.bootstrap.bootstrap import get_bootstrap from sourcerykit.db._intercepts import ( select_intercept_by_call_ref, @@ -7,7 +9,6 @@ ) from sourcerykit.errors import SourceryKitBootstrapError from sourcerykit.logger import get_logger -from sourcerykit.provably.service import service from sourcerykit.utils.validation import validate_length _log = get_logger(__name__) diff --git a/src/sourcerykit/provably/__init__.py b/src/sourcerykit/provably/__init__.py deleted file mode 100644 index 8a519d0..0000000 --- a/src/sourcerykit/provably/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -"""Provably backend HTTP client (httpx-based).""" - -__all__: list[str] = [] # internal HTTP client — not part of the public API diff --git a/src/sourcerykit/provably/_answer_model.py b/src/sourcerykit/provably/_answer_model.py deleted file mode 100644 index 9119a3f..0000000 --- a/src/sourcerykit/provably/_answer_model.py +++ /dev/null @@ -1,79 +0,0 @@ -"""Data validation and serialization contracts for the Rust API layers.""" - -from typing import Any, Literal - -import msgspec -from pydantic import BaseModel, Field, model_validator - -from sourcerykit.logger import get_logger - -_log = get_logger(__name__) - -__all__ = ["QueryAnswer", "TabularData", "AggregateAnswer", "ResultsetAnswer"] - - -def _safe_deserialize(cell: Any) -> Any: - """Helper to cleanly parse stringified JSON containers using high-speed C decoding.""" - if isinstance(cell, str) and cell.strip().startswith(("{", "[")): - try: - return msgspec.json.decode(cell) - except Exception as e: - _log.debug("safe_deserialize_fallback", error=str(e)) - return cell - return cell - - -class TabularData(BaseModel): - columns: list[dict[str, Any]] = Field(default_factory=list) - rows: list[list[Any]] = Field(default_factory=list) - - def extract_value(self) -> Any: - """Unpacks tabular results down to the core target payload.""" - if not self.rows: - return {"columns": self.columns, "rows": self.rows} - - row0 = self.rows[0] - # Extracted column names are normalized to lowercase - col_names = [str(c.get("name") or "").lower() for c in self.columns] - - # Target 1: Extract intercept record payloads cleanly if present - if "raw_response" in col_names: - return _safe_deserialize(row0[col_names.index("raw_response")]) - - # Target 2: Single cell scalar - if len(col_names) == 1 and len(row0) == 1: - return _safe_deserialize(row0[0]) - - return {"columns": self.columns, "rows": self.rows} - - -class AggregateAnswer(BaseModel): - type: Literal["aggregate"] - value: str - - def extract_value(self) -> Any: - return _safe_deserialize(self.value) - - -class ResultsetAnswer(BaseModel): - type: Literal["resultset"] - value: TabularData - - def extract_value(self) -> Any: - return self.value.extract_value() - - -class QueryAnswer(BaseModel): - """Wrapper mapping the structure of QueryAnswer enum.""" - - root: AggregateAnswer | ResultsetAnswer = Field(..., discriminator="type") - - @model_validator(mode="before") - @classmethod - def wrap_root(cls, data: Any) -> Any: - if isinstance(data, dict) and "root" not in data: - return {"root": data} - return data - - def flatten(self) -> Any: - return self.root.extract_value() diff --git a/src/sourcerykit/provably/_api.py b/src/sourcerykit/provably/_api.py deleted file mode 100644 index 40d1175..0000000 --- a/src/sourcerykit/provably/_api.py +++ /dev/null @@ -1,405 +0,0 @@ -"""Provably API — named methods for every endpoint. - -:class:`ProvablyAPI` covers five resource groups: -- **Middlewares** — create the Provably middleware for an org -- **Databases / Schemas / Tables / Columns** — onboard and inspect the connected database -- **Collections** — manage query collections -- **Integrations** — register external integrations -- **Queries & Proofs** — run queries, generate proofs, poll status -""" - -import functools -import uuid -from typing import Any - -from sourcerykit.config import Settings, get_settings -from sourcerykit.provably._http import get_http - - -class ProvablyAPI: - """Provably API endpoints.""" - - def __init__(self, settings: Settings | None = None) -> None: - s = settings or get_settings() - self.org_id = s.org_id - self.app = s.provably_app - - def _generic_path(self) -> str: - return "/api/v1" - - def _org_path(self) -> str: - return f"/api/v1/organizations/{self.org_id}" - - # ------------------------------------------------------------------ - # Feedback - # ------------------------------------------------------------------ - async def create_feedback(self, body: dict[str, Any], files: dict[str, Any] | None = None) -> None: - """ - Send a feedback. - """ - path = "/api/v1/feedback" - - await get_http().post_multipart(path, body, files=files) - return - - # ------------------------------------------------------------------ - # Organizations - # ------------------------------------------------------------------ - - async def list_organizations(self) -> list[dict[str, Any]]: - """ - List all organizations accessible to the authenticated user (via API key). - - Returns: - list[dict[str, Any]]: List of organization objects (each contains at least ``id`` and ``name``). - """ - path = "/api/v1/organizations" - result: list[dict[str, Any]] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Sandboxes - # ------------------------------------------------------------------ - - async def create_sandbox(self, org_id: uuid.UUID, *, token: str | None = None) -> dict[str, Any]: - """ - Create a hosted sandbox database for the given organisation. - - Args: - org_id: The ID of the organisation that owns the sandbox. - token: Optional JWT token for authentication (used during init). - - Returns: - dict[str, Any]: Sandbox record with ``status`` and ``connection_uri``. - """ - path = "/api/v1/sandboxes" - result: dict[str, Any] = await get_http().post(path, {"org_id": str(org_id)}, token=token) - return result - - async def get_sandbox(self, *, token: str | None = None) -> dict[str, Any]: - """ - Retrieve the current sandbox for the authenticated user. - - Args: - token: Optional JWT token for authentication (used during init). - - Returns: - dict[str, Any]: Sandbox record with ``status`` and ``connection_uri``. - """ - path = "/api/v1/sandboxes" - result: dict[str, Any] = await get_http().get(path, token=token) - return result - - async def delete_sandbox(self, *, token: str | None = None) -> None: - """ - Delete the sandbox for the authenticated user. - - Args: - token: Optional JWT token for authentication (used during init). - """ - path = "/api/v1/sandboxes" - await get_http().delete(path, token=token) - - # ------------------------------------------------------------------ - # Middlewares - # ------------------------------------------------------------------ - - async def create_middleware(self) -> dict[str, Any]: - """ - Create the Provably middleware for the configured org. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares/provably" - - result: dict[str, Any] = await get_http().post(path) - return result - - async def list_middlewares(self) -> list[dict[str, Any]]: - """ - List all middlewares. - - Returns: - Any: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares" - - result: list[dict[str, Any]] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Databases - # ------------------------------------------------------------------ - - async def create_database(self, middleware_id: uuid.UUID, body: dict[str, Any]) -> dict[str, Any]: - """ - Onboard a database to a middleware. - - Args: - middleware_id: The ID of the middleware to attach the database to. - body: The database connection payload. - - Returns: - httpx.Response: The raw HTTP response from the API. - """ - path = f"{self._org_path()}/middlewares/{middleware_id}/databases" - - result: dict[str, Any] = await get_http().post(path, body) - return result - - async def list_databases(self, middleware_id: uuid.UUID) -> list[dict[str, Any]]: - """ - List all databases attached to a middleware. - - Args: - middleware_id: The ID of the middleware to query. - - Returns: - Any: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares/{middleware_id}/databases" - - result: list[dict[str, Any]] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Schemas / Tables / Columns - # ------------------------------------------------------------------ - - async def list_columns_from_database( - self, - middleware_id: uuid.UUID, - database_id: uuid.UUID, - schema_id: uuid.UUID, - table_id: uuid.UUID, - ) -> list[dict[str, Any]]: - """ - List all columns in a table. - - Args: - middleware_id: The ID of the middleware owning the database. - database_id: The ID of the database containing the schema. - schema_id: The ID of the schema containing the table. - table_id: The ID of the table to inspect. - - Returns: - list[dict[str, Any]]: The raw JSON response from the API. - """ - path = ( - f"{self._org_path()}/middlewares/{middleware_id}" - f"/databases/{database_id}/schemas/{schema_id}/tables/{table_id}/columns" - ) - result: list[dict[str, Any]] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Data - # ------------------------------------------------------------------ - - async def get_data(self) -> dict[str, Any]: - """ - Retrieve data for the configured org. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/data" - result: dict[str, Any] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Collections - # ------------------------------------------------------------------ - - async def list_collections(self) -> list[dict[str, Any]]: - """ - List all collections for the configured org. - - Returns: - list[dict[str, Any]]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/collections" - result: list[dict[str, Any]] = await get_http().get(path) - return result - - async def create_collection(self, body: dict[str, Any]) -> dict[str, Any]: - """ - Create a new collection for the configured org. - - Args: - body: The collection creation payload. - - Returns: - httpx.Response: The raw HTTP response from the API. - """ - path = f"{self._org_path()}/collections" - - result: dict[str, Any] = await get_http().post(path, body) - return result - - # ------------------------------------------------------------------ - # Integrations - # ------------------------------------------------------------------ - - async def ensure_integration(self, body: dict[str, Any]) -> dict[str, Any]: - """ - Idempotent get-or-create an integration for the configured org. - - Reuses an existing enabled integration with the same name linked to the - exact requested collections, returning its full key instead of minting a - duplicate. - - Args: - body: The integration registration payload. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/integrations/ensure" - - result: dict[str, Any] = await get_http().post(path, body) - return result - - async def list_integrations(self, query: str | None = None) -> list[dict[str, Any]]: - """ - List all integrations for the configured org. - - Args: - query: Optional search string to filter integrations by name. - - Returns: - list[dict[str, Any]]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/integrations" - params = {"query": query} if query is not None else None - result: list[dict[str, Any]] = await get_http().get(path, params=params) - return result - - # ------------------------------------------------------------------ - # Preprocess - # ------------------------------------------------------------------ - - async def start_preprocess(self, middleware_id: uuid.UUID, table_id: uuid.UUID) -> dict[str, Any]: - """ - Start a preprocessing job for a table. - - Args: - middleware_id: The ID of the middleware owning the table. - table_id: The ID of the table to preprocess. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares/{middleware_id}/tables/{table_id}/preprocess" - - result: dict[str, Any] = await get_http().post(path, {"force": True}) - return result - - async def get_preprocess_status(self, middleware_id: uuid.UUID, table_id: uuid.UUID) -> dict[str, Any]: - """ - Get the preprocessing status for a table. - - Args: - middleware_id: The ID of the middleware owning the table. - table_id: The ID of the table to check. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares/{middleware_id}/tables/{table_id}/preprocess" - - result: dict[str, Any] = await get_http().get(path) - return result - - # ------------------------------------------------------------------ - # Queries / Proofs - # ------------------------------------------------------------------ - - async def run_query(self, middleware_id: uuid.UUID, collection_id: uuid.UUID, sql: str) -> dict[str, Any]: - """ - Run a SQL query through a middleware and request a proof. - - Args: - middleware_id: The ID of the middleware to execute the query against. - collection_id: The ID of the collection to associate the query with. - sql: The SQL query string to execute. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/middlewares/{middleware_id}/query" - - result: dict[str, Any] = await get_http().post( - path, {"query": sql, "require_proof": True, "collection_id": str(collection_id)} - ) - return result - - async def get_query(self, query_id: uuid.UUID, *, api_key: str | None = None) -> dict[str, Any]: - """ - Retrieve a query record by ID. - - Args: - query_id: The ID of the query to retrieve. - api_key: Optional API key override for this request. - - Returns: - dict[str, Any]: The raw JSON response from the API. - """ - path = f"{self._org_path()}/queries/{query_id}" - - result: dict[str, Any] = await get_http().get(path, api_key=api_key) - return result - - async def get_query_proof(self, proof_id: uuid.UUID, *, api_key: str | None = None) -> bytes: - """ - Download the raw proof data by proof ID. - - Args: - proof_id: The ID of the proof to retrieve. - api_key: Optional API key override for this request. - - Returns: - bytes: The raw proof content. - """ - path = f"{self._generic_path()}/proof_requests/{proof_id}/download" - - return await get_http().get_raw(path, api_key=api_key) - - async def verify_proof(self, query_id: uuid.UUID, *, api_key: str | None = None) -> dict[str, Any]: - """ - Request verification for an existing query proof. - - Args: - query_id: The unique identifier of the query whose proof needs - verification. - api_key: Optional API key override for this request. - - Returns: - dict[str, Any]: A response confirming the verification - task has been successfully initiated. - - Raises: - ProvablyAPIError: If the query does not exist or verification - cannot be initiated. - """ - path = f"{self._org_path()}/queries/{query_id}/verify" - - result: dict[str, Any] = await get_http().post(path, {}, api_key=api_key) - return result - - # ------------------------------------------------------------------ - # URL helpers - # ------------------------------------------------------------------ - - def query_record_url(self, query_record_id: uuid.UUID) -> str: - """Provably Data Admin URL for a query record.""" - if not query_record_id: - raise ValueError("query_record_id is required") - return f"{self.app}/org/{self.org_id}/query-record/{query_record_id}" - - -@functools.lru_cache(maxsize=1) -def get_api() -> ProvablyAPI: - """Return the shared :class:`ProvablyAPI`, constructed on first call.""" - return ProvablyAPI() diff --git a/src/sourcerykit/provably/_auth_api.py b/src/sourcerykit/provably/_auth_api.py deleted file mode 100644 index 6337220..0000000 --- a/src/sourcerykit/provably/_auth_api.py +++ /dev/null @@ -1,186 +0,0 @@ -"""Provably Auth API — OAuth tokens, user, API key and organisation endpoints. - -:class:`ProvablyAuthAPI` covers three resource groups: -- **OAuth** — exchange an authorization code and rotate refresh tokens -- **User** — retrieve the current authenticated user -- **Organisations** — create and list organisations - -The browser OAuth flow itself lives in -:mod:`sourcerykit.provably.oauth_login`. -""" - -import functools -from dataclasses import dataclass -from enum import StrEnum -from typing import Any - -from sourcerykit.provably._http import ProvablyHTTPClient - -OAUTH_CLIENT_ID = "sourcerykit-cli" -OAUTH_SCOPE = "read write" -LOOPBACK_PORT = 8910 -REDIRECT_URI = f"http://127.0.0.1:{LOOPBACK_PORT}/callback" - - -@dataclass(slots=True) -class OAuthTokens: - """Tokens issued by the OAuth token endpoint.""" - - access_token: str - refresh_token: str | None - - -class OrganizationType(StrEnum): - DEMOGRAPHICS = "demographics" - E_COMMERCE = "e-commerce" - SOCIAL_MEDIA = "social-media" - HEALTH_AND_FITNESS = "health-and-fitness" - CLIMATE_AND_WEATHER = "climate-and-weather" - EDUCATION = "education" - FINANCIAL = "financial" - REAL_ESTATE = "real-estate" - ENERGY_CONSUMPTION = "energy-consumption" - SPORTS = "sports" - RETAIL = "retail" - HEALTHCARE = "healthcare" - CRYPTOCURRENCY = "cryptocurrency" - GOVERNMENT = "government" - ENTERTAINMENT = "entertainment" - - -@dataclass(slots=True) -class Organization: - handle: str - name: str - organization_type: OrganizationType - - -class ProvablyAuthAPI: - """Provably Auth API endpoints.""" - - def __init__(self) -> None: - self._http = ProvablyHTTPClient(pre_auth=True) - - def _user_path(self) -> str: - return "/api/v1/user" - - def _org_path(self) -> str: - return "/api/v1/organizations" - - # ------------------------------------------------------------------ - # OAuth - # ------------------------------------------------------------------ - - async def exchange_code(self, code: str, verifier: str) -> dict[str, Any]: - """ - Exchange an authorization code for tokens (public client, no secret). - - Args: - code: The authorization code from the redirect. - verifier: The PKCE code verifier. - - Returns: - dict[str, Any]: The raw JSON response (contains ``access_token`` - and optionally ``refresh_token``). - """ - path = "/api/v1/auth/oauth/token" - - result: dict[str, Any] = await self._http.post_form( - path, - { - "grant_type": "authorization_code", - "code": code, - "redirect_uri": REDIRECT_URI, - "client_id": OAUTH_CLIENT_ID, - "code_verifier": verifier, - }, - ) - return result - - async def refresh_tokens(self, refresh_token: str) -> dict[str, Any]: - """ - Rotate tokens: exchange a refresh token for a new access+refresh pair. - - Args: - refresh_token: The refresh token to redeem. - - Returns: - dict[str, Any]: The raw JSON response (contains ``access_token`` - and optionally ``refresh_token``). - """ - path = "/api/v1/auth/oauth/refresh" - - result: dict[str, Any] = await self._http.post_form( - path, - { - "grant_type": "refresh_token", - "refresh_token": refresh_token, - "client_id": OAUTH_CLIENT_ID, - }, - ) - return result - - # ------------------------------------------------------------------ - # User - # ------------------------------------------------------------------ - - async def get_current_user(self, token: str) -> dict[str, Any]: - """ - Retrieve the current authenticated user. - - Args: - token: OAuth access token (Bearer). - - Returns: - dict[str, Any]: The raw JSON response from the API (contains ``email``). - """ - path = f"{self._user_path()}/current" - - result: dict[str, Any] = await self._http.get(path, token=token) - return result - - # ------------------------------------------------------------------ - # Organization - # ------------------------------------------------------------------ - - async def create_organization(self, token: str, organization: Organization) -> dict[str, Any]: - """ - Create a new organisation. - - Args: - token: JWT Bearer token obtained from ``login``. - organization: The organisation details. - - Returns: - dict[str, Any]: The raw JSON response from the API (contains ``id``). - """ - payload = { - "handle": organization.handle, - "name": organization.name, - "type": organization.organization_type.value, - } - path = f"{self._org_path()}" - - result: dict[str, Any] = await self._http.post_multipart(path, payload, token=token) - return result - - async def get_organizations(self, token: str) -> list[dict[str, Any]]: - """ - List all organisations accessible to the authenticated user. - - Args: - token: JWT Bearer token obtained from ``login``. - - Returns: - list[dict[str, Any]]: List of organisation objects (each contains at least ``id`` and ``name``). - """ - path = f"{self._org_path()}" - - result: list[dict[str, Any]] = await self._http.get(path, token=token) - return result - - -@functools.lru_cache(maxsize=1) -def get_api() -> ProvablyAuthAPI: - """Return the shared :class:`ProvablyAuthAPI`, constructed on first call.""" - return ProvablyAuthAPI() diff --git a/src/sourcerykit/provably/_errors.py b/src/sourcerykit/provably/_errors.py deleted file mode 100644 index 3247d75..0000000 --- a/src/sourcerykit/provably/_errors.py +++ /dev/null @@ -1,164 +0,0 @@ -from collections.abc import AsyncIterator -from contextlib import asynccontextmanager - -import httpx - -from sourcerykit.errors import SourceryKitError -from sourcerykit.logger import get_logger - -_log = get_logger(__name__) - - -class ProvablyError(SourceryKitError): - """Base exception for all SourceryKit Provably errors.""" - - pass - - -class ProvablyAPIError(ProvablyError): - """Raised when the Provably API returns a 4xx or 5xx response.""" - - def __init__(self, message: str, status_code: int | None = None, response_body: str | None = None): - super().__init__(message) - self.status_code = status_code - self.response_body = response_body - - -class ProvablyConnectionError(ProvablyError): - """Raised when the Provably API is unreachable (Network/Timeout).""" - - pass - - -class ProvablyDataError(ProvablyError): - """Raised when the API response is malformed or invalid.""" - - pass - - -class ProvablyAuthError(ProvablyAPIError): - """Base exception for authentication / account errors.""" - - pass - - -class ProvablyResourceAlreadyExistsError(ProvablyAuthError): - """Raised when a resource (account, organisation handle, etc.) already exists.""" - - pass - - -class ProvablyUnauthorizedError(ProvablyAuthError): - """Raised when credentials are invalid or the request is unauthorized (HTTP 401).""" - - pass - - -class ProvablyNotFoundError(ProvablyAPIError): - """Raised when the requested resource does not exist (HTTP 404).""" - - pass - - -@asynccontextmanager -async def provably_auth_error_handler(service: str) -> AsyncIterator[None]: - """ - Standardizes error handling and logging across auth service methods. - - Args: - service: A slug representing the operation. - """ - service_name = service.replace("_", " ") - try: - yield - except httpx.HTTPStatusError as e: - status = e.response.status_code - body = e.response.text - _log.error( - f"provably_auth_rejected_{service}", - status_code=status, - path=str(e.request.url), - response=body[:500], - ) - if status == 400: - try: - e.response.json().get("description", "") - except Exception: - raise - if status == 401: - raise ProvablyUnauthorizedError( - message=f"Invalid credentials for {service_name}.", - status_code=status, - response_body=body, - ) from e - if status == 409: - raise ProvablyResourceAlreadyExistsError( - message=f"Resource already exists for {service_name}: {body}", - status_code=status, - response_body=body, - ) from e - raise ProvablyAuthError( - message=f"Provably API rejected {service_name}: {body}", - status_code=status, - response_body=body, - ) from e - - except (ValueError, TypeError, KeyError) as e: - _log.error(f"provably_data_invalid_{service}", error=str(e)) - raise ProvablyDataError(f"Provably API returned invalid data for {service_name}: {e}") from e - - except httpx.RequestError as e: - _log.error(f"provably_network_unreachable_{service}", error=str(e)) - raise ProvablyConnectionError(f"Could not reach Provably API to perform {service_name}.") from e - - except Exception as e: - _log.error(f"provably_unexpected_error_{service}", error=str(e)) - raise ProvablyError(f"Unexpected error during {service_name}: {e}") from e - - -@asynccontextmanager -async def provably_error_handler(service: str) -> AsyncIterator[None]: - """ - Standardizes error handling and logging across service methods. - - Args: - service: A slug representing the operation. - """ - service_name = service.replace("_", " ") - try: - yield - except httpx.HTTPStatusError as e: - status = e.response.status_code - body = e.response.text - # Log API failures - _log.error( - f"provably_api_rejected_{service}", - status_code=status, - path=str(e.request.url), - response=body[:500], - ) - if status == 404: - raise ProvablyNotFoundError( - message=f"Resource not found for {service_name}: {body}", - status_code=status, - response_body=body, - ) from e - raise ProvablyAPIError( - message=f"Provably API rejected {service_name}: {body}", - status_code=status, - response_body=body, - ) from e - - except (ValueError, TypeError, KeyError) as e: - # Log data corruption or unexpected schema changes - _log.error(f"provably_data_invalid_{service}", error=str(e)) - raise ProvablyDataError(f"Provably API returned invalid data for {service_name}: {e}") from e - - except httpx.RequestError as e: - # Log network-level failures - _log.error(f"provably_network_unreachable_{service}", error=str(e)) - raise ProvablyConnectionError(f"Could not reach Provably API to perform {service_name}.") from e - - except Exception as e: - _log.error(f"provably_unexpected_error_{service}", error=str(e)) - raise ProvablyError(f"Unexpected error during {service_name}: {e}") from e diff --git a/src/sourcerykit/provably/_http.py b/src/sourcerykit/provably/_http.py deleted file mode 100644 index 27fbfe5..0000000 --- a/src/sourcerykit/provably/_http.py +++ /dev/null @@ -1,289 +0,0 @@ -"""HTTP client for the Provably API.""" - -import asyncio -import functools -import json -import os -from typing import Any - -import httpx -from dotenv import set_key - -from sourcerykit.config import ( - CONFIG_FILE, - Settings, - get_bootstrap_settings, - get_settings, - load_app_dir_config, - load_local_env, - save_app_dir_config, -) -from sourcerykit.intercept._self_egress import provably_self_egress -from sourcerykit.logger import get_logger - -_log = get_logger(__name__) - - -def _token_store() -> str | None: - """Return the app's token-store path (SOURCERYKIT_TOKEN_STORE) or None for global JSON. - - Resolved per-call so the app can set it at startup regardless of import order. - """ - return os.getenv("SOURCERYKIT_TOKEN_STORE") or None - - -def _persist_tokens(access_token: str, refresh_token: str | None) -> None: - """Persist rotated tokens to the configured store (app .env or global JSON).""" - store = _token_store() - if store: - set_key(store, "PROVABLY_ACCESS_TOKEN", access_token) - set_key(store, "PROVABLY_REFRESH_TOKEN", refresh_token or "") - os.environ["PROVABLY_ACCESS_TOKEN"] = access_token - if refresh_token: - os.environ["PROVABLY_REFRESH_TOKEN"] = refresh_token - load_local_env.cache_clear() - else: - save_app_dir_config(token=access_token, refresh_token=refresh_token) - get_settings.cache_clear() - - -def _clear_refresh_token() -> None: - """Drop the stored refresh token from the configured store.""" - store = _token_store() - if store: - set_key(store, "PROVABLY_REFRESH_TOKEN", "") - os.environ.pop("PROVABLY_REFRESH_TOKEN", None) - load_local_env.cache_clear() - else: - payload = load_app_dir_config() - payload.pop("refresh_token", None) - CONFIG_FILE.write_text(json.dumps(payload)) - load_app_dir_config.cache_clear() - get_settings.cache_clear() - - -async def _refresh_session() -> str | None: - """Rotate the stored OAuth refresh token once; return the new access token. - - Returns None (and clears the stored refresh token) when no refresh token - is configured or rotation fails. - """ - from sourcerykit.provably.auth_service import auth_service - - refresh = get_settings().refresh_token - if not refresh: - return None - try: - tokens = await auth_service.refresh_tokens(str(refresh)) - except Exception: - _log.warning("oauth_refresh_failed", detail="dropping stored refresh token") - _clear_refresh_token() - return None - - _persist_tokens(tokens.access_token, tokens.refresh_token) - return tokens.access_token - - -class ProvablyHTTPClient: - """Httpx wrapper for the Provably API. - - All requests are wrapped in ``provably_self_egress()`` so SDK-internal - traffic bypasses the trust gate and the intercept recorder. - """ - - def __init__(self, settings: Settings | None = None, *, pre_auth: bool = False) -> None: - self._client: httpx.AsyncClient | None = None - self._client_loop: asyncio.AbstractEventLoop | None = None - self._post_auth = not pre_auth - - if pre_auth: - self.base_url = get_bootstrap_settings() - self._headers = {"Content-Type": "application/json"} - else: - s = settings or get_settings() - self.base_url = s.provably_api.rstrip("/") - self._headers = { - "Content-Type": "application/json", - } - - def _get_client(self) -> httpx.AsyncClient: - """Return a shared AsyncClient, recreating it when the event loop has changed.""" - try: - loop: asyncio.AbstractEventLoop | None = asyncio.get_running_loop() - except RuntimeError: - loop = None - - if self._client is None or self._client.is_closed or loop is not self._client_loop: - self._client = httpx.AsyncClient() - self._client_loop = loop - - return self._client - - async def _request( - self, - method: str, - path: str, - *, - timeout: float = 60.0, - api_key: str | None = None, - token: str | None = None, - **kwargs: Any, - ) -> httpx.Response: - - headers = {**self._headers} - - if "files" in kwargs or "data" in kwargs: - headers.pop("Content-Type", None) - - if token is not None: - headers["Authorization"] = f"Bearer {token}" - headers.pop("x-api-key", None) - - elif api_key is not None: - headers["x-api-key"] = api_key - - with provably_self_egress(): - return await self._get_client().request( - method, f"{self.base_url}{path}", headers=headers, timeout=timeout, **kwargs - ) - - async def _fetch( - self, - method: str, - path: str, - *, - api_key: str | None = None, - token: str | None = None, - _oauth_retried: bool = False, - **kwargs: Any, - ) -> Any: - # Sentinel kwarg guards the single refresh-retry; never forwarded to httpx. - kwargs.pop("_oauth_retried", None) - _log.debug("provably_api_request", method=method, path=path) - # Post-auth requests authenticate with the OAuth access token unless an - # explicit credential is supplied - if token is None and api_key is None and self._post_auth: - token = get_settings().access_token - try: - response = await self._request(method, path, api_key=api_key, token=token, **kwargs) - response.raise_for_status() - - if not response.content or not response.content.strip(): - return {} - - try: - result = response.json() - _log.debug("provably_api_response_ok", method=method, path=path, status=response.status_code) - return result - except ValueError: - _log.debug( - "provably_api_response_not_json", - method=method, - path=path, - body=response.text[:200], - ) - return {} - - except httpx.HTTPStatusError as e: - if e.response.status_code == 401 and token is not None and not _oauth_retried: - new_token = await _refresh_session() - if new_token is not None: - _log.info("provably_api_token_refreshed", method=method, path=path) - return await self._fetch( - method, path, api_key=api_key, token=new_token, _oauth_retried=True, **kwargs - ) - _log.error( - "provably_api_rejected", - method=method, - path=path, - status_code=e.response.status_code, - body=e.response.text[:500], - ) - raise - except httpx.RequestError as e: - _log.error("provably_api_network_error", method=method, path=path, error=str(e)) - raise - except httpx.HTTPError as e: - _log.error("provably_api_unexpected_error", method=method, path=path, error=str(e)) - raise - - async def get( - self, - path: str, - params: dict[str, Any] | None = None, - *, - api_key: str | None = None, - token: str | None = None, - ) -> Any: - return await self._fetch("GET", path, api_key=api_key, token=token, params=params) - - async def get_raw( - self, - path: str, - *, - api_key: str | None = None, - token: str | None = None, - ) -> bytes: - """GET that returns raw response bytes instead of parsed JSON.""" - _log.debug("provably_api_request_raw", method="GET", path=path) - if token is None and api_key is None and self._post_auth: - token = get_settings().access_token - try: - response = await self._request("GET", path, api_key=api_key, token=token) - response.raise_for_status() - return response.content - except httpx.HTTPStatusError as e: - if e.response.status_code == 401 and token is not None: - new_token = await _refresh_session() - if new_token is not None: - response = await self._request("GET", path, api_key=api_key, token=new_token) - response.raise_for_status() - return response.content - raise - - async def post( - self, - path: str, - json: dict[str, Any] | None = None, - *, - api_key: str | None = None, - token: str | None = None, - ) -> Any: - return await self._fetch("POST", path, api_key=api_key, token=token, json=json or {}) - - async def post_form( - self, - path: str, - data: dict[str, Any], - *, - token: str | None = None, - ) -> Any: - """POST with an ``application/x-www-form-urlencoded`` body.""" - return await self._fetch("POST", path, token=token, data=data) - - async def post_multipart( - self, - path: str, - data: dict[str, Any], - *, - files: dict[str, Any] | None = None, - token: str | None = None, - ) -> Any: - processed_payload = {key: (None, str(value)) for key, value in data.items()} - if files: - processed_payload.update(files) - return await self._fetch("POST", path, token=token, files=processed_payload) - - async def delete( - self, - path: str, - *, - token: str | None = None, - ) -> Any: - return await self._fetch("DELETE", path, token=token) - - -@functools.lru_cache(maxsize=1) -def get_http() -> ProvablyHTTPClient: - """Return the shared :class:`ProvablyHTTPClient`, constructed on first call.""" - return ProvablyHTTPClient() diff --git a/src/sourcerykit/provably/auth_service.py b/src/sourcerykit/provably/auth_service.py deleted file mode 100644 index e28c31c..0000000 --- a/src/sourcerykit/provably/auth_service.py +++ /dev/null @@ -1,131 +0,0 @@ -""" -Provably auth service layer -""" - -import uuid -from typing import Any - -from sourcerykit.provably._api import get_api as get_main_api -from sourcerykit.provably._auth_api import OAuthTokens, Organization, get_api -from sourcerykit.provably._errors import provably_auth_error_handler - - -class ProvablyAuthService: - """High-level service for account and organisation management.""" - - # ------------------------------------------------------------------ - # OAuth - # ------------------------------------------------------------------ - - async def exchange_code(self, code: str, verifier: str) -> OAuthTokens: - """Exchange an authorization code for tokens (public client, no secret). - - Args: - code: The authorization code from the redirect. - verifier: The PKCE code verifier. - - Returns: - OAuthTokens: The issued access and refresh tokens. - - Raises: - ProvablyAuthError: On API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("oauth_token_exchange"): - result = await get_api().exchange_code(code, verifier) - return OAuthTokens(access_token=result["access_token"], refresh_token=result.get("refresh_token")) - - async def refresh_tokens(self, refresh_token: str) -> OAuthTokens: - """Rotate tokens: exchange a refresh token for a new access+refresh pair. - - Args: - refresh_token: The refresh token to redeem. - - Returns: - OAuthTokens: The new access and refresh tokens. - - Raises: - ProvablyAuthError: On API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("oauth_refresh"): - result = await get_api().refresh_tokens(refresh_token) - return OAuthTokens(access_token=result["access_token"], refresh_token=result.get("refresh_token")) - - # ------------------------------------------------------------------ - # User - # ------------------------------------------------------------------ - - async def get_user_email(self, token: str) -> str: - """Retrieve the email of the authenticated user. - - Args: - token: OAuth access token (Bearer) from ``browser_login``. - - Returns: - str: The user's email address. - - Raises: - ProvablyAuthError: On API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("get_user_email"): - result = await get_api().get_current_user(token) - return str(result["email"]) - - # ------------------------------------------------------------------ - # Organisation - # ------------------------------------------------------------------ - - async def create_organization(self, token: str, organization: Organization) -> uuid.UUID: - """Create a new organisation and return its ID. - - Args: - token: JWT Bearer token from ``login``. - organization: Organisation details. - - Returns: - uuid.UUID: The ID of the newly created organisation. - - Raises: - ProvablyResourceAlreadyExistsError: If the handle is already taken. - ProvablyAuthError: On other API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("create_organization"): - result = await get_api().create_organization(token, organization) - return uuid.UUID(str(result["id"])) - - async def get_organizations(self, token: str) -> list[dict[str, Any]]: - """List organisations accessible to the authenticated user. - - Args: - token: JWT token from ``login``. - - Returns: - list[dict[str, Any]]: List of organisation. - - Raises: - ProvablyAuthError: On API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("get_organizations"): - result = await get_api().get_organizations(token) - return result - - async def list_organizations(self) -> list[dict[str, Any]]: - """List organisations accessible to the authenticated user (via API key). - - Returns: - list[dict[str, Any]]: List of organisation objects (each contains at least ``id`` and ``name``). - - Raises: - ProvablyAuthError: On API errors. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_auth_error_handler("list_organizations"): - result = await get_main_api().list_organizations() - return result - - -auth_service = ProvablyAuthService() diff --git a/src/sourcerykit/provably/oauth_login.py b/src/sourcerykit/provably/oauth_login.py deleted file mode 100644 index 531adbd..0000000 --- a/src/sourcerykit/provably/oauth_login.py +++ /dev/null @@ -1,202 +0,0 @@ -"""OAuth2 login for the SourceryKit CLI — public client with PKCE (RFC 7636/8252). - -:func:`browser_login` opens ``{provably_app}/consent?…``; the web app drives -sign-in and consent and redirects to this CLI's loopback listener -(``http://127.0.0.1:8910/callback``). The client is public: no secret, PKCE -S256 is the only proof. - -This module is pure OAuth orchestration — the token HTTP calls live in -:mod:`sourcerykit.provably.auth_service`. -""" - -import asyncio -import base64 -import hashlib -import html -import os -import secrets -import threading -import urllib.parse -import webbrowser -from collections.abc import Callable -from http.server import BaseHTTPRequestHandler, HTTPServer - -from sourcerykit.config import DEFAULT_PROVABLY_APP_URL, load_app_dir_config -from sourcerykit.logger import get_logger -from sourcerykit.provably._auth_api import ( - LOOPBACK_PORT, - OAUTH_CLIENT_ID, - OAUTH_SCOPE, - REDIRECT_URI, - OAuthTokens, -) -from sourcerykit.provably.auth_service import auth_service - -_log = get_logger(__name__) - -# How long the browser's request waits for the token exchange to finish. -SETTLE_TIMEOUT = 30.0 - - -def pkce_pair() -> tuple[str, str]: - """Return ``(verifier, challenge)`` using the S256 method.""" - verifier = secrets.token_urlsafe(48) - digest = hashlib.sha256(verifier.encode()).digest() - challenge = base64.urlsafe_b64encode(digest).rstrip(b"=").decode() - return verifier, challenge - - -def consent_page_url() -> str: - """URL of the web app consent page. - - When ``SOURCERYKIT_PROVABLY_APP_URL`` (or ``provably_app``) is set it is - used verbatim as the consent page URL — point it at your app's consent page - in dev. When unset, falls back to the production app's ``/consent`` page. - """ - raw = os.environ.get("SOURCERYKIT_PROVABLY_APP_URL") or "" - if not raw: - raw = str(load_app_dir_config().get("provably_app", "")) - raw = raw.strip().rstrip("/") - return raw or f"{DEFAULT_PROVABLY_APP_URL}/consent" - - -# ---------------------------------------------------------------------- -# Browser flow -# ---------------------------------------------------------------------- - - -class _LoopbackCallbackHandler(BaseHTTPRequestHandler): - """Captures ?code/&state on GET /callback and holds its reply. - - The reply waits until the code has been traded for tokens, so what the - browser reads is the real outcome rather than "received it". The consent - page reads that reply across origins, which needs the allow-origin header. - """ - - result: dict[str, str] = {} - done = threading.Event() - settled = threading.Event() - failure: str | None = None - allowed_origin: str = "" - - def do_GET(self) -> None: # noqa: N802 - stdlib API - parsed = urllib.parse.parse_qsl(urllib.parse.urlsplit(self.path).query) - type(self).result = dict(parsed) - type(self).done.set() - - type(self).settled.wait(SETTLE_TIMEOUT) - - failure = type(self).failure - if not type(self).settled.is_set(): - failure = "the CLI stopped waiting" - - self.send_response(200 if failure is None else 500) - self.send_header("Content-Type", "text/html; charset=utf-8") - if type(self).allowed_origin: - self.send_header("Access-Control-Allow-Origin", type(self).allowed_origin) - self.end_headers() - self.wfile.write(_callback_page(failure)) - - def log_message(self, format: str, *args: object) -> None: # silence stderr - pass - - -def _callback_page(failure: str | None) -> bytes: - """What the browser shows when the CLI has finished with the code.""" - if failure is None: - body = "

Logged in!

You can close this window." - else: - body = f"

Sign-in failed

{html.escape(failure)}" - return f"{body}".encode() - - -def _origin_of(url: str) -> str: - """Scheme and host of a URL, or "" when it has neither.""" - parts = urllib.parse.urlsplit(url) - return f"{parts.scheme}://{parts.netloc}" if parts.scheme and parts.netloc else "" - - -async def _wait_for_loopback_code( - allowed_origin: str = "", - timeout: float = 300.0, -) -> tuple[dict[str, str], Callable[[str | None], None]]: - """Bind the registered loopback port and wait for the OAuth redirect. - - The browser is still waiting for its reply when this returns. Call the - returned ``settle`` with ``None`` once the code has been traded for tokens, - or with a reason when it failed, and the browser is told which. ``settle`` - also stops the listener, so it has to run exactly once. - """ - server = HTTPServer(("127.0.0.1", LOOPBACK_PORT), _LoopbackCallbackHandler) - _LoopbackCallbackHandler.result = {} - _LoopbackCallbackHandler.failure = None - _LoopbackCallbackHandler.allowed_origin = allowed_origin - _LoopbackCallbackHandler.done.clear() - _LoopbackCallbackHandler.settled.clear() - thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.2}, daemon=True) - thread.start() - - def settle(failure: str | None) -> None: - _LoopbackCallbackHandler.failure = failure - _LoopbackCallbackHandler.settled.set() - server.shutdown() - server.server_close() - thread.join(SETTLE_TIMEOUT) - - try: - await asyncio.wait_for(asyncio.to_thread(_LoopbackCallbackHandler.done.wait), timeout) - except BaseException: - settle("the CLI stopped waiting") - raise - - return _LoopbackCallbackHandler.result, settle - - -async def browser_login(consent_page: str | None = None, *, machine_id: str | None = None) -> OAuthTokens: - """Open the web app's consent page and complete the loopback redirect flow. - - Args: - consent_page: Full URL of the consent page. Defaults to - :func:`consent_page_url`. - machine_id: Opaque, stable id of this machine. When given it rides - along in the consent page address, so the page can tell whether - the machine is already paired. Never the raw hardware id. - """ - verifier, challenge = pkce_pair() - state = secrets.token_urlsafe(16) - - query = { - "response_type": "code", - "client_id": OAUTH_CLIENT_ID, - "redirect_uri": REDIRECT_URI, - "scope": OAUTH_SCOPE, - "state": state, - "code_challenge": challenge, - "code_challenge_method": "S256", - } - if machine_id: - query["machine_id"] = machine_id - page = (consent_page or consent_page_url()).split("?", 1)[0] - consent_url = f"{page}?" + urllib.parse.urlencode(query) - - _log.info("oauth_browser_opening", consent_url=consent_url) - print(f"Sign-in page: {consent_url}") - webbrowser.open(consent_url) - - callback, settle = await _wait_for_loopback_code(_origin_of(page)) - - try: - if callback.get("state") != state: - raise ValueError("OAuth state mismatch — aborting") - code = callback.get("code", "") - if not code: - error = callback.get("error", "unknown error") - raise ValueError(f"authorization denied: {error}") - - tokens = await auth_service.exchange_code(code, verifier) - except BaseException as error: - settle(str(error) or error.__class__.__name__) - raise - - settle(None) - return tokens diff --git a/src/sourcerykit/provably/service.py b/src/sourcerykit/provably/service.py deleted file mode 100644 index 94c8c17..0000000 --- a/src/sourcerykit/provably/service.py +++ /dev/null @@ -1,672 +0,0 @@ -""" -Provably service layer -""" - -import asyncio -import uuid -from typing import Any - -from sourcerykit.db._engine import ConnectionInfo -from sourcerykit.db._schema import INTERCEPTS_TABLE -from sourcerykit.logger import get_logger -from sourcerykit.provably._api import get_api -from sourcerykit.provably._errors import ProvablyNotFoundError, provably_error_handler - -_log = get_logger(__name__) - - -class ProvablyService: - """High-level service for managing Provably resources.""" - - # ------------------------------------------------------------------ - # Feedback - # ------------------------------------------------------------------ - async def create_feedback(self, description: str, file: bytes | None) -> None: - """Send a feedback.""" - - feedback_body = {"description": description} - - file_payload = {} - - if file: - file_payload["files"] = ("attachment.dat", file) - - async with provably_error_handler("create_feedback"): - return await get_api().create_feedback(feedback_body, files=file_payload) - - # ------------------------------------------------------------------ - # Sandboxes - # ------------------------------------------------------------------ - - async def create_sandbox(self, org_id: uuid.UUID, *, token: str | None = None) -> str: - """Create a hosted sandbox database for the given organisation. - - Args: - org_id: The ID of the organisation that owns the sandbox. - token: Optional JWT token for authentication (used during init). - - Returns: - str: The connection URI for the new sandbox. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - async with provably_error_handler("create_sandbox"): - result = await get_api().create_sandbox(org_id, token=token) - uri = result.get("connection_uri") - if not uri: - raise ValueError("create_sandbox response missing 'connection_uri'") - return str(uri) - - async def get_sandbox(self, *, token: str | None = None) -> dict[str, Any] | None: - """Retrieve the current sandbox for the authenticated user. - - Args: - token: Optional JWT token for authentication (used during init). - - Returns: - dict[str, Any] | None: Sandbox dict with ``status`` and - ``connection_uri`` keys, or ``None`` if no sandbox exists. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - try: - async with provably_error_handler("get_sandbox"): - return await get_api().get_sandbox(token=token) - except ProvablyNotFoundError: - return None - - async def get_sandbox_connection_uri(self, *, token: str | None = None) -> str | None: - """Return the connection URI of an active sandbox, or ``None``. - - Args: - token: Optional JWT token for authentication (used during init). - - Returns: - str | None: The connection URI if a sandbox exists and is - active, ``None`` otherwise. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - sandbox = await self.get_sandbox(token=token) - if not sandbox: - return None - status = sandbox.get("status", "").lower() - if status in ("active", "provisioning"): - return sandbox.get("connection_uri") - return None - - async def get_sandbox_status(self, postgres_url: str) -> tuple[dict[str, Any] | None, bool]: - """Fetch sandbox and check if *postgres_url* matches its connection URI. - - Returns: - (sandbox_data, is_sandbox) — ``is_sandbox`` is ``True`` when the - configured postgres_url points at the sandbox. - """ - sandbox = await self.get_sandbox() - if not sandbox: - return None, False - uri = sandbox.get("connection_uri") - if not uri: - return None, False - is_sandbox = ConnectionInfo.from_url(postgres_url).same_server(ConnectionInfo.from_url(uri)) - return sandbox, is_sandbox - - async def delete_sandbox(self, *, token: str | None = None) -> None: - """Delete the sandbox for the authenticated user. - - Args: - token: Optional JWT token for authentication (used during init). - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_error_handler("delete_sandbox"): - await get_api().delete_sandbox(token=token) - - # ------------------------------------------------------------------ - # Middleware - # ------------------------------------------------------------------ - - async def create_middleware(self) -> uuid.UUID: - """Create the Provably middleware. - - Returns: - uuid.UUID: The ID of the middleware. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - async with provably_error_handler("create_middleware"): - result = await get_api().create_middleware() - return uuid.UUID(str(result["id"])) - - async def get_middleware_id(self) -> uuid.UUID: - """Find and return the ID of the existing Provably middleware. - - Returns: - uuid.UUID: The ID of the middleware named 'Provably Middleware'. - - Raises: - ValueError: If no middleware with that name exists. - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - - async with provably_error_handler("get_middleware_id"): - middlewares = await get_api().list_middlewares() - - try: - match = next(md for md in middlewares if md.get("name") == "Provably Middleware") - return uuid.UUID(str(match["id"])) - except StopIteration: - raise ValueError("Middleware with name 'Provably Middleware' not found.") - - # ------------------------------------------------------------------ - # Database - # ------------------------------------------------------------------ - - # Will create database, schema and table - async def create_database(self, middleware_id: uuid.UUID, database: ConnectionInfo) -> uuid.UUID: - """Register a new database with the middleware. - - Args: - middleware_id: The ID of the middleware to attach the database to. - database: Connection details for the database to register. - - Returns: - uuid.UUID: The ID of the newly created database. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - - async with provably_error_handler("create_database"): - result = await get_api().create_database(middleware_id, database.to_dict()) - return uuid.UUID(str(result["id"])) - - async def get_database_id(self, middleware_id: uuid.UUID, database: ConnectionInfo) -> uuid.UUID: - """Find and return the ID of an existing database by name. - - Args: - middleware_id: The ID of the middleware owning the database. - database: Connection info whose name is used to locate the database. - - Returns: - uuid.UUID: The ID of the matching database. - - Raises: - ValueError: If no database with the given name exists in the middleware. - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - - async with provably_error_handler("get_database_id"): - databases = await get_api().list_databases(middleware_id) - - try: - match = next(db for db in databases if db.get("name") == database.name) - return uuid.UUID(str(match["id"])) - except StopIteration: - raise ValueError(f"Database with name '{database.name}' not found in middleware {middleware_id}") - - # ------------------------------------------------------------------ - # Collections - # ------------------------------------------------------------------ - - async def create_collection( - self, - middleware_id: uuid.UUID, - database_id: uuid.UUID, - schema_id: uuid.UUID, - table_id: uuid.UUID, - columns: list[uuid.UUID], - name: str, - ) -> uuid.UUID: - """Create a new query collection. - - Args: - middleware_id: The ID of the middleware owning the table. - database_id: The ID of the database containing the table. - schema_id: The ID of the schema containing the table. - table_id: The ID of the table to base the collection on. - columns: List of column IDs to enable for the collection. - name: Collection name (the project name). - - Returns: - uuid.UUID: The ID of the newly created collection. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - - collection = { - "name": name, - "publicity_status": "private", - "middleware_id": str(middleware_id), - "database_id": str(database_id), - "is_descriptions_generated": False, - "entities": [], - "integrations": [], - "query_price": 0, - "is_general_sql_queries_enabled": True, - "schema_id": str(schema_id), - "table_id": str(table_id), - "enabled_columns": [{"id": str(c)} for c in columns], - } - - async with provably_error_handler("create_collection"): - result = await get_api().create_collection(collection) - return uuid.UUID(str(result["id"])) - - async def get_collection_id(self, name: str) -> uuid.UUID: - """Find and return the ID of an existing collection by name. - - Args: - name: The collection name to look up. - - Returns: - uuid.UUID: The ID of the matching collection. - - Raises: - ValueError: If no matching collection is found. - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_error_handler("get_collection_id"): - collections = await get_api().list_collections() - - try: - match = next(collection for collection in collections if collection.get("name") == name) - return uuid.UUID(str(match["id"])) - except StopIteration: - raise ValueError(f"Collection with name '{name}' not found") - - async def list_collections(self) -> list[dict[str, Any]]: - """Return all collections. - - Returns: - list[dict[str, Any]]: Raw collection dicts from the API. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_error_handler("list_collections"): - return await get_api().list_collections() - - # ------------------------------------------------------------------ - # Data - # ------------------------------------------------------------------ - - async def get_database_schema_id_and_table_id( - self, middleware_id: uuid.UUID, database: ConnectionInfo - ) -> dict[str, uuid.UUID]: - """Locate the schema and table IDs for the intercepts table within a database. - - Args: - middleware_id: The ID of the middleware owning the database. - database: Connection info used to identify the target database by name. - - Returns: - dict[str, uuid.UUID]: A dict with keys ``schema_id`` and ``table_id``. - - Raises: - ValueError: If the middleware, database, or table cannot be found. - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_error_handler("get_data"): - data = await get_api().get_data() - middlewares = data.get("middlewares", []) - - # Find the Middleware - mw = next((m for m in middlewares if m.get("id") == str(middleware_id)), None) - if not mw: - raise ValueError(f"Middleware {middleware_id} not found in response") - - # Find the Database - db = next((d for d in mw.get("databases", []) if d.get("name") == database.name), None) - if not db: - raise ValueError(f"Database '{database.name}' not found in middleware {middleware_id}") - - # Find the Schema and Table - # We assume the table exists within one of the schemas of this database - for schema in db.get("schemas", []): - table = next((t for t in schema.get("tables", []) if t.get("name") == INTERCEPTS_TABLE), None) - - if table: - return {"schema_id": uuid.UUID(str(schema["id"])), "table_id": uuid.UUID(str(table["id"]))} - - # Table not found - raise ValueError(f"Table '{INTERCEPTS_TABLE}' not found in any schema for database '{database.name}'") - - # ------------------------------------------------------------------ - # Columns - # ------------------------------------------------------------------ - - async def get_columns_from_database( - self, - middleware_id: uuid.UUID, - database_id: uuid.UUID, - schema_id: uuid.UUID, - table_id: uuid.UUID, - ) -> list[uuid.UUID]: - """Retrieve the column IDs for a specific table. - - Args: - middleware_id: The ID of the middleware owning the database. - database_id: The ID of the database containing the table. - schema_id: The ID of the schema containing the table. - table_id: The ID of the table whose columns to retrieve. - - Returns: - list[uuid.UUID]: List of column IDs. - - Raises: - ValueError: If any column is missing a valid ``id`` field. - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - async with provably_error_handler("get_columns_from_database"): - columns = await get_api().list_columns_from_database(middleware_id, database_id, schema_id, table_id) - - try: - return [uuid.UUID(str(col["id"])) for col in columns] - except (KeyError, ValueError) as e: - raise ValueError("One or more columns missing a valid 'id' field") from e - - # ------------------------------------------------------------------ - # Integrations - # ------------------------------------------------------------------ - - async def ensure_integration(self, collection_id: uuid.UUID) -> tuple[uuid.UUID, str]: - """Idempotent get-or-create of the intercepts integration for a collection. - - Reuses an existing enabled integration for the collection (returning its - current key) instead of minting a duplicate on every bootstrap. - - Args: - collection_id: The ID of the collection to associate with the integration. - - Returns: - tuple[uuid.UUID, str]: The ID and full API key of the integration. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - integration = { - "description": INTERCEPTS_TABLE, - "is_enabled": True, - "name": INTERCEPTS_TABLE, - "role": "developer", - "type": "agent", - "collections": [str(collection_id)], - } - - async with provably_error_handler("ensure_integration"): - result = await get_api().ensure_integration(integration) - api_key = result.get("api_key") - if not api_key: - raise ValueError("ensure_integration response missing 'api_key'") - return uuid.UUID(str(result["id"])), str(api_key) - - # ------------------------------------------------------------------ - # Preprocess - # ------------------------------------------------------------------ - - async def start_preprocess(self, middleware_id: uuid.UUID, table_id: uuid.UUID) -> uuid.UUID: - """Start a preprocessing job for a table. - - Args: - middleware_id: The ID of the middleware owning the table. - table_id: The ID of the table to preprocess. - - Returns: - uuid.UUID: The ID of the started preprocessing job. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - ProvablyDataError: If the response is malformed. - """ - async with provably_error_handler("start_preprocess"): - result = await get_api().start_preprocess(middleware_id, table_id) - return uuid.UUID(str(result["id"])) - - async def get_preprocess_completed(self, middleware_id: uuid.UUID, table_id: uuid.UUID, timeout: int = 60) -> None: - """ - Polls the preprocess status until it reaches 'completed'. - - Args: - middleware_id: The ID of the middleware. - table_id: The ID of the table being preprocessed. - timeout: maximum seconds to wait (default 60 seconds). - - Raises: - RuntimeError: If the status becomes 'error'. - TimeoutError: If the process exceeds the timeout. - """ - - start_time = asyncio.get_running_loop().time() - - api_client = get_api() - current_delay = 0.05 - - while (asyncio.get_running_loop().time() - start_time) < timeout: - async with provably_error_handler("get_preprocess_status"): - preprocess = await api_client.get_preprocess_status(middleware_id, table_id) - - status = preprocess.get("status") - - if status == "completed": - _log.info("preprocess_finished", table_id=str(table_id)) - return - - if status == "error": - error_detail = preprocess.get("error", preprocess.get("status_detail")) - raise RuntimeError(f"Table preprocessing failed: {error_detail}") - - # If status is 'pending' or 'processing', wait and try again - _log.debug("preprocess_in_progress", table_id=str(table_id), status=status) - await asyncio.sleep(current_delay) - current_delay = min(current_delay * 2, 0.1) - - raise TimeoutError(f"Preprocessing for table {table_id} timed out after {timeout}s") - - async def get_preprocess_status_only(self, middleware_id: uuid.UUID, table_id: uuid.UUID) -> str: - """Get the current preprocessing status without waiting. - - Args: - middleware_id: The ID of the middleware. - table_id: The ID of the table. - - Returns: - str: The current status ('pending', 'processing', 'completed', 'error', or 'unknown'). - """ - async with provably_error_handler("get_preprocess_status"): - preprocess = await get_api().get_preprocess_status(middleware_id, table_id) - status: str = preprocess.get("status", "unknown") - return status - - # ------------------------------------------------------------------ - # Queries / Proofs - # ------------------------------------------------------------------ - - async def run_query(self, middleware_id: uuid.UUID, collection_id: uuid.UUID, sql: str) -> uuid.UUID: - """Run a SQL query through a middleware and request a proof. - - Args: - middleware_id: The ID of the middleware to execute the query against. - collection_id: The ID of the collection to associate the query with. - sql: The SQL query string to execute. - - Returns: - dict: The raw JSON response from the API. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - _log.info("run_query_started", middleware_id=str(middleware_id), collection_id=str(collection_id)) - async with provably_error_handler("run_query"): - result = await get_api().run_query(middleware_id, collection_id, sql) - return uuid.UUID(str(result["query_id"])) - - async def get_query(self, query_id: uuid.UUID) -> dict[str, Any]: - """Retrieve a query record by ID.""" - async with provably_error_handler("get_query"): - return await get_api().get_query(query_id) - - async def get_query_proof(self, proof_id: uuid.UUID) -> bytes: - """Download the full proof data for a given proof ID.""" - async with provably_error_handler("get_query_proof"): - return await get_api().get_query_proof(proof_id) - - async def wait_for_proof_computation(self, query_id: uuid.UUID, timeout: int = 60) -> dict[str, Any]: - """ - Polls the query status until it reaches a terminal state (completed or failed). - - Args: - query_id: The unique identifier for the query. - timeout: Maximum seconds to wait for the proof/result. - - Returns: - dict[str, Any]: The final query result and proof data. - - Raises: - ProvablyAPIError: If the server rejects a status check. - RuntimeError: If the query fails on the backend. - TimeoutError: If the terminal state isn't reached within the timeout. - """ - start_time = asyncio.get_running_loop().time() - - api_client = get_api() - current_delay = 0.05 - - while (asyncio.get_running_loop().time() - start_time) < timeout: - async with provably_error_handler("wait_for_proof_computation"): - data = await api_client.get_query(query_id) - - proof = data.get("proof") - - # If proof is not null, check the internal status - if proof: - status = proof.get("status") - if status == "Completed": - _log.info("proof_generation_success", query_id=str(query_id)) - return data - - if status == "Failed": - _log.error("proof_generation_failed", query_id=str(query_id)) - raise RuntimeError(f"Provably proof generation failed for query {query_id}") - - # If status is 'Pending' (or proof is still null), we continue waiting - _log.debug("proof_generation_pending", query_id=str(query_id)) - await asyncio.sleep(current_delay) - current_delay = min(current_delay * 2, 0.1) - - raise TimeoutError(f"Timed out waiting for proof {query_id} after {timeout}s") - - async def verify_proof(self, query_id: uuid.UUID, integration_api_key: str) -> None: - """Run a SQL query through a middleware and request a proof. - - Args: - query_id: The ID of the query whose proof to verify. - integration_api_key: The integration API key used to authenticate this request. - - Returns: - dict: The raw JSON response from the API. - - Raises: - ProvablyAPIError: If the server rejects the request. - ProvablyConnectionError: If the network is unreachable. - """ - _log.info("verify_proof_started", query_id=str(query_id)) - async with provably_error_handler("verify_proof"): - await get_api().verify_proof(query_id, api_key=integration_api_key) - - async def wait_for_proof_verification( - self, query_id: uuid.UUID, integration_api_key: str, timeout: int = 60 - ) -> dict[str, Any]: - """ - Polls the query until the proof verification_status reaches 'Verified'. - - Args: - query_id: The identifier for the query. - integration_api_key: The integration API key used to authenticate polling requests. - timeout: Maximum seconds to wait for verification. - - Returns: - dict[str, Any]: The full response containing the Verified ProofInfo. - - Raises: - RuntimeError: If verification_status becomes 'Failed'. - TimeoutError: If verification doesn't complete within the timeout. - """ - start_time = asyncio.get_running_loop().time() - - api_client = get_api() - current_delay = 0.05 - - while (asyncio.get_running_loop().time() - start_time) < timeout: - async with provably_error_handler("wait_for_proof_verification"): - data = await api_client.get_query(query_id, api_key=integration_api_key) - - proof = data.get("proof") - - # If proof is not null, check the internal status - if proof: - v_status = proof.get("verification_status") - if v_status == "Verified": - _log.info("proof_verification_success", query_id=str(query_id)) - return data - - if v_status == "Failed": - _log.error("proof_verification_failed", query_id=str(query_id)) - raise RuntimeError(f"Provably proof verification failed for query {query_id}") - - # If status is 'Unverified' or 'Verifying', continue polling - _log.debug( - "proof_verification_pending", - query_id=str(query_id), - status=proof.get("verification_status") if proof else "null", - ) - await asyncio.sleep(current_delay) - current_delay = min(current_delay * 2, 0.1) - - raise TimeoutError(f"Verification for query {query_id} timed out after {timeout}s") - - # ------------------------------------------------------------------ - # URL helpers - # ------------------------------------------------------------------ - - def query_record_url(self, query_id: uuid.UUID) -> str: - """Build the Provably Data Admin URL for a query record. - - Args: - query_id: The ID of the query record. - - Returns: - str: The full URL to the query record in the Provably admin UI. - """ - return get_api().query_record_url(query_id) - - -# Shared singleton -service = ProvablyService() diff --git a/src/sourcerykit/ui/server.py b/src/sourcerykit/ui/server.py index 52c5fbe..4ada55c 100644 --- a/src/sourcerykit/ui/server.py +++ b/src/sourcerykit/ui/server.py @@ -11,6 +11,7 @@ from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import FileResponse from fastapi.staticfiles import StaticFiles +from provably.service import service from sourcerykit.db._engine import get_engine from sourcerykit.db._traces import ( @@ -18,7 +19,6 @@ select_trace_by_id_prefix, select_trace_intercepts_by_trace_id, ) -from sourcerykit.provably.service import service from sourcerykit.utils import extract_actual _log = logging.getLogger(__name__) diff --git a/tests/e2e/test_evaluate_handoff_e2e.py b/tests/e2e/test_evaluate_handoff_e2e.py index 55e348d..32da426 100644 --- a/tests/e2e/test_evaluate_handoff_e2e.py +++ b/tests/e2e/test_evaluate_handoff_e2e.py @@ -22,12 +22,13 @@ from unittest.mock import AsyncMock import pytest +from provably import ProvablyConfig +from provably._api import ProvablyAPI +from provably._http import ProvablyHTTPClient from sourcerykit.config import Settings from sourcerykit.errors import SourceryKitTrustError from sourcerykit.evaluator.evaluator import evaluate_handoff -from sourcerykit.provably._api import ProvablyAPI -from sourcerykit.provably._http import ProvablyHTTPClient from sourcerykit.schemas import HandoffClaim, HandoffPayload, Outcome, VerificationMode from sourcerykit.schemas.agent_response import ClaimedValue from tests.e2e.conftest import FakeHttpServer @@ -86,20 +87,25 @@ def _provably_settings(fake_server: FakeHttpServer) -> Settings: def _wired_service(_provably_settings: Settings, monkeypatch: pytest.MonkeyPatch) -> None: """Patch the Provably service layer to use a real HTTP client pointed at the fake server. - The evaluator imports ``service`` from ``sourcerykit.provably.service`` and calls + The evaluator imports ``service`` from ``provably.service`` and calls ``service.verify_proof`` and ``service.wait_for_proof_verification``. Those methods call ``get_api()`` which calls ``get_http()``. We patch: - - ``sourcerykit.provably._api.get_http`` to return our ProvablyHTTPClient - - ``sourcerykit.provably.service.get_api`` to return our ProvablyAPI + - ``provably._api.get_http`` to return our ProvablyHTTPClient + - ``provably.service.get_api`` to return our ProvablyAPI - ``sourcerykit.evaluator.evaluator.update_trace`` """ - http_client = ProvablyHTTPClient(settings=_provably_settings) - api = ProvablyAPI(settings=_provably_settings) + config = ProvablyConfig( + api_url=_provably_settings.provably_api, + app_url=_provably_settings.provably_app, + org_id=_provably_settings.org_id, + ) + http_client = ProvablyHTTPClient(config) + api = ProvablyAPI(config) - monkeypatch.setattr("sourcerykit.provably._api.get_http", lambda: http_client) - monkeypatch.setattr("sourcerykit.provably.service.get_api", lambda: api) + monkeypatch.setattr("provably._api.get_http", lambda: http_client) + monkeypatch.setattr("provably.service.get_api", lambda: api) monkeypatch.setattr( "sourcerykit.evaluator.evaluator.update_trace", AsyncMock(), diff --git a/tests/unit/test_answer_model.py b/tests/unit/test_answer_model.py deleted file mode 100644 index af75781..0000000 --- a/tests/unit/test_answer_model.py +++ /dev/null @@ -1,113 +0,0 @@ -"""Tests for sourcerykit.provably._answer_model.""" - -import pytest -from pydantic import ValidationError - -from sourcerykit.provably._answer_model import AggregateAnswer, QueryAnswer, ResultsetAnswer, TabularData - -# --------------------------------------------------------------------------- -# TabularData -# --------------------------------------------------------------------------- - - -class TestTabularData: - def test_empty_rows_returns_dict(self) -> None: - td = TabularData(columns=[{"name": "col1"}], rows=[]) - result = td.extract_value() - assert isinstance(result, dict) - assert result["rows"] == [] - - def test_single_cell_scalar_returned(self) -> None: - td = TabularData(columns=[{"name": "score"}], rows=[[42]]) - result = td.extract_value() - assert result == 42 - - def test_raw_response_column_parsed_as_json(self) -> None: - td = TabularData( - columns=[{"name": "raw_response"}], - rows=[['{"status": "ok", "value": 1}']], - ) - result = td.extract_value() - assert result == {"status": "ok", "value": 1} - - def test_raw_response_non_json_string_returned_as_is(self) -> None: - td = TabularData( - columns=[{"name": "raw_response"}], - rows=[["plain string"]], - ) - result = td.extract_value() - assert result == "plain string" - - def test_multi_column_returns_full_tabular_dict(self) -> None: - td = TabularData( - columns=[{"name": "a"}, {"name": "b"}], - rows=[[1, 2]], - ) - result = td.extract_value() - assert result == {"columns": [{"name": "a"}, {"name": "b"}], "rows": [[1, 2]]} - - -# --------------------------------------------------------------------------- -# AggregateAnswer -# --------------------------------------------------------------------------- - - -class TestAggregateAnswer: - def test_extract_value_returns_string(self) -> None: - aa = AggregateAnswer(type="aggregate", value="100") - assert aa.extract_value() == "100" - - def test_extract_value_parses_json_string(self) -> None: - aa = AggregateAnswer(type="aggregate", value='{"key": "val"}') - result = aa.extract_value() - assert result == {"key": "val"} - - def test_type_must_be_aggregate(self) -> None: - with pytest.raises(ValidationError): - AggregateAnswer(type="other", value="x") # pyright: ignore[reportArgumentType] - - -# --------------------------------------------------------------------------- -# ResultsetAnswer -# --------------------------------------------------------------------------- - - -class TestResultsetAnswer: - def test_extract_value_delegates_to_tabular(self) -> None: - rs = ResultsetAnswer( - type="resultset", - value=TabularData(columns=[{"name": "score"}], rows=[[99]]), - ) - assert rs.extract_value() == 99 - - -# --------------------------------------------------------------------------- -# QueryAnswer -# --------------------------------------------------------------------------- - - -class TestQueryAnswer: - def test_flatten_aggregate(self) -> None: - qa = QueryAnswer.model_validate({"type": "aggregate", "value": "42"}) - assert qa.flatten() == "42" - - def test_flatten_resultset_single_cell(self) -> None: - qa = QueryAnswer.model_validate( - { - "type": "resultset", - "value": { - "columns": [{"name": "count"}], - "rows": [[7]], - }, - } - ) - assert qa.flatten() == 7 - - def test_model_validate_wraps_root_automatically(self) -> None: - """wrap_root validator: top-level dict without 'root' key is wrapped.""" - qa = QueryAnswer.model_validate({"type": "aggregate", "value": "x"}) - assert qa.root is not None - - def test_invalid_type_raises_validation_error(self) -> None: - with pytest.raises(ValidationError): - QueryAnswer.model_validate({"type": "unknown", "value": "x"}) diff --git a/tests/unit/test_api.py b/tests/unit/test_api.py deleted file mode 100644 index 160a733..0000000 --- a/tests/unit/test_api.py +++ /dev/null @@ -1,69 +0,0 @@ -"""Tests for sourcerykit.provably._api.ProvablyAPI.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -from sourcerykit.provably._api import ProvablyAPI - - -def _make_api() -> tuple[ProvablyAPI, MagicMock]: - """Return a ProvablyAPI with mocked settings and HTTP client.""" - settings = MagicMock() - settings.org_id = uuid.uuid4() - settings.provably_app = "https://app.provably.ai" - api = ProvablyAPI(settings=settings) - return api, settings - - -class TestProvablyAPICreateFeedback: - async def test_calls_post_multipart(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.post_multipart = AsyncMock() - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.create_feedback({"description": "test"}, files={"file": ("f.txt", b"data")}) - - mock_http.post_multipart.assert_called_once_with( - "/api/v1/feedback", - {"description": "test"}, - files={"file": ("f.txt", b"data")}, - ) - - async def test_without_files(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.post_multipart = AsyncMock() - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.create_feedback({"description": "test"}) - - mock_http.post_multipart.assert_called_once_with( - "/api/v1/feedback", - {"description": "test"}, - files=None, - ) - - -class TestProvablyAPIListOrganizations: - async def test_returns_list(self) -> None: - api, _ = _make_api() - orgs = [{"id": str(uuid.uuid4()), "name": "Org1"}] - mock_http = MagicMock() - mock_http.get = AsyncMock(return_value=orgs) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - result = await api.list_organizations() - - assert result == orgs - mock_http.get.assert_called_once_with("/api/v1/organizations") - - async def test_returns_empty_list(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.get = AsyncMock(return_value=[]) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - result = await api.list_organizations() - - assert result == [] diff --git a/tests/unit/test_auth_api.py b/tests/unit/test_auth_api.py deleted file mode 100644 index 4ed78f3..0000000 --- a/tests/unit/test_auth_api.py +++ /dev/null @@ -1,110 +0,0 @@ -"""Tests for sourcerykit.provably._auth_api.ProvablyAuthAPI.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -from sourcerykit.provably._auth_api import ( - OAUTH_CLIENT_ID, - REDIRECT_URI, - Organization, - OrganizationType, - ProvablyAuthAPI, -) - -_TOKEN = "test-jwt-token" -_ORG = Organization(handle="my-org", name="My Org", organization_type=OrganizationType.EDUCATION) - - -def _make_api() -> tuple[ProvablyAuthAPI, MagicMock]: - """Return a ProvablyAuthAPI with its _http client fully mocked.""" - with patch("sourcerykit.provably._auth_api.ProvablyHTTPClient"): - api = ProvablyAuthAPI() - - mock_http = MagicMock() - api._http = mock_http - return api, mock_http - - -class TestProvablyAuthAPIOAuth: - async def test_exchange_code_posts_form(self) -> None: - api, mock_http = _make_api() - mock_http.post_form = AsyncMock(return_value={"access_token": "at", "refresh_token": "rt"}) - - result = await api.exchange_code("CODE", "VERIFIER") - - mock_http.post_form.assert_called_once_with( - "/api/v1/auth/oauth/token", - { - "grant_type": "authorization_code", - "code": "CODE", - "redirect_uri": REDIRECT_URI, - "client_id": OAUTH_CLIENT_ID, - "code_verifier": "VERIFIER", - }, - ) - assert result == {"access_token": "at", "refresh_token": "rt"} - - async def test_refresh_tokens_posts_form(self) -> None: - api, mock_http = _make_api() - mock_http.post_form = AsyncMock(return_value={"access_token": "new-at", "refresh_token": "new-rt"}) - - result = await api.refresh_tokens("old-rt") - - mock_http.post_form.assert_called_once_with( - "/api/v1/auth/oauth/refresh", - { - "grant_type": "refresh_token", - "refresh_token": "old-rt", - "client_id": OAUTH_CLIENT_ID, - }, - ) - assert result == {"access_token": "new-at", "refresh_token": "new-rt"} - - -class TestProvablyAuthAPIUser: - async def test_get_current_user_calls_get_with_token(self) -> None: - api, mock_http = _make_api() - mock_http.get = AsyncMock(return_value={"email": "user@example.com"}) - - result = await api.get_current_user(_TOKEN) - - mock_http.get.assert_called_once_with("/api/v1/user/current", token=_TOKEN) - assert result == {"email": "user@example.com"} - - -class TestProvablyAuthAPIOrganization: - async def test_create_organization_calls_post_multipart_with_token(self) -> None: - org_id = str(uuid.uuid4()) - api, mock_http = _make_api() - mock_http.post_multipart = AsyncMock(return_value={"id": org_id}) - - result = await api.create_organization(_TOKEN, _ORG) - - mock_http.post_multipart.assert_called_once_with( - "/api/v1/organizations", - { - "handle": _ORG.handle, - "name": _ORG.name, - "type": _ORG.organization_type.value, - }, - token=_TOKEN, - ) - assert result == {"id": org_id} - - async def test_get_organizations_calls_get_with_token(self) -> None: - orgs = [{"id": str(uuid.uuid4()), "name": "My Org"}] - api, mock_http = _make_api() - mock_http.get = AsyncMock(return_value=orgs) - - result = await api.get_organizations(_TOKEN) - - mock_http.get.assert_called_once_with("/api/v1/organizations", token=_TOKEN) - assert result == orgs - - async def test_get_organizations_returns_empty_list(self) -> None: - api, mock_http = _make_api() - mock_http.get = AsyncMock(return_value=[]) - - result = await api.get_organizations(_TOKEN) - - assert result == [] diff --git a/tests/unit/test_auth_service.py b/tests/unit/test_auth_service.py deleted file mode 100644 index b3b043c..0000000 --- a/tests/unit/test_auth_service.py +++ /dev/null @@ -1,170 +0,0 @@ -"""Tests for sourcerykit.provably.auth_service.ProvablyAuthService.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -from sourcerykit.provably._auth_api import OAuthTokens, Organization, OrganizationType -from sourcerykit.provably._errors import ( - ProvablyConnectionError, - ProvablyResourceAlreadyExistsError, - ProvablyUnauthorizedError, -) -from sourcerykit.provably.auth_service import ProvablyAuthService - -_TOKEN = "test-jwt-token" -_ORG = Organization(handle="my-org", name="My Org", organization_type=OrganizationType.EDUCATION) -_ORG_ID = uuid.uuid4() - - -def _make_service() -> tuple[ProvablyAuthService, MagicMock]: - """Return a ProvablyAuthService with `get_api` patched to a mock.""" - service = ProvablyAuthService() - mock_api = MagicMock() - return service, mock_api - - -class TestProvablyAuthServiceOAuth: - async def test_exchange_code_returns_tokens(self) -> None: - service, mock_api = _make_service() - mock_api.exchange_code = AsyncMock(return_value={"access_token": "at", "refresh_token": "rt"}) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.exchange_code("CODE", "VERIFIER") - - assert isinstance(result, OAuthTokens) - assert result.access_token == "at" - assert result.refresh_token == "rt" - - async def test_refresh_tokens_rotates(self) -> None: - service, mock_api = _make_service() - mock_api.refresh_tokens = AsyncMock(return_value={"access_token": "new-at", "refresh_token": "new-rt"}) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.refresh_tokens("old-rt") - - assert result.access_token == "new-at" - assert result.refresh_token == "new-rt" - - -class TestProvablyAuthServiceUserEmail: - async def test_returns_email(self) -> None: - service, mock_api = _make_service() - mock_api.get_current_user = AsyncMock(return_value={"email": "user@example.com"}) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.get_user_email(_TOKEN) - - assert result == "user@example.com" - - async def test_missing_email_raises_data_error(self) -> None: - from sourcerykit.provably._errors import ProvablyDataError - - service, mock_api = _make_service() - mock_api.get_current_user = AsyncMock(return_value={}) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - with pytest.raises(ProvablyDataError): - await service.get_user_email(_TOKEN) - - -class TestProvablyAuthServiceOrganization: - async def test_create_organization_returns_uuid(self) -> None: - service, mock_api = _make_service() - mock_api.create_organization = AsyncMock(return_value={"id": str(_ORG_ID)}) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.create_organization(_TOKEN, _ORG) - - assert result == _ORG_ID - - async def test_create_organization_already_exists_raises_error(self) -> None: - service, mock_api = _make_service() - mock_request = httpx.Request("POST", "https://api.provably.ai/api/v1/organizations") - mock_response = httpx.Response(409, request=mock_request, text="Conflict") - mock_api.create_organization = AsyncMock( - side_effect=httpx.HTTPStatusError("409", request=mock_request, response=mock_response) - ) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - with pytest.raises(ProvablyResourceAlreadyExistsError): - await service.create_organization(_TOKEN, _ORG) - - async def test_get_organizations_returns_list(self) -> None: - orgs = [{"id": str(_ORG_ID), "name": "My Org"}] - service, mock_api = _make_service() - mock_api.get_organizations = AsyncMock(return_value=orgs) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.get_organizations(_TOKEN) - - assert result == orgs - - async def test_get_organizations_returns_empty_list(self) -> None: - service, mock_api = _make_service() - mock_api.get_organizations = AsyncMock(return_value=[]) - - with patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api): - result = await service.get_organizations(_TOKEN) - - assert result == [] - - -class TestProvablyAuthServiceListOrganizations: - async def test_returns_list(self) -> None: - service, mock_api = _make_service() - orgs = [{"id": str(_ORG_ID), "name": "My Org"}] - mock_main_api = MagicMock() - mock_main_api.list_organizations = AsyncMock(return_value=orgs) - - with ( - patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api), - patch("sourcerykit.provably.auth_service.get_main_api", return_value=mock_main_api), - ): - result = await service.list_organizations() - - assert result == orgs - - async def test_returns_empty_list(self) -> None: - service, mock_api = _make_service() - mock_main_api = MagicMock() - mock_main_api.list_organizations = AsyncMock(return_value=[]) - - with ( - patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api), - patch("sourcerykit.provably.auth_service.get_main_api", return_value=mock_main_api), - ): - result = await service.list_organizations() - - assert result == [] - - async def test_connection_error(self) -> None: - service, mock_api = _make_service() - mock_main_api = MagicMock() - req = httpx.Request("GET", "https://api.provably.ai/api/v1/organizations") - mock_main_api.list_organizations = AsyncMock(side_effect=httpx.ConnectError("refused", request=req)) - - with ( - patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api), - patch("sourcerykit.provably.auth_service.get_main_api", return_value=mock_main_api), - ): - with pytest.raises(ProvablyConnectionError): - await service.list_organizations() - - async def test_unauthorized(self) -> None: - service, mock_api = _make_service() - mock_main_api = MagicMock() - mock_request = httpx.Request("GET", "https://api.provably.ai/api/v1/organizations") - mock_response = httpx.Response(401, request=mock_request, text="Unauthorized") - mock_main_api.list_organizations = AsyncMock( - side_effect=httpx.HTTPStatusError("401", request=mock_request, response=mock_response) - ) - - with ( - patch("sourcerykit.provably.auth_service.get_api", return_value=mock_api), - patch("sourcerykit.provably.auth_service.get_main_api", return_value=mock_main_api), - ): - with pytest.raises(ProvablyUnauthorizedError): - await service.list_organizations() diff --git a/tests/unit/test_cli.py b/tests/unit/test_cli.py index 19bdb39..468b031 100644 --- a/tests/unit/test_cli.py +++ b/tests/unit/test_cli.py @@ -1,9 +1,10 @@ """Tests for sourcerykit.cli — helper functions.""" -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import ANY, AsyncMock, MagicMock, patch import pytest import typer +from provably import OAuthTokens, ProvablyConnectionError from sourcerykit.cli.init import ( _run_oauth_browser, @@ -17,8 +18,6 @@ require_settings, run_connectivity_check, ) -from sourcerykit.provably._auth_api import OAuthTokens -from sourcerykit.provably._errors import ProvablyConnectionError _VALID_POSTGRES_URL = "postgresql://user:pass@1.2.3.4:5432/mydb" @@ -109,9 +108,9 @@ def test_returns_false_on_psycopg_connect_failure(self) -> None: class TestRunOauthBrowser: def test_logs_in_and_runs_post_auth_with_flags(self) -> None: - tokens = OAuthTokens(access_token="at", refresh_token="rt") + tokens = OAuthTokens(access_token="at", refresh_token="rt", client_id="sourcerykit-cli") with ( - patch("sourcerykit.cli.init.browser_login", new=AsyncMock(return_value=tokens)), + patch("sourcerykit.cli.init.browser_login", new=AsyncMock(return_value=tokens)) as mock_login, patch("sourcerykit.cli.init.service.get_user_email", new=AsyncMock(return_value="user@example.com")), patch("sourcerykit.cli.init.save_app_dir_config") as mock_save, patch("sourcerykit.cli.init._execute_post_auth_phases", return_value=True) as mock_phases, @@ -124,6 +123,8 @@ def test_logs_in_and_runs_post_auth_with_flags(self) -> None: sandbox=True, ) + # The SDK has no default client; sourcerykit must name its own. + mock_login.assert_awaited_once_with(ANY, client_id="sourcerykit-cli", port=8910) mock_save.assert_called_once_with(token="at", refresh_token="rt", email="user@example.com") mock_phases.assert_called_once_with( "at", diff --git a/tests/unit/test_db_helpers.py b/tests/unit/test_db_helpers.py index 15cfcac..aded047 100644 --- a/tests/unit/test_db_helpers.py +++ b/tests/unit/test_db_helpers.py @@ -1,35 +1,7 @@ -"""Tests for sourcerykit.db helpers — ConnectionInfo, _intercepts, _trusted_endpoints.""" +"""Tests for sourcerykit.db helpers — _intercepts, _trusted_endpoints.""" import uuid -from sourcerykit.db._engine import ConnectionInfo - -# --------------------------------------------------------------------------- -# ConnectionInfo -# --------------------------------------------------------------------------- - - -class TestConnectionInfo: - def test_to_dict_returns_all_fields(self) -> None: - ci = ConnectionInfo( - name="mydb", - username="user", - password="secret", - provider="postgresql", - uri="postgresql://user:secret@localhost:5432/mydb", - ) - d = ci.to_dict() - assert d["name"] == "mydb" - assert d["username"] == "user" - assert d["password"] == "secret" - assert d["provider"] == "postgresql" - assert d["uri"] == "postgresql://user:secret@localhost:5432/mydb" - - def test_to_dict_keys(self) -> None: - ci = ConnectionInfo(name="n", username="u", password="p", provider="prov", uri="u://x") - assert set(ci.to_dict().keys()) == {"name", "username", "password", "provider", "uri"} - - # --------------------------------------------------------------------------- # DB intercept SQL builders # --------------------------------------------------------------------------- diff --git a/tests/unit/test_doctor.py b/tests/unit/test_doctor.py index ca31616..f06cd8c 100644 --- a/tests/unit/test_doctor.py +++ b/tests/unit/test_doctor.py @@ -3,6 +3,8 @@ import uuid from unittest.mock import AsyncMock, MagicMock, patch +from provably import ProvablyConnectionError, ProvablyUnauthorizedError + from sourcerykit.cli.doctor import ( _check_bootstrap_ids, _check_database, @@ -15,7 +17,6 @@ run_doctor, ) from sourcerykit.config import Settings -from sourcerykit.provably._errors import ProvablyConnectionError, ProvablyUnauthorizedError _ORG_ID = uuid.uuid4() _TOKEN = "some-access-token" diff --git a/tests/unit/test_errors.py b/tests/unit/test_errors.py deleted file mode 100644 index cf90754..0000000 --- a/tests/unit/test_errors.py +++ /dev/null @@ -1,212 +0,0 @@ -"""Tests for sourcerykit.provably._errors — error hierarchy and provably_error_handler.""" - -import httpx -import pytest - -from sourcerykit.errors import SourceryKitError -from sourcerykit.provably._errors import ( - ProvablyAPIError, - ProvablyAuthError, - ProvablyConnectionError, - ProvablyDataError, - ProvablyError, - ProvablyNotFoundError, - ProvablyResourceAlreadyExistsError, - ProvablyUnauthorizedError, - provably_auth_error_handler, - provably_error_handler, -) - -# --------------------------------------------------------------------------- -# Error class hierarchy -# --------------------------------------------------------------------------- - - -class TestErrorHierarchy: - def test_provably_error_is_sourcerykit_error(self) -> None: - assert issubclass(ProvablyError, SourceryKitError) - - def test_api_error_is_provably_error(self) -> None: - assert issubclass(ProvablyAPIError, ProvablyError) - - def test_connection_error_is_provably_error(self) -> None: - assert issubclass(ProvablyConnectionError, ProvablyError) - - def test_data_error_is_provably_error(self) -> None: - assert issubclass(ProvablyDataError, ProvablyError) - - def test_api_error_stores_status_code(self) -> None: - err = ProvablyAPIError("bad request", status_code=400, response_body="oops") - assert err.status_code == 400 - assert err.response_body == "oops" - - def test_api_error_status_code_defaults_to_none(self) -> None: - err = ProvablyAPIError("bad request") - assert err.status_code is None - - def test_not_found_is_api_error(self) -> None: - assert issubclass(ProvablyNotFoundError, ProvablyAPIError) - - def test_not_found_stores_status_code(self) -> None: - err = ProvablyNotFoundError("not found", status_code=404, response_body="Not Found") - assert err.status_code == 404 - - -# --------------------------------------------------------------------------- -# provably_error_handler -# --------------------------------------------------------------------------- - - -class TestProvablyErrorHandler: - async def test_no_exception_passes_through(self) -> None: - async with provably_error_handler("test_op"): - pass # must not raise - - async def test_http_status_error_raises_api_error(self) -> None: - mock_request = httpx.Request("GET", "https://api.provably.ai/test") - mock_response = httpx.Response(422, request=mock_request, text="Unprocessable") - http_err = httpx.HTTPStatusError("422", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyAPIError) as exc_info: - async with provably_error_handler("run_query"): - raise http_err - - assert exc_info.value.status_code == 422 - assert "run query" in str(exc_info.value).lower() - - async def test_http_404_raises_not_found_error(self) -> None: - mock_request = httpx.Request("GET", "https://api.provably.ai/test") - mock_response = httpx.Response(404, request=mock_request, text="Not Found") - http_err = httpx.HTTPStatusError("404", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyNotFoundError) as exc_info: - async with provably_error_handler("get_preprocess_status"): - raise http_err - - assert exc_info.value.status_code == 404 - - async def test_value_error_raises_data_error(self) -> None: - with pytest.raises(ProvablyDataError): - async with provably_error_handler("get_collection"): - raise ValueError("unexpected key in response") - - async def test_key_error_raises_data_error(self) -> None: - with pytest.raises(ProvablyDataError): - async with provably_error_handler("get_database"): - raise KeyError("missing_field") - - async def test_request_error_raises_connection_error(self) -> None: - req = httpx.Request("GET", "https://api.provably.ai/") - with pytest.raises(ProvablyConnectionError): - async with provably_error_handler("create_middleware"): - raise httpx.ConnectError("connection refused", request=req) - - async def test_unexpected_exception_raises_provably_error(self) -> None: - with pytest.raises(ProvablyError): - async with provably_error_handler("run_query"): - raise RuntimeError("something unexpected") - - -# --------------------------------------------------------------------------- -# Auth error class hierarchy -# --------------------------------------------------------------------------- - - -class TestAuthErrorHierarchy: - def test_auth_error_is_api_error(self) -> None: - assert issubclass(ProvablyAuthError, ProvablyAPIError) - - def test_resource_already_exists_is_auth_error(self) -> None: - assert issubclass(ProvablyResourceAlreadyExistsError, ProvablyAuthError) - - def test_unauthorized_is_auth_error(self) -> None: - assert issubclass(ProvablyUnauthorizedError, ProvablyAuthError) - - def test_auth_error_stores_status_code_and_body(self) -> None: - err = ProvablyAuthError("bad auth", status_code=403, response_body="Forbidden") - assert err.status_code == 403 - assert err.response_body == "Forbidden" - - def test_unauthorized_error_stores_status_code(self) -> None: - err = ProvablyUnauthorizedError("wrong credentials", status_code=401, response_body="Unauthorized") - assert err.status_code == 401 - - def test_resource_already_exists_is_sourcerykit_error(self) -> None: - assert issubclass(ProvablyResourceAlreadyExistsError, SourceryKitError) - - -# --------------------------------------------------------------------------- -# provably_auth_error_handler -# --------------------------------------------------------------------------- - - -class TestProvablyAuthErrorHandler: - async def test_no_exception_passes_through(self) -> None: - async with provably_auth_error_handler("test_op"): - pass # must not raise - - async def test_http_401_raises_unauthorized_error(self) -> None: - mock_request = httpx.Request("GET", "https://api.provably.ai/api/v1/user/current") - mock_response = httpx.Response(401, request=mock_request, text="Unauthorized") - http_err = httpx.HTTPStatusError("401", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyUnauthorizedError) as exc_info: - async with provably_auth_error_handler("oauth_userinfo"): - raise http_err - - assert exc_info.value.status_code == 401 - assert "userinfo" in str(exc_info.value).lower() - - async def test_http_400_raises_auth_error(self) -> None: - mock_request = httpx.Request("POST", "https://api.provably.ai/api/v1/organizations") - mock_response = httpx.Response(400, request=mock_request, json={"description": "Invalid request"}) - http_err = httpx.HTTPStatusError("400", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyAuthError) as exc_info: - async with provably_auth_error_handler("create_organization"): - raise http_err - - assert exc_info.value.status_code == 400 - - async def test_http_500_raises_auth_error(self) -> None: - mock_request = httpx.Request("GET", "https://api.provably.ai/api/v1/user/current") - mock_response = httpx.Response(500, request=mock_request, text="Internal Server Error") - http_err = httpx.HTTPStatusError("500", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyAuthError) as exc_info: - async with provably_auth_error_handler("oauth_userinfo"): - raise http_err - - assert exc_info.value.status_code == 500 - - async def test_http_409_raises_resource_already_exists(self) -> None: - mock_request = httpx.Request("POST", "https://api.provably.ai/api/v1/organizations") - mock_response = httpx.Response(409, request=mock_request, text="Conflict") - http_err = httpx.HTTPStatusError("409", request=mock_request, response=mock_response) - - with pytest.raises(ProvablyResourceAlreadyExistsError) as exc_info: - async with provably_auth_error_handler("create_organization"): - raise http_err - - assert exc_info.value.status_code == 409 - - async def test_value_error_raises_data_error(self) -> None: - with pytest.raises(ProvablyDataError): - async with provably_auth_error_handler("get_api_key"): - raise ValueError("unexpected key in response") - - async def test_key_error_raises_data_error(self) -> None: - with pytest.raises(ProvablyDataError): - async with provably_auth_error_handler("get_api_key"): - raise KeyError("api_key") - - async def test_type_error_raises_data_error(self) -> None: - with pytest.raises(ProvablyDataError): - async with provably_auth_error_handler("create_organization"): - raise TypeError("cannot convert") - - async def test_request_error_raises_connection_error(self) -> None: - req = httpx.Request("GET", "https://api.provably.ai/api/v1/user/current") - with pytest.raises(ProvablyConnectionError): - async with provably_auth_error_handler("oauth_userinfo"): - raise httpx.ConnectError("connection refused", request=req) diff --git a/tests/unit/test_http.py b/tests/unit/test_http.py deleted file mode 100644 index 792f556..0000000 --- a/tests/unit/test_http.py +++ /dev/null @@ -1,264 +0,0 @@ -"""Tests for sourcerykit.provably._http.ProvablyHTTPClient.""" - -import os -from collections.abc import Generator -from pathlib import Path -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -from sourcerykit.config import Settings -from sourcerykit.provably._http import ProvablyHTTPClient, _clear_refresh_token, _persist_tokens - -_ORG = "00000000-0000-0000-0000-000000000001" - - -@pytest.fixture(autouse=True) -def _env(monkeypatch: pytest.MonkeyPatch) -> Generator[None, None, None]: - """Provide env so _fetch's default Bearer resolution (get_settings) succeeds.""" - from sourcerykit.config import get_settings, load_app_dir_config, load_local_env - - get_settings.cache_clear() - load_app_dir_config.cache_clear() - load_local_env.cache_clear() - monkeypatch.setenv("PROVABLY_ACCESS_TOKEN", "test-access-token") - monkeypatch.setenv("SOURCERYKIT_ORG_ID", _ORG) - yield - get_settings.cache_clear() - load_app_dir_config.cache_clear() - load_local_env.cache_clear() - - -def _make_settings(api_url: str = "https://api.provably.ai") -> Settings: - import uuid - - return Settings( - access_token="test-access-token", - org_id=uuid.UUID(_ORG), - postgres_url="postgresql://user:pass@localhost/db", - provably_api=api_url, - ) - - -def _make_client(api_url: str = "https://api.provably.ai") -> ProvablyHTTPClient: - return ProvablyHTTPClient(settings=_make_settings(api_url)) - - -class TestProvablyHTTPClientInit: - def test_base_url_stripped_of_trailing_slash(self) -> None: - client = _make_client("https://api.provably.ai/") - assert client.base_url == "https://api.provably.ai" - - def test_headers_do_not_include_auth(self) -> None: - client = _make_client() - assert "x-api-key" not in client._headers - assert "Authorization" not in client._headers - - def test_headers_include_content_type(self) -> None: - client = _make_client() - assert client._headers["Content-Type"] == "application/json" - - -class TestProvablyHTTPClientGet: - async def test_get_returns_parsed_json(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b'{"result": "ok"}' - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {"result": "ok"} - - with patch.object(client, "_request", AsyncMock(return_value=mock_response)): - result = await client.get("/test-path") - assert result == {"result": "ok"} - - async def test_get_raises_on_http_status_error(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.status_code = 404 - mock_response.text = "Not Found" - error = httpx.HTTPStatusError("404", request=MagicMock(), response=mock_response) - - with patch.object(client, "_request", AsyncMock(side_effect=error)): - with pytest.raises(httpx.HTTPStatusError): - await client.get("/missing") - - async def test_get_passes_params(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b"{}" - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {} - - with patch.object(client, "_request", AsyncMock(return_value=mock_response)) as mock_req: - await client.get("/path", params={"key": "val"}) - _, kwargs = mock_req.call_args - assert kwargs.get("params") == {"key": "val"} - - -class TestProvablyHTTPClientAuth: - async def test_post_auth_injects_access_token_as_bearer(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b'{"result": "ok"}' - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {"result": "ok"} - - with ( - patch("sourcerykit.provably._http.get_settings", return_value=_make_settings()), - patch.object(client, "_request", AsyncMock(return_value=mock_response)) as mock_req, - ): - await client.get("/path") - - _, kwargs = mock_req.call_args - assert kwargs.get("token") == "test-access-token" - - async def test_api_key_override_suppresses_bearer(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b'{"ok": true}' - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {"ok": True} - - with ( - patch("sourcerykit.provably._http.get_settings", return_value=_make_settings()), - patch.object(client, "_request", AsyncMock(return_value=mock_response)) as mock_req, - ): - await client.get("/path", api_key="i-zk-123") - - _, kwargs = mock_req.call_args - assert kwargs.get("api_key") == "i-zk-123" - assert kwargs.get("token") is None - - -class TestProvablyHTTPClientPost: - async def test_post_returns_empty_dict_on_empty_body(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b"" - mock_response.raise_for_status = MagicMock() - - with patch.object(client, "_request", AsyncMock(return_value=mock_response)): - result = await client.post("/some/path", json={"payload": "data"}) - assert result == {} - - async def test_post_with_custom_api_key(self) -> None: - client = _make_client() - mock_response = MagicMock() - mock_response.content = b'{"ok": true}' - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {"ok": True} - - with patch.object(client, "_request", AsyncMock(return_value=mock_response)) as mock_req: - await client.post("/path", json={}, api_key="custom-key") - _, kwargs = mock_req.call_args - assert kwargs.get("api_key") == "custom-key" - - -class TestProvablyHTTPClientPreAuth: - def test_pre_auth_sets_content_type_only(self) -> None: - with patch("sourcerykit.provably._http.get_bootstrap_settings", return_value="https://api.provably.ai"): - client = ProvablyHTTPClient(pre_auth=True) - assert "Content-Type" in client._headers - assert "x-api-key" not in client._headers - - def test_pre_auth_base_url_from_bootstrap_settings(self) -> None: - with patch("sourcerykit.provably._http.get_bootstrap_settings", return_value="https://custom.provably.ai"): - client = ProvablyHTTPClient(pre_auth=True) - assert client.base_url == "https://custom.provably.ai" - - async def test_pre_auth_token_sets_authorization_header(self) -> None: - with patch("sourcerykit.provably._http.get_bootstrap_settings", return_value="https://api.provably.ai"): - client = ProvablyHTTPClient(pre_auth=True) - - mock_response = MagicMock() - mock_response.content = b'{"token": "abc"}' - mock_response.raise_for_status = MagicMock() - mock_response.json.return_value = {"token": "abc"} - - with patch.object(client, "_request", AsyncMock(return_value=mock_response)) as mock_req: - await client.get("/api/v1/user/key", token="my-jwt-token") - _, kwargs = mock_req.call_args - assert kwargs.get("token") == "my-jwt-token" - - -class TestTokenStorePersistence: - def test_persist_tokens_writes_to_app_store_and_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - store = tmp_path / "app.env" - monkeypatch.setenv("SOURCERYKIT_TOKEN_STORE", str(store)) - - with ( - patch("sourcerykit.provably._http.get_settings") as mock_gs, - patch("sourcerykit.provably._http.load_local_env") as mock_lle, - patch("sourcerykit.provably._http.save_app_dir_config") as mock_save, - ): - _persist_tokens("at2", "rt2") - - mock_save.assert_not_called() - assert os.environ["PROVABLY_ACCESS_TOKEN"] == "at2" - assert os.environ["PROVABLY_REFRESH_TOKEN"] == "rt2" - content = store.read_text() - assert "PROVABLY_ACCESS_TOKEN" in content and "at2" in content - assert "PROVABLY_REFRESH_TOKEN" in content and "rt2" in content - mock_lle.cache_clear.assert_called_once() - mock_gs.cache_clear.assert_called_once() - - def test_persist_tokens_uses_global_json_when_no_store(self, monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.delenv("SOURCERYKIT_TOKEN_STORE", raising=False) - with patch("sourcerykit.provably._http.save_app_dir_config") as mock_save: - _persist_tokens("at", None) - mock_save.assert_called_once_with(token="at", refresh_token=None) - - def test_clear_refresh_token_from_app_store(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - store = tmp_path / "app.env" - monkeypatch.setenv("SOURCERYKIT_TOKEN_STORE", str(store)) - store.write_text("PROVABLY_REFRESH_TOKEN=rt\n") - os.environ["PROVABLY_REFRESH_TOKEN"] = "rt" - - _clear_refresh_token() - - assert "PROVABLY_REFRESH_TOKEN" not in os.environ - assert "PROVABLY_REFRESH_TOKEN=" in store.read_text() - - -class TestProvablyHTTPClientOAuthRefresh: - def _make_401(self) -> MagicMock: - resp = MagicMock() - resp.status_code = 401 - resp.text = "expired" - req = httpx.Request("GET", "http://x") - resp.raise_for_status.side_effect = httpx.HTTPStatusError("401", request=req, response=resp) - return resp - - def _make_ok(self) -> MagicMock: - resp = MagicMock() - resp.content = b'{"ok": true}' - resp.raise_for_status = MagicMock() - resp.json.return_value = {"ok": True} - return resp - - async def test_401_refreshes_once_and_retries(self) -> None: - client = _make_client() - calls = 0 - - async def fake_request(method: str, path: str, **kwargs: object) -> MagicMock: - nonlocal calls - calls += 1 - return self._make_ok() if calls > 1 else self._make_401() - - client._request = fake_request # type: ignore[method-assign] - - with patch("sourcerykit.provably._http._refresh_session", AsyncMock(return_value="new-token")) as refresh: - result = await client.get("/api/v1/data", token="old-token") - - assert result == {"ok": True} - refresh.assert_awaited_once() - assert calls == 2 - - async def test_401_without_refresh_token_raises(self) -> None: - client = _make_client() - client._request = AsyncMock(return_value=self._make_401()) # type: ignore[method-assign] - - with patch("sourcerykit.provably._http._refresh_session", AsyncMock(return_value=None)): - with pytest.raises(httpx.HTTPStatusError): - await client.get("/api/v1/data", token="old-token") diff --git a/tests/unit/test_oauth_login.py b/tests/unit/test_oauth_login.py deleted file mode 100644 index cc0bb1b..0000000 --- a/tests/unit/test_oauth_login.py +++ /dev/null @@ -1,126 +0,0 @@ -"""Tests for sourcerykit.provably.oauth_login.""" - -import asyncio -import base64 -import contextlib -import hashlib -import socket -import urllib.error -import urllib.request -from collections.abc import Callable - -import pytest - -from sourcerykit.provably import oauth_login -from sourcerykit.provably.oauth_login import _origin_of, _wait_for_loopback_code, pkce_pair - - -def _free_port() -> int: - """A port nothing else holds, so the suite does not fight the real CLI.""" - with contextlib.closing(socket.socket()) as probe: - probe.bind(("127.0.0.1", 0)) - return int(probe.getsockname()[1]) - - -@pytest.fixture -def loopback_port(monkeypatch: pytest.MonkeyPatch) -> int: - port = _free_port() - monkeypatch.setattr(oauth_login, "LOOPBACK_PORT", port) - return port - - -def test_pkce_pair_s256_verifiable() -> None: - verifier, challenge = pkce_pair() - assert verifier and challenge - expected = base64.urlsafe_b64encode(hashlib.sha256(verifier.encode()).digest()).rstrip(b"=").decode() - assert challenge == expected - - -def test_browser_login_carries_the_machine_id_in_the_consent_url( - monkeypatch: pytest.MonkeyPatch, -) -> None: - import asyncio - import urllib.parse - - from sourcerykit.provably import oauth_login - - opened: list[str] = [] - monkeypatch.setattr("sourcerykit.provably.oauth_login.webbrowser.open", lambda url: opened.append(url)) - - async def no_callback( - allowed_origin: str = "", - timeout: float = 300.0, - ) -> tuple[dict[str, str], Callable[[str | None], None]]: - return {}, lambda failure: None - - monkeypatch.setattr(oauth_login, "_wait_for_loopback_code", no_callback) - - async def run(**kwargs: str) -> None: - try: - await oauth_login.browser_login("https://app.example/consent", **kwargs) - except ValueError: - pass # empty callback; only the opened URL matters here - - asyncio.run(run(machine_id="opaque-123")) - with_id = dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(opened[0]).query)) - assert with_id["machine_id"] == "opaque-123" - assert with_id["response_type"] == "code" - - asyncio.run(run()) - without = dict(urllib.parse.parse_qsl(urllib.parse.urlsplit(opened[1]).query)) - assert "machine_id" not in without - - -def test_origin_of() -> None: - assert _origin_of("https://app.example/consent") == "https://app.example" - assert _origin_of("https://app.example/consent?a=1") == "https://app.example" - assert _origin_of("not a url") == "" - - -async def _open_callback(port: int, reply: dict[str, object]) -> None: - """Stand in for the browser: fetch /callback and record what came back.""" - request = urllib.request.Request(f"http://127.0.0.1:{port}/callback?code=abc&state=xyz") - - def fetch() -> None: - try: - with urllib.request.urlopen(request, timeout=10) as answer: - reply["status"] = answer.status - reply["origin"] = answer.headers.get("Access-Control-Allow-Origin") - reply["body"] = answer.read().decode() - except urllib.error.HTTPError as failed: - reply["status"] = failed.code - reply["origin"] = failed.headers.get("Access-Control-Allow-Origin") - reply["body"] = failed.read().decode() - - await asyncio.to_thread(fetch) - - -@pytest.mark.asyncio -async def test_callback_reply_waits_and_reports_success(loopback_port: int) -> None: - reply: dict[str, object] = {} - browser = asyncio.create_task(_open_callback(loopback_port, reply)) - - callback, settle = await _wait_for_loopback_code("https://app.example") - - assert callback == {"code": "abc", "state": "xyz"} - assert not reply, "the browser must still be waiting for its reply" - - settle(None) - await browser - - assert reply["status"] == 200 - assert reply["origin"] == "https://app.example" - assert "Logged in!" in str(reply["body"]) - - -@pytest.mark.asyncio -async def test_callback_reply_reports_failure(loopback_port: int) -> None: - reply: dict[str, object] = {} - browser = asyncio.create_task(_open_callback(loopback_port, reply)) - - _, settle = await _wait_for_loopback_code("https://app.example") - settle("the code was refused") - await browser - - assert reply["status"] == 500 - assert "the code was refused" in str(reply["body"]) diff --git a/tests/unit/test_provably_glue.py b/tests/unit/test_provably_glue.py new file mode 100644 index 0000000..45fd24f --- /dev/null +++ b/tests/unit/test_provably_glue.py @@ -0,0 +1,103 @@ +"""Tests for sourcerykit._provably — how sourcerykit plugs into the Provably SDK.""" + +import os +from pathlib import Path +from unittest.mock import patch + +import pytest +from provably import OAuthTokens +from provably._config import current + +from sourcerykit import _provably +from sourcerykit._provably import SourceryKitTokenStore, sdk_config +from sourcerykit.errors import SourceryKitConfigError +from sourcerykit.intercept._self_egress import is_self_egress + +_ORG = "00000000-0000-0000-0000-000000000001" + + +class TestSdkIsConfigured: + def test_importing_sourcerykit_points_the_sdk_at_its_settings_store_and_egress(self) -> None: + setup = current() + assert setup.config is sdk_config + assert isinstance(setup.tokens, SourceryKitTokenStore) + # Without this, the interceptor would record and gate the SDK's own calls. + with setup.egress(): + assert is_self_egress() + + def test_config_before_setup_has_urls_and_no_org(self) -> None: + with ( + patch.object(_provably, "get_settings", side_effect=SourceryKitConfigError("not set up")), + patch.object(_provably, "get_bootstrap_settings", return_value="https://api.example"), + patch.object(_provably, "get_bootstrap_app_url", return_value="https://app.example"), + ): + config = sdk_config() + assert (config.api_url, config.app_url, config.org_id) == ("https://api.example", "https://app.example", None) + + +class TestTokenStorePersistence: + def test_save_writes_to_app_store_and_env(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + store = tmp_path / "app.env" + monkeypatch.setenv("SOURCERYKIT_TOKEN_STORE", str(store)) + + with ( + patch.object(_provably, "get_settings") as mock_gs, + patch.object(_provably, "load_local_env") as mock_lle, + patch.object(_provably, "save_app_dir_config") as mock_save, + ): + SourceryKitTokenStore().save( + OAuthTokens(access_token="at2", refresh_token="rt2", client_id="sourcerykit-cli") + ) + + mock_save.assert_not_called() + assert os.environ["PROVABLY_ACCESS_TOKEN"] == "at2" + assert os.environ["PROVABLY_REFRESH_TOKEN"] == "rt2" + content = store.read_text() + assert "PROVABLY_ACCESS_TOKEN" in content and "at2" in content + assert "PROVABLY_REFRESH_TOKEN" in content and "rt2" in content + mock_lle.cache_clear.assert_called_once() + mock_gs.cache_clear.assert_called_once() + + def test_save_uses_global_json_when_no_store(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SOURCERYKIT_TOKEN_STORE", raising=False) + with patch.object(_provably, "save_app_dir_config") as mock_save: + SourceryKitTokenStore().save( + OAuthTokens(access_token="at", refresh_token=None, client_id="sourcerykit-cli") + ) + mock_save.assert_called_once_with(token="at", refresh_token=None) + + def test_clear_refresh_token_from_app_store(self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + store = tmp_path / "app.env" + monkeypatch.setenv("SOURCERYKIT_TOKEN_STORE", str(store)) + store.write_text("PROVABLY_REFRESH_TOKEN=rt\n") + monkeypatch.setenv("PROVABLY_REFRESH_TOKEN", "rt") + + SourceryKitTokenStore().clear_refresh_token() + + assert "PROVABLY_REFRESH_TOKEN" not in os.environ + assert "PROVABLY_REFRESH_TOKEN=" in store.read_text() + + def test_load_reads_the_session_from_settings(self, monkeypatch: pytest.MonkeyPatch) -> None: + from sourcerykit.config import get_settings + + monkeypatch.setenv("PROVABLY_ACCESS_TOKEN", "at") + monkeypatch.setenv("PROVABLY_REFRESH_TOKEN", "rt") + monkeypatch.setenv("SOURCERYKIT_ORG_ID", _ORG) + get_settings.cache_clear() + try: + # A refresh must name the client; the SDK has no default, so the store supplies it. + assert SourceryKitTokenStore().load() == OAuthTokens( + access_token="at", refresh_token="rt", client_id="sourcerykit-cli" + ) + finally: + get_settings.cache_clear() + + +def test_consent_page_has_its_own_setting(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.delenv("SOURCERYKIT_CONSENT_URL", raising=False) + # The app URL is the base of query-record links, so it must not move the consent page. + monkeypatch.setenv("SOURCERYKIT_PROVABLY_APP_URL", "https://app.example") + assert _provably.consent_page_url() == "https://switchboard.provably.ai/consent" + + monkeypatch.setenv("SOURCERYKIT_CONSENT_URL", "http://localhost:3000/consent") + assert _provably.consent_page_url() == "http://localhost:3000/consent" diff --git a/tests/unit/test_sandbox_api.py b/tests/unit/test_sandbox_api.py deleted file mode 100644 index 889cb32..0000000 --- a/tests/unit/test_sandbox_api.py +++ /dev/null @@ -1,107 +0,0 @@ -"""Tests for sourcerykit.provably._api — sandbox API methods.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -from sourcerykit.provably._api import ProvablyAPI - - -def _make_api() -> tuple[ProvablyAPI, MagicMock]: - """Return a ProvablyAPI with mocked settings and HTTP client.""" - settings = MagicMock() - settings.org_id = uuid.uuid4() - settings.provably_app = "https://app.provably.ai" - api = ProvablyAPI(settings=settings) - return api, settings - - -# --------------------------------------------------------------------------- -# create_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyAPICreateSandbox: - async def test_posts_org_id(self) -> None: - api, settings = _make_api() - mock_http = MagicMock() - mock_http.post = AsyncMock(return_value={"status": "active", "connection_uri": "postgresql://sandbox/db"}) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - result = await api.create_sandbox(settings.org_id) - - mock_http.post.assert_called_once_with( - "/api/v1/sandboxes", - {"org_id": str(settings.org_id)}, - token=None, - ) - assert result["status"] == "active" - - async def test_passes_token(self) -> None: - api, settings = _make_api() - mock_http = MagicMock() - mock_http.post = AsyncMock(return_value={"status": "provisioning"}) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.create_sandbox(settings.org_id, token="jwt-abc") - - mock_http.post.assert_called_once_with( - "/api/v1/sandboxes", - {"org_id": str(settings.org_id)}, - token="jwt-abc", - ) - - -# --------------------------------------------------------------------------- -# get_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyAPIGetSandbox: - async def test_returns_sandbox_record(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - sandbox = {"status": "active", "connection_uri": "postgresql://sandbox/db"} - mock_http.get = AsyncMock(return_value=sandbox) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - result = await api.get_sandbox() - - mock_http.get.assert_called_once_with("/api/v1/sandboxes", token=None) - assert result == sandbox - - async def test_passes_token(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.get = AsyncMock(return_value={}) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.get_sandbox(token="jwt-xyz") - - mock_http.get.assert_called_once_with("/api/v1/sandboxes", token="jwt-xyz") - - -# --------------------------------------------------------------------------- -# delete_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyAPIDeleteSandbox: - async def test_calls_delete(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.delete = AsyncMock(return_value=None) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.delete_sandbox() - - mock_http.delete.assert_called_once_with("/api/v1/sandboxes", token=None) - - async def test_passes_token(self) -> None: - api, _ = _make_api() - mock_http = MagicMock() - mock_http.delete = AsyncMock(return_value=None) - - with patch("sourcerykit.provably._api.get_http", return_value=mock_http): - await api.delete_sandbox(token="jwt-del") - - mock_http.delete.assert_called_once_with("/api/v1/sandboxes", token="jwt-del") diff --git a/tests/unit/test_sandbox_service.py b/tests/unit/test_sandbox_service.py deleted file mode 100644 index 469ff9b..0000000 --- a/tests/unit/test_sandbox_service.py +++ /dev/null @@ -1,187 +0,0 @@ -"""Tests for sourcerykit.provably.service — sandbox service methods.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -import pytest - -from sourcerykit.provably._errors import ProvablyDataError -from sourcerykit.provably.service import ProvablyService - - -def _make_service() -> tuple[ProvablyService, MagicMock]: - """Return a ProvablyService with a mocked API.""" - service = ProvablyService() - mock_api = MagicMock() - return service, mock_api - - -# --------------------------------------------------------------------------- -# create_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyServiceCreateSandbox: - async def test_returns_uri(self) -> None: - service, mock_api = _make_service() - org_id = uuid.uuid4() - mock_api.create_sandbox = AsyncMock( - return_value={"status": "active", "connection_uri": "postgresql://sandbox/db"} - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.create_sandbox(org_id) - - assert result == "postgresql://sandbox/db" - - async def test_missing_uri_raises(self) -> None: - service, mock_api = _make_service() - org_id = uuid.uuid4() - mock_api.create_sandbox = AsyncMock(return_value={"status": "provisioning"}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - with pytest.raises(ProvablyDataError, match="connection_uri"): - await service.create_sandbox(org_id) - - async def test_passes_token(self) -> None: - service, mock_api = _make_service() - org_id = uuid.uuid4() - mock_api.create_sandbox = AsyncMock(return_value={"connection_uri": "postgresql://sandbox/db"}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - await service.create_sandbox(org_id, token="jwt-abc") - - mock_api.create_sandbox.assert_called_once_with(org_id, token="jwt-abc") - - -# --------------------------------------------------------------------------- -# get_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyServiceGetSandbox: - async def test_returns_sandbox_data(self) -> None: - service, mock_api = _make_service() - sandbox = {"status": "active", "connection_uri": "postgresql://sandbox/db"} - mock_api.get_sandbox = AsyncMock(return_value=sandbox) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.get_sandbox() - - assert result == sandbox - - async def test_not_found_returns_none(self) -> None: - service = ProvablyService() - with patch.object(service, "get_sandbox", return_value=None): - result = await service.get_sandbox() - - assert result is None - - -# --------------------------------------------------------------------------- -# get_sandbox_connection_uri -# --------------------------------------------------------------------------- - - -class TestProvablyServiceGetSandboxConnectionUri: - async def test_active_returns_uri(self) -> None: - service, mock_api = _make_service() - mock_api.get_sandbox = AsyncMock(return_value={"status": "active", "connection_uri": "postgresql://sandbox/db"}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.get_sandbox_connection_uri() - - assert result == "postgresql://sandbox/db" - - async def test_provisioning_returns_uri(self) -> None: - service, mock_api = _make_service() - mock_api.get_sandbox = AsyncMock( - return_value={"status": "provisioning", "connection_uri": "postgresql://sandbox/db"} - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.get_sandbox_connection_uri() - - assert result == "postgresql://sandbox/db" - - async def test_expired_returns_none(self) -> None: - service, mock_api = _make_service() - mock_api.get_sandbox = AsyncMock( - return_value={"status": "expired", "connection_uri": "postgresql://sandbox/db"} - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.get_sandbox_connection_uri() - - assert result is None - - async def test_no_sandbox_returns_none(self) -> None: - service = ProvablyService() - with patch.object(service, "get_sandbox", return_value=None): - result = await service.get_sandbox_connection_uri() - - assert result is None - - -# --------------------------------------------------------------------------- -# get_sandbox_status -# --------------------------------------------------------------------------- - - -class TestProvablyServiceGetSandboxStatus: - async def test_matching_url_returns_true(self) -> None: - service, mock_api = _make_service() - uri = "postgresql://user:pass@sandbox.provably.ai:5432/mydb" - mock_api.get_sandbox = AsyncMock(return_value={"status": "active", "connection_uri": uri}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - sandbox, is_sandbox = await service.get_sandbox_status(uri) - - assert is_sandbox is True - assert sandbox is not None - assert sandbox["status"] == "active" - - async def test_different_url_returns_false(self) -> None: - service, mock_api = _make_service() - sandbox_uri = "postgresql://user:pass@sandbox.provably.ai:5432/mydb" - personal_uri = "postgresql://user:pass@myhost:5432/mydb" - mock_api.get_sandbox = AsyncMock(return_value={"status": "active", "connection_uri": sandbox_uri}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - sandbox, is_sandbox = await service.get_sandbox_status(personal_uri) - - assert is_sandbox is False - assert sandbox is not None - - async def test_no_sandbox_returns_none_false(self) -> None: - service = ProvablyService() - with patch.object(service, "get_sandbox", return_value=None): - sandbox, is_sandbox = await service.get_sandbox_status("postgresql://host/db") - - assert sandbox is None - assert is_sandbox is False - - -# --------------------------------------------------------------------------- -# delete_sandbox -# --------------------------------------------------------------------------- - - -class TestProvablyServiceDeleteSandbox: - async def test_calls_api_delete(self) -> None: - service, mock_api = _make_service() - mock_api.delete_sandbox = AsyncMock(return_value=None) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - await service.delete_sandbox() - - mock_api.delete_sandbox.assert_awaited_once() - - async def test_passes_token(self) -> None: - service, mock_api = _make_service() - mock_api.delete_sandbox = AsyncMock(return_value=None) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - await service.delete_sandbox(token="jwt-del") - - mock_api.delete_sandbox.assert_called_once_with(token="jwt-del") diff --git a/tests/unit/test_service.py b/tests/unit/test_service.py deleted file mode 100644 index 45a1030..0000000 --- a/tests/unit/test_service.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Tests for sourcerykit.provably.service.ProvablyService.""" - -import uuid -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import pytest - -from sourcerykit.provably._errors import ProvablyAPIError, ProvablyDataError -from sourcerykit.provably.service import ProvablyService - - -def _make_service() -> tuple[ProvablyService, MagicMock]: - """Return a ProvablyService with a mocked API.""" - service = ProvablyService() - mock_api = MagicMock() - return service, mock_api - - -# --------------------------------------------------------------------------- -# create_feedback -# --------------------------------------------------------------------------- - - -class TestProvablyServiceCreateFeedback: - async def test_success_without_file(self) -> None: - service, mock_api = _make_service() - mock_api.create_feedback = AsyncMock(return_value=None) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - await service.create_feedback("Great product!", None) - - mock_api.create_feedback.assert_called_once_with( - {"description": "Great product!"}, - files={}, - ) - - async def test_success_with_file(self) -> None: - service, mock_api = _make_service() - mock_api.create_feedback = AsyncMock(return_value=None) - file_bytes = b"some file content" - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - await service.create_feedback("Bug report", file_bytes) - - mock_api.create_feedback.assert_called_once_with( - {"description": "Bug report"}, - files={"files": ("attachment.dat", file_bytes)}, - ) - - async def test_api_error(self) -> None: - service, mock_api = _make_service() - mock_request = httpx.Request("POST", "https://api.provably.ai/api/v1/feedback") - mock_response = httpx.Response(500, request=mock_request, text="Internal Server Error") - mock_api.create_feedback = AsyncMock( - side_effect=httpx.HTTPStatusError("500", request=mock_request, response=mock_response) - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - with pytest.raises(ProvablyAPIError): - await service.create_feedback("test", None) - - -# --------------------------------------------------------------------------- -# list_collections -# --------------------------------------------------------------------------- - - -class TestProvablyServiceListCollections: - async def test_returns_list(self) -> None: - service, mock_api = _make_service() - collections = [{"id": str(uuid.uuid4()), "name": "my-project"}] - mock_api.list_collections = AsyncMock(return_value=collections) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.list_collections() - - assert result == collections - - async def test_returns_empty_list(self) -> None: - service, mock_api = _make_service() - mock_api.list_collections = AsyncMock(return_value=[]) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.list_collections() - - assert result == [] - - async def test_api_error(self) -> None: - service, mock_api = _make_service() - mock_request = httpx.Request("GET", "https://api.provably.ai/api/v1/collections") - mock_response = httpx.Response(403, request=mock_request, text="Forbidden") - mock_api.list_collections = AsyncMock( - side_effect=httpx.HTTPStatusError("403", request=mock_request, response=mock_response) - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - with pytest.raises(ProvablyAPIError): - await service.list_collections() - - -# --------------------------------------------------------------------------- -# ensure_integration -# --------------------------------------------------------------------------- - - -class TestProvablyServiceEnsureIntegration: - async def test_returns_id_and_key(self) -> None: - service, mock_api = _make_service() - collection_id = uuid.uuid4() - integration_id = uuid.uuid4() - mock_api.ensure_integration = AsyncMock( - return_value={"id": str(integration_id), "api_key": "i-zk-key-123", "collections": [str(collection_id)]} - ) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - result = await service.ensure_integration(collection_id) - - assert result == (integration_id, "i-zk-key-123") - call_body = mock_api.ensure_integration.call_args.args[0] - assert call_body["collections"] == [str(collection_id)] - - async def test_missing_key_raises_data_error(self) -> None: - service, mock_api = _make_service() - mock_api.ensure_integration = AsyncMock(return_value={"id": str(uuid.uuid4())}) - - with patch("sourcerykit.provably.service.get_api", return_value=mock_api): - with pytest.raises(ProvablyDataError): - await service.ensure_integration(uuid.uuid4()) diff --git a/uv.lock b/uv.lock index c238d6f..faeb106 100644 --- a/uv.lock +++ b/uv.lock @@ -1858,6 +1858,21 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f5/cd/785c64ed382f3f04201870267b02783f63b4678c2acfddc177a3ebcc2727/propcache-0.5.4-py3-none-any.whl", hash = "sha256:62c60aec739ed00124573cce1178138fd690c7676352d67a37328c1cf51d7468", size = 16338, upload-time = "2026-09-16T00:17:13.106Z" }, ] +[[package]] +name = "provably-sdk" +version = "0.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "httpx" }, + { name = "msgspec" }, + { name = "pydantic" }, + { name = "structlog" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/17/48/8586bacc9a2450a8fc2b20f6aec6e4f1b9c311aa38cd8acae5163cfd78dd/provably_sdk-0.3.1.tar.gz", hash = "sha256:dcbc261ff21b60e77aa43668e744c22bafd9d6f23bce749b7ab6db37a029aa6a", size = 68115, upload-time = "2026-09-23T13:19:15.871Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/80/0a/ed8f26f908e97e9b4fc90d953211f5e8b4ef48c52430963e773143f0f569/provably_sdk-0.3.1-py3-none-any.whl", hash = "sha256:870e62f1ed4c09e98d78f5d2aacf4abf66d9043ac5c71f24b74921aa2a97e635", size = 26231, upload-time = "2026-09-23T13:19:14.51Z" }, +] + [[package]] name = "psycopg" version = "3.3.5" @@ -2418,6 +2433,7 @@ dependencies = [ { name = "httpx" }, { name = "jsonschema" }, { name = "msgspec" }, + { name = "provably-sdk" }, { name = "psycopg", extra = ["binary"] }, { name = "pydantic" }, { name = "python-dotenv" }, @@ -2460,6 +2476,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.10" }, { name = "openai-agents", marker = "extra == 'dev'", specifier = ">=0.0.3" }, { name = "pre-commit", marker = "extra == 'dev'", specifier = ">=4.0" }, + { name = "provably-sdk", specifier = ">=0.3.1,<1" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.1" }, { name = "pydantic", specifier = ">=2.6" }, { name = "pytest", marker = "extra == 'dev'", specifier = ">=8.0" },