pyspiral 0.1.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/api/__init__.py

@@ -0,0 +1,221 @@
+import abc
+import base64
+import logging
+import os
+import re
+import socket
+import time
+from collections.abc import Iterable, Iterator
+from typing import TYPE_CHECKING, Annotated, Generic, TypeVar
+
+import betterproto
+import httpx
+import pyarrow as pa
+from httpx import HTTPStatusError
+from pydantic import (
+    BaseModel,
+    BeforeValidator,
+    Field,
+    GetPydanticSchema,
+    PlainSerializer,
+    StringConstraints,
+    TypeAdapter,
+)
+
+if TYPE_CHECKING:
+    from .admin import AdminService
+    from .filesystems import FileSystemService
+    from .organizations import OrganizationService
+    from .projects import ProjectService
+    from .tables import TableService
+    from .tokens import TokenService
+    from .workloads import WorkloadService
+
+log = logging.getLogger(__name__)
+
+RE_ID = re.compile("[a-zA-Z0-9-]+")
+
+OrganizationId = Annotated[str, StringConstraints(min_length=1, max_length=64)]  # , pattern=RE_ID)]
+ProjectId = Annotated[str, StringConstraints(min_length=1, max_length=64)]  # , pattern=RE_ID)]
+RoleId = str
+
+#: Annotations to implement pa.Schema serde with byte arrays
+ArrowSchema = Annotated[
+    pa.Schema,
+    GetPydanticSchema(lambda tp, handler: handler(object)),
+    BeforeValidator(
+        lambda v: v
+        if isinstance(v, pa.Schema)
+        else pa.ipc.read_schema(pa.ipc.read_message(base64.urlsafe_b64decode(v)))
+    ),
+    PlainSerializer(lambda schema: base64.urlsafe_b64encode(schema.serialize().to_pybytes())),
+]
+
+E = TypeVar("E")
+
+
+class PagedRequest(BaseModel):
+    page_token: str | None = None
+    page_size: int = 50
+
+
+class PagedResponse(BaseModel, Generic[E]):
+    items: list[E] = Field(default_factory=list)
+    next_page_token: str | None = None
+
+
+PagedReqT = TypeVar("PagedReqT", bound=PagedRequest)
+
+
+class Paged(Iterable[E], Generic[E]):
+    def __init__(
+        self,
+        client: "_Client",
+        path: str,
+        request: PagedRequest,
+        response_cls: type[PagedResponse[E]],
+    ):
+        self._client = client
+        self._path = path
+        self._request: PagedRequest = request
+        self._response_cls = response_cls
+
+        self._response: PagedResponse[E] = client.put(path, request, response_cls)
+
+    @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
+
+            self._request = self._request.model_copy(update=dict(page_token=self._response.next_page_token))
+            self._response = self._client.put(self._path, self._request, self._response_cls)
+
+
+class ServiceBase:
+    def __init__(self, client: "_Client"):
+        self.client = client
+
+
+class Authn:
+    """An abstract class for credential providers."""
+
+    @abc.abstractmethod
+    def token(self) -> str | None:
+        """Return a token, if available."""
+
+
+class _Client:
+    RequestT = TypeVar("RequestT")
+    ResponseT = TypeVar("ResponseT")
+
+    def __init__(self, http: httpx.Client, authn: Authn):
+        self.http = http
+        self.authn = authn
+
+    def post(self, path: str, req: RequestT, response_cls: type[ResponseT]) -> ResponseT:
+        return self.request("POST", path, req, response_cls)
+
+    def put(self, path: str, req: RequestT, response_cls: type[ResponseT]) -> ResponseT:
+        return self.request("PUT", path, req, response_cls)
+
+    def request(self, method: str, path: str, req: RequestT, response_cls: type[ResponseT]) -> ResponseT:
+        if isinstance(req, betterproto.Message):
+            try:
+                req = dict(content=bytes(req))
+            except:
+                raise
+        else:
+            req = dict(json=TypeAdapter(req.__class__).dump_python(req, mode="json") if req is not None else None)
+
+        token = self.authn.token()
+        resp = self.http.request(method, path, headers={"authorization": f"Bearer {token}" if token else None}, **req)
+
+        try:
+            resp.raise_for_status()
+        except HTTPStatusError as e:
+            # Enrich the exception with the response body
+            raise HTTPStatusError(f"{str(e)}: {resp.text}", request=e.request, response=e.response)
+
+        if issubclass(response_cls, betterproto.Message):
+            return response_cls().parse(resp.content)
+        else:
+            return TypeAdapter(response_cls).validate_python(resp.json())
+
+    def paged(self, path: str, req: PagedRequest, response_cls: type[PagedResponse[E]]) -> Paged[E]:
+        return Paged(self, path, req, response_cls)
+
+
+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=60,
+                # 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 file_system(self) -> "FileSystemService":
+        from .filesystems import FileSystemService
+
+        return FileSystemService(self.client)
+
+    @property
+    def organization(self) -> "OrganizationService":
+        from .organizations import OrganizationService
+
+        return OrganizationService(self.client)
+
+    @property
+    def token(self) -> "TokenService":
+        from .tokens import TokenService
+
+        return TokenService(self.client)
+
+    @property
+    def project(self) -> "ProjectService":
+        from .projects import ProjectService
+
+        return ProjectService(self.client)
+
+    @property
+    def table(self) -> "TableService":
+        from .tables import TableService
+
+        return TableService(self.client)
+
+    @property
+    def workload(self) -> "WorkloadService":
+        from .workloads import WorkloadService
+
+        return WorkloadService(self.client)
+
+
+def wait_for_port(port: int, host: str = "localhost", timeout: float = 5.0):
+    """Wait until a port starts accepting TCP connections."""
+    start_time = time.time()
+    while True:
+        try:
+            with socket.create_connection((host, port), timeout=timeout):
+                break
+        except OSError as ex:
+            time.sleep(0.01)
+            if time.time() - start_time >= timeout:
+                raise TimeoutError(
+                    f"Waited too long for the port {port} on host {host} to start accepting " "connections."
+                ) from ex

spiral/api/admin.py

@@ -0,0 +1,29 @@
+from pydantic import BaseModel
+
+from . import OrganizationId, Paged, PagedRequest, PagedResponse, ServiceBase
+
+
+class SyncOrgs:
+    class Request(PagedRequest): ...
+
+    class Response(PagedResponse[OrganizationId]): ...
+
+
+class Membership(BaseModel):
+    user_id: str
+    organization_id: str
+
+
+class SyncMemberships:
+    class Request(PagedRequest):
+        organization_id: OrganizationId | None = None
+
+    class Response(PagedResponse[Membership]): ...
+
+
+class AdminService(ServiceBase):
+    def sync_orgs(self, request: SyncOrgs.Request) -> Paged[SyncOrgs.Response]:
+        return self.client.paged("/admin/sync-orgs", request, SyncOrgs.Response)
+
+    def sync_memberships(self, request: SyncMemberships.Request) -> Paged[Membership]:
+        return self.client.paged("/admin/sync-memberships", request, SyncMemberships.Response)

spiral/api/filesystems.py

@@ -0,0 +1,125 @@
+from datetime import timedelta
+from enum import Enum
+from typing import Annotated, Literal
+
+from pydantic import AfterValidator, BaseModel, Field
+
+from . import ProjectId, ServiceBase
+
+
+class BuiltinFileSystem(BaseModel):
+    type: Literal["builtin"] = "builtin"
+    provider: str
+
+
+FileSystem = Annotated[BuiltinFileSystem, Field(discriminator="type")]
+
+
+def _validate_file_path(path: str) -> str:
+    if "//" in path:
+        raise ValueError("FilePath must not contain multiple slashes")
+    if not path.startswith("/"):
+        raise ValueError("FilePath must start with /")
+    if path.endswith("/"):
+        raise ValueError("FilePath must not end with /")
+    return path
+
+
+FilePath = Annotated[
+    str,
+    AfterValidator(_validate_file_path),
+    Field(description="A path to an individual file in the file system. Must not end with a slash."),
+]
+
+
+def _validate_prefix(path: str) -> str:
+    if "//" in path:
+        raise ValueError("Prefix must not contain multiple slashes")
+    if not path.startswith("/"):
+        raise ValueError("Prefix must start with /")
+    if not path.endswith("/"):
+        raise ValueError("Prefix must end with /")
+    return path
+
+
+Prefix = Annotated[
+    str,
+    AfterValidator(_validate_prefix),
+    Field(description="A prefix in a file system. Must end with a slash."),
+]
+
+
+class Mode(Enum):
+    READ_ONLY = "ro"
+    READ_WRITE = "rw"
+
+
+class Mount(BaseModel):
+    id: str
+    project_id: ProjectId
+    prefix: Prefix
+    mode: Mode
+    principal: str
+
+
+class GetFileSystem:
+    class Request(BaseModel):
+        project_id: ProjectId
+
+    class Response(BaseModel):
+        file_system: FileSystem
+
+
+class UpdateFileSystem:
+    class Request(BaseModel):
+        project_id: ProjectId
+        file_system: FileSystem
+
+    class Response(BaseModel):
+        file_system: FileSystem
+
+
+class ListProviders:
+    class Response(BaseModel):
+        providers: list[str]
+
+
+class CreateMount:
+    class Request(BaseModel):
+        project_id: ProjectId
+        prefix: Prefix
+        mode: Mode
+        principal: str
+
+    class Response(BaseModel):
+        mount: Mount
+
+
+class CreateMountToken:
+    class Request(BaseModel):
+        mount_id: str
+        mode: Mode
+        path: FilePath | Prefix
+        ttl: int = Field(default_factory=lambda: int(timedelta(hours=1).total_seconds()))
+
+    class Response(BaseModel):
+        token: str
+
+
+class FileSystemService(ServiceBase):
+    def get_file_system(self, request: GetFileSystem.Request) -> GetFileSystem.Response:
+        """Get the file system currently configured for a project."""
+        return self.client.put("/file-system/get", request, GetFileSystem.Response)
+
+    def update_file_system(self, request: UpdateFileSystem.Request) -> UpdateFileSystem.Response:
+        """Update the file system for a project."""
+        return self.client.put("/file-system/update", request, UpdateFileSystem.Response)
+
+    def list_providers(self) -> ListProviders.Response:
+        return self.client.put("/file-system/list-providers", None, ListProviders.Response)
+
+    def create_mount(self, request: CreateMount.Request) -> CreateMount.Response:
+        return self.client.post("/file-system/create-mount", request, CreateMount.Response)
+
+    def create_mount_token(self, request: CreateMountToken.Request) -> CreateMountToken.Response:
+        return self.client.post("/file-system/create-mount-token", request, CreateMountToken.Response)

spiral/api/organizations.py

@@ -0,0 +1,90 @@
+from enum import Enum
+
+from pydantic import BaseModel, EmailStr, Field
+
+from . import OrganizationId, Paged, PagedRequest, PagedResponse, ServiceBase
+
+
+class OrganizationRole(Enum):
+    OWNER = "owner"
+    MEMBER = "member"
+    GUEST = "guest"
+
+
+class Organization(BaseModel):
+    id: OrganizationId
+    name: str | None = Field(
+        default=None,
+        description="Optional human-readable name for the organization",
+    )
+
+
+class OrganizationMembership(BaseModel):
+    organization: Organization
+    role: str = Field(description="The user's role in the organization")
+
+
+class CreateOrganization:
+    class Request(BaseModel):
+        name: str | None = Field(
+            default=None,
+            description="Optional human-readable name for the organization",
+        )
+
+    class Response(BaseModel):
+        organization: Organization
+
+
+class ListUserMemberships:
+    class Request(PagedRequest):
+        """List the organization memberships of the current user."""
+
+    class Response(PagedResponse[OrganizationMembership]):
+        """The user's organization memberships."""
+
+
+class PortalLink:
+    class Intent(Enum):
+        SSO = "sso"
+        DIRECTORY = "directory"
+        AUDIT_LOGS = "audit-logs"
+        LOG_STREAMS = "log-streams"
+        DOMAIN_VERIFICATION = "domain-verification"
+
+    class Request(BaseModel):
+        intent: "PortalLink.Intent"
+
+    class Response(BaseModel):
+        url: str
+
+
+class InviteUser:
+    class Request(BaseModel):
+        email: EmailStr
+        role: OrganizationRole
+        expires_in_days: int = 7
+
+    class Response(BaseModel):
+        invite_id: str
+
+
+class OrganizationService(ServiceBase):
+    def create_organization(self, request: CreateOrganization.Request) -> CreateOrganization.Response:
+        """Create a new organization."""
+        return self.client.post("/organization/create", request, CreateOrganization.Response)
+
+    def list_user_memberships(self) -> Paged[OrganizationMembership]:
+        """List organizations that the user is a member of."""
+        return self.client.paged(
+            "/organization/list-user-memberships",
+            ListUserMemberships.Request(),
+            ListUserMemberships.Response,
+        )
+
+    def portal_link(self, request: PortalLink.Request) -> PortalLink.Response:
+        """Get a link to the organization configuration portal."""
+        return self.client.put("/organization/portal-link", request, PortalLink.Response)
+
+    def invite_user(self, request: InviteUser.Request) -> InviteUser.Response:
+        """Invite a user to the organization."""
+        return self.client.post("/organization/invite-user", request, InviteUser.Response)

spiral/api/projects.py

@@ -0,0 +1,160 @@
+from enum import Enum
+from typing import Annotated, Literal, Union
+
+from pydantic import BaseModel, Field
+
+from . import OrganizationId, Paged, PagedRequest, PagedResponse, ProjectId, RoleId, ServiceBase
+from .organizations import OrganizationRole
+
+
+class Project(BaseModel):
+    id: ProjectId
+    organization_id: OrganizationId
+    name: str | None = None
+
+
+class Grant(BaseModel):
+    id: str
+    project_id: ProjectId
+    role_id: RoleId
+    principal: str
+
+
+class CreateProject:
+    class Request(BaseModel):
+        organization_id: OrganizationId
+        id_prefix: str | None = Field(default=None, description="Optional prefix for a random project ID")
+        name: str | None = Field(default=None, description="Optional human-readable name for the project")
+
+        id: str | None = Field(
+            default=None,
+            description="Exact project ID to use. Requires elevated permissions.",
+        )
+
+    class Response(BaseModel):
+        project: Project
+
+
+class GetProject:
+    class Request(BaseModel):
+        id: ProjectId
+
+    class Response(BaseModel):
+        project: Project
+
+
+class ListProjects:
+    class Request(PagedRequest): ...
+
+    class Response(PagedResponse[Project]): ...
+
+
+class ListGrants:
+    class Request(PagedRequest):
+        project_id: ProjectId
+
+    class Response(PagedResponse[Grant]): ...
+
+
+class GrantRole:
+    class Request(BaseModel):
+        project_id: ProjectId
+        role_id: RoleId
+        principal: Annotated[
+            Union[
+                "GrantRole.OrgRolePrincipal",
+                "GrantRole.OrgUserPrincipal",
+                "GrantRole.WorkloadPrincipal",
+                "GrantRole.GitHubPrincipal",
+                "GrantRole.ModalPrincipal",
+            ],
+            Field(discriminator="type"),
+        ]
+
+    class Response(BaseModel): ...
+
+    class OrgRolePrincipal(BaseModel):
+        type: Literal["org"] = "org"
+
+        org_id: OrganizationId
+        role: OrganizationRole
+
+    class OrgUserPrincipal(BaseModel):
+        type: Literal["org_user"] = "org_user"
+
+        org_id: OrganizationId
+        user_id: str
+
+    class WorkloadPrincipal(BaseModel):
+        type: Literal["workload"] = "workload"
+
+        workload_id: str
+
+    class GitHubClaim(Enum):
+        environment = "environment"
+        ref = "ref"
+        sha = "sha"
+        repository = "repository"
+        repository_owner = "repository_owner"
+        actor_id = "actor_id"
+        repository_visibility = "repository_visibility"
+        repository_id = "repository_id"
+        repository_owner_id = "repository_owner_id"
+        run_id = "run_id"
+        run_number = "run_number"
+        run_attempt = "run_attempt"
+        runner_environment = "runner_environment"
+        actor = "actor"
+        workflow = "workflow"
+        head_ref = "head_ref"
+        base_ref = "base_ref"
+        event_name = "event_name"
+        ref_type = "ref_type"
+        job_workflow_ref = "job_workflow_ref"
+
+    class GitHubPrincipal(BaseModel):
+        type: Literal["github"] = "github"
+
+        org: str
+        repo: str
+        conditions: dict["GrantRole.GitHubClaim", str] | None = None
+
+    class ModalClaim(Enum):
+        workspace_id = "workspace_id"
+        environment_id = "environment_id"
+        environment_name = "environment_name"
+        app_id = "app_id"
+        app_name = "app_name"
+        function_id = "function_id"
+        function_name = "function_name"
+        container_id = "container_id"
+
+    class ModalPrincipal(BaseModel):
+        type: Literal["modal"] = "modal"
+
+        workspace_id: str
+        # Environments are sub-divisions of workspaces. Name is unique within a workspace.
+        # See https://modal.com/docs/guide/environments
+        environment_name: str
+        # A Modal App is a group of functions and classes that are deployed together.
+        # See https://modal.com/docs/guide/apps. Nick and Marko discussed having an app_name
+        # here as well and decided to leave it out for now with the assumption that people
+        # will want to authorize the whole Modal environment to access Spiral (their data).
+        conditions: dict["GrantRole.ModalClaim", str] | None = None
+
+
+class ProjectService(ServiceBase):
+    def get(self, request: GetProject.Request) -> GetProject.Response:
+        return self.client.put("/project/get", request, GetProject.Response)
+
+    def create(self, request: CreateProject.Request) -> CreateProject.Response:
+        return self.client.post("/project/create", request, CreateProject.Response)
+
+    def list(self) -> Paged[Project]:
+        return self.client.paged("/project/list", ListProjects.Request(), ListProjects.Response)
+
+    def list_grants(self, request: ListGrants.Request) -> Paged[Grant]:
+        return self.client.paged("/project/list-grants", request, ListGrants.Response)
+
+    def grant_role(self, request: GrantRole.Request) -> GrantRole.Response:
+        return self.client.post("/project/grant-role", request, GrantRole.Response)