pyspiral 0.6.6__cp312-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

spiral/api/__init__.py

@@ -0,0 +1,78 @@
+import os
+from typing import TYPE_CHECKING
+
+import httpx
+
+from .client import _Client
+
+if TYPE_CHECKING:
+    from spiral.core.client import Authn
+
+    from .admin import AdminService
+    from .filesystems import FileSystemService
+    from .key_space_indexes import KeySpaceIndexesService
+    from .organizations import OrganizationService
+    from .projects import ProjectService
+    from .telemetry import TelemetryService
+    from .text_indexes import TextIndexesService
+    from .workloads import WorkloadService
+
+
+class SpiralAPI:
+    def __init__(self, authn: "Authn", base_url: str | None = None):
+        self.base_url = base_url or os.environ.get("SPIRAL_URL", "https://api.spiraldb.com")
+        self.client = _Client(
+            httpx.Client(
+                base_url=self.base_url,
+                timeout=None if ("PYTEST_VERSION" in os.environ or bool(os.environ.get("SPIRAL_DEV", None))) else 60,
+            ),
+            authn,
+        )
+
+    @property
+    def _admin(self) -> "AdminService":
+        from .admin import AdminService
+
+        return AdminService(self.client)
+
+    @property
+    def organization(self) -> "OrganizationService":
+        from .organizations import OrganizationService
+
+        return OrganizationService(self.client)
+
+    @property
+    def project(self) -> "ProjectService":
+        from .projects import ProjectService
+
+        return ProjectService(self.client)
+
+    @property
+    def file_system(self) -> "FileSystemService":
+        from .filesystems import FileSystemService
+
+        return FileSystemService(self.client)
+
+    @property
+    def workload(self) -> "WorkloadService":
+        from .workloads import WorkloadService
+
+        return WorkloadService(self.client)
+
+    @property
+    def text_indexes(self) -> "TextIndexesService":
+        from .text_indexes import TextIndexesService
+
+        return TextIndexesService(self.client)
+
+    @property
+    def key_space_indexes(self) -> "KeySpaceIndexesService":
+        from .key_space_indexes import KeySpaceIndexesService
+
+        return KeySpaceIndexesService(self.client)
+
+    @property
+    def telemetry(self) -> "TelemetryService":
+        from .telemetry import TelemetryService
+
+        return TelemetryService(self.client)

spiral/api/admin.py

@@ -0,0 +1,15 @@
+from .client import Paged, PagedResponse, ServiceBase
+from .organizations import OrgMembership
+from .types import OrgId
+
+
+class AdminService(ServiceBase):
+    def sync_memberships(self, org_id: OrgId | None = None) -> Paged[OrgMembership]:
+        params = {}
+        if org_id:
+            params["org_id"] = str(org_id)
+        return self.client.paged("/v1/admin/sync-memberships", PagedResponse[OrgMembership], params=params)
+
+    def sync_orgs(self) -> Paged[OrgId]:
+        params = {}
+        return self.client.paged("/v1/admin/sync-orgs", PagedResponse[OrgId], params=params)

spiral/api/client.py

@@ -0,0 +1,164 @@
+import logging
+from collections.abc import Iterable, Iterator, Mapping
+from typing import Any, TypeVar
+
+import httpx
+from httpx import HTTPStatusError
+from pydantic import BaseModel, Field, TypeAdapter
+
+from spiral.core.authn import Authn
+
+log = logging.getLogger(__name__)
+
+
+E = TypeVar("E")
+
+
+class PagedRequest(BaseModel):
+    page_token: str | None = None
+    page_size: int = 50
+
+
+class PagedResponse[E](BaseModel):
+    items: list[E] = Field(default_factory=list)
+    next_page_token: str | None = None
+
+
+PagedReqT = TypeVar("PagedReqT", bound=PagedRequest)
+
+
+class Paged[E](Iterable[E]):
+    def __init__(
+        self,
+        client: "_Client",
+        path: str,
+        page_token: str | None,
+        page_size: int,
+        response_cls: type[PagedResponse[E]],
+        params: Mapping[str, str] | None = None,
+    ):
+        self._client = client
+        self._path = path
+        self._response_cls = response_cls
+
+        self._page_size = page_size
+
+        self._params = params or {}
+        if page_token is not None:
+            self._params["page_token"] = str(page_token)
+        # TODO(marko): Support paging.
+        # if page_size is not None:
+        #     self._params["page_size"] = str(page_size)
+
+        self._response: PagedResponse[E] = client.get(path, response_cls, params=self._params)
+
+    @property
+    def page(self) -> PagedResponse[E]:
+        return self._response
+
+    def __iter__(self) -> Iterator[E]:
+        while True:
+            yield from self._response.items
+
+            if self._response.next_page_token is None:
+                break
+
+            params = self._params.copy()
+            params["page_token"] = self._response.next_page_token
+            self._response = self._client.get(self._path, self._response_cls, params=params)
+
+
+class ServiceBase:
+    def __init__(self, client: "_Client"):
+        self.client = client
+
+
+class SpiralHTTPError(Exception):
+    def __init__(self, body: str, code: int):
+        super().__init__(body)
+        self.body = body
+        self.code = code
+
+
+class _Client:
+    RequestT = TypeVar("RequestT")
+    ResponseT = TypeVar("ResponseT")
+
+    def __init__(self, http: httpx.Client, authn: Authn):
+        self.http = http
+        self.authn = authn
+
+    def get(
+        self, path: str, response_cls: type[ResponseT], *, params: Mapping[str, str | list[str]] | None = None
+    ) -> ResponseT:
+        return self.request("GET", path, None, response_cls, params=params)
+
+    def post(
+        self,
+        path: str,
+        req: RequestT,
+        response_cls: type[ResponseT],
+        *,
+        params: Mapping[str, str | list[str]] | None = None,
+    ) -> ResponseT:
+        return self.request("POST", path, req, response_cls, params=params)
+
+    def put(
+        self,
+        path: str,
+        req: RequestT,
+        response_cls: type[ResponseT],
+        *,
+        params: Mapping[str, str | list[str]] | None = None,
+    ) -> ResponseT:
+        return self.request("PUT", path, req, response_cls, params=params)
+
+    def delete(
+        self, path: str, response_cls: type[ResponseT], *, params: Mapping[str, str | list[str]] | None = None
+    ) -> ResponseT:
+        return self.request("DELETE", path, None, response_cls, params=params)
+
+    def request(
+        self,
+        method: str,
+        path: str,
+        req: RequestT | None,
+        response_cls: type[ResponseT],
+        *,
+        params: Mapping[str, str | list[str]] | None = None,
+    ) -> ResponseT:
+        req_data: dict[str, Any] = {}
+        if req is not None:
+            req_data = dict(json=TypeAdapter(req.__class__).dump_python(req, mode="json", exclude_none=True))
+
+        token = self.authn.token()
+        resp = self.http.request(
+            method,
+            path,
+            params=params or {},
+            headers={"authorization": f"Bearer {token.expose_secret()}"} if token else None,
+            **req_data,
+        )
+
+        try:
+            resp.raise_for_status()
+        except HTTPStatusError as e:
+            # Enrich the exception with the response body
+            raise SpiralHTTPError(body=resp.text, code=resp.status_code) from e
+
+        if response_cls == type[None]:
+            assert resp.text == ""
+            return None
+
+        return TypeAdapter(response_cls).validate_python(resp.json())
+
+    def paged(
+        self,
+        path: str,
+        response_cls: type[PagedResponse[E]],
+        *,
+        page_token: str | None = None,
+        page_size: int = 50,
+        params: Mapping[str, str] | None = None,
+    ) -> Paged[E]:
+        return Paged(self, path, page_token, page_size, response_cls, params)

spiral/api/filesystems.py

@@ -0,0 +1,134 @@
+from enum import Enum
+from types import NoneType
+from typing import Annotated, Literal
+
+from pydantic import AfterValidator, BaseModel, Field
+
+from .client import Paged, PagedResponse, ServiceBase
+from .types import ProjectId
+
+
+def _validate_directory_path(path: str) -> str:
+    if not path.startswith("/"):
+        raise ValueError("Directory path must start with a slash.")
+    if not path.endswith("/"):
+        raise ValueError("Directory path must not end with a slash.")
+    return path
+
+
+DirectoryPath = Annotated[str, AfterValidator(_validate_directory_path)]
+FilePath = str  # Path or directory
+FsLoc = str
+
+
+class BuiltinFileSystem(BaseModel):
+    """Spiral supports several builtin file systems in different cloud provider regions."""
+
+    type: Literal["builtin"] = "builtin"
+    provider: str
+
+
+class UpstreamFileSystem(BaseModel):
+    """File system that points to another project, usually a "file system" project.
+
+    Upstream project must have an external file system configured,
+    and not a builtin file system or another upstream file system.
+    """
+
+    type: Literal["upstream"] = "upstream"
+    project_id: ProjectId
+
+
+class S3FileSystem(BaseModel):
+    """File system backed by an S3-compatible bucket."""
+
+    type: Literal["s3"] = "s3"
+    endpoint: str | None = None
+    region: str
+    bucket: str
+    directory: DirectoryPath | None = None
+    # ARN of the role to assume when accessing the bucket https://docs.spiraldb.com/filesystems#aws
+    # role_arn: str | None = None
+    role_arn: str  # TODO(marko): Make optional once we support third-party S3-compatible storage.
+
+
+class GCSFileSystem(BaseModel):
+    """File system backed by a Google Cloud Storage bucket."""
+
+    type: Literal["gcs"] = "gcs"
+    region: str
+    bucket: str
+    directory: DirectoryPath | None = None
+
+
+FileSystem = Annotated[
+    BuiltinFileSystem | UpstreamFileSystem | S3FileSystem | GCSFileSystem, Field(discriminator="type")
+]
+
+
+class Mode(str, Enum):
+    READ_ONLY = "ro"
+    READ_WRITE = "rw"
+
+
+class Mount(BaseModel):
+    """Mount grants permission to a Spiral resource to use a specific directory within the file system."""
+
+    id: str
+    project_id: ProjectId
+    directory: DirectoryPath
+    mode: Mode
+    principal: str
+
+
+class CreateMountRequest(BaseModel):
+    directory: DirectoryPath
+    mode: Mode
+    principal: str
+
+
+class CreateMountResponse(BaseModel):
+    mount: Mount
+
+
+class GetMountAndFileSystemResponse(BaseModel):
+    mount: Mount
+    file_system: FileSystem
+    fs_loc: FsLoc
+
+
+class FileSystemService(ServiceBase):
+    """Service for file system operations."""
+
+    def list_providers(self) -> list[str]:
+        """List builtin providers."""
+        response = self.client.get("/v1/file-systems/builtin-providers", dict)
+        return response.get("providers", [])
+
+    def update_file_system(self, project_id: ProjectId, request: FileSystem) -> FileSystem:
+        """Update project's default file system."""
+        return self.client.post(f"/v1/file-systems/{project_id}", request, FileSystem)
+
+    def get_file_system(self, project_id: ProjectId) -> FileSystem:
+        """Get project's default file system."""
+        return self.client.get(f"/v1/file-systems/{project_id}", FileSystem)
+
+    def create_mount(self, project_id: ProjectId, request: CreateMountRequest) -> CreateMountResponse:
+        """Create a mount."""
+        return self.client.post(f"/v1/file-systems/{project_id}/mounts", request, CreateMountResponse)
+
+    def list_mounts(self, project_id: ProjectId) -> Paged[Mount]:
+        """List active mounts in project's file system."""
+        return self.client.paged(f"/v1/file-systems/{project_id}/mounts", PagedResponse[Mount])
+
+    def get_mount(self, mount_id: str) -> Mount:
+        """Get a mount."""
+        return self.client.get(f"/v1/mounts/{mount_id}", Mount)
+
+    def get_mount_and_file_system(self, mount_id: str) -> GetMountAndFileSystemResponse:
+        """Get the mount and its associated file system."""
+        return self.client.get(f"/v1/mounts/{mount_id}/with-filesystem", GetMountAndFileSystemResponse)
+
+    def remove_mount(self, mount_id: str) -> None:
+        """Remove mount."""
+        return self.client.delete(f"/v1/mounts/{mount_id}", NoneType)

spiral/api/key_space_indexes.py

@@ -0,0 +1,23 @@
+from pydantic import BaseModel
+
+from .client import ServiceBase
+from .types import IndexId, WorkerId
+from .workers import ResourceClass
+
+
+class SyncIndexRequest(BaseModel):
+    """Request to sync a text index."""
+
+    resources: ResourceClass
+
+
+class SyncIndexResponse(BaseModel):
+    worker_id: WorkerId
+
+
+class KeySpaceIndexesService(ServiceBase):
+    """Service for key space index operations."""
+
+    def sync_index(self, index_id: IndexId, request: SyncIndexRequest) -> SyncIndexResponse:
+        """Start a job to sync an index."""
+        return self.client.post(f"/v1/key-space-indexes/{index_id}/sync", request, SyncIndexResponse)

spiral/api/organizations.py

@@ -0,0 +1,77 @@
+from enum import Enum
+
+from pydantic import BaseModel
+
+from .client import Paged, PagedResponse, ServiceBase
+from .types import OrgId
+
+
+class OrgRole(str, Enum):
+    OWNER = "owner"
+    MEMBER = "member"
+    GUEST = "guest"
+
+
+class Org(BaseModel):
+    id: OrgId
+    name: str | None = None
+
+
+class OrgMembership(BaseModel):
+    user_id: str
+    org: Org
+    role: str
+
+
+class CreateOrgRequest(BaseModel):
+    name: str | None = None
+
+
+class CreateOrgResponse(BaseModel):
+    org: Org
+
+
+class PortalLinkIntent(str, Enum):
+    SSO = "sso"
+    DIRECTORY_SYNC = "directory-sync"
+    AUDIT_LOGS = "audit-logs"
+    LOG_STREAMS = "log-streams"
+    DOMAIN_VERIFICATION = "domain-verification"
+
+
+class PortalLinkRequest(BaseModel):
+    intent: PortalLinkIntent
+
+
+class PortalLinkResponse(BaseModel):
+    url: str
+
+
+class InviteUserRequest(BaseModel):
+    email: str
+    role: OrgRole
+    expires_in_days: int | None = 7
+
+
+class InviteUserResponse(BaseModel):
+    invite_id: str
+
+
+class OrganizationService(ServiceBase):
+    """Service for organization operations."""
+
+    def create(self, request: CreateOrgRequest) -> CreateOrgResponse:
+        """Create a new organization."""
+        return self.client.post("/v1/organizations", request, CreateOrgResponse)
+
+    def list_memberships(self) -> Paged[OrgMembership]:
+        """List organization memberships."""
+        return self.client.paged("/v1/organizations", PagedResponse[OrgMembership])
+
+    def invite_user(self, request: InviteUserRequest) -> InviteUserResponse:
+        """Invite a user to the organization."""
+        return self.client.post("/v1/organizations/invite-user", request, InviteUserResponse)
+
+    def portal_link(self, request: PortalLinkRequest) -> PortalLinkResponse:
+        """Get configuration portal link for the organization."""
+        return self.client.put("/v1/organizations/portal-link", request, PortalLinkResponse)