pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/api/client.py

@@ -0,0 +1,160 @@
+import logging
+from collections.abc import Iterable, Iterator, Mapping
+from typing import Any, Generic, TypeVar
+
+import httpx
+from httpx import HTTPStatusError
+from pydantic import BaseModel, Field, TypeAdapter
+
+from spiral.core.client import Authn
+
+log = logging.getLogger(__name__)
+
+
+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,
+        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"))
+
+        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
+
+        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,153 @@
+from enum import Enum
+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
+
+
+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 = "https://s3.amazonaws.com"
+    region: str = "auto"
+    bucket: str
+    directory: DirectoryPath | None
+
+
+class GCSFileSystem(BaseModel):
+    """File system backed by a Google Cloud Storage bucket."""
+
+    type: Literal["gcs"] = "gcs"
+    region: str
+    bucket: str
+    directory: DirectoryPath | 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 AWSSecretAccessKey(BaseModel):
+    """AWS secret access key credentials to be used with an S3 file system.
+    The access key must have read/write access to the bucket specified in the file system.
+    """
+
+    access_key_id: str
+    secret_access_key: str
+
+
+class UpdateS3FileSystem(S3FileSystem):
+    credentials: AWSSecretAccessKey
+
+
+class GCPServiceAccount(BaseModel):
+    """Google Cloud Platform service account credentials to be used with a GCS file system.
+    The service account must have read/write access to the bucket specified in the file system.
+    """
+
+    service_account: str
+
+
+class UpdateGCSFileSystem(GCSFileSystem):
+    credentials: GCPServiceAccount
+
+
+UpdateFileSystemRequest = Annotated[
+    BuiltinFileSystem | UpstreamFileSystem | UpdateS3FileSystem | UpdateGCSFileSystem, Field(discriminator="type")
+]
+
+
+class UpdateFileSystemResponse(BaseModel):
+    file_system: FileSystem
+
+
+class CreateMountRequest(BaseModel):
+    directory: DirectoryPath
+    mode: Mode
+    principal: str
+
+
+class CreateMountResponse(BaseModel):
+    mount: Mount
+
+
+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: UpdateFileSystemRequest) -> UpdateFileSystemResponse:
+        """Update project's default file system."""
+        return self.client.post(f"/v1/file-systems/{project_id}", request, UpdateFileSystemResponse)
+
+    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 remove_mount(self, mount_id: str) -> None:
+        """Remove mount."""
+        return self.client.delete(f"/v1/mounts/{mount_id}", None)

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)

spiral/api/projects.py

@@ -0,0 +1,197 @@
+from typing import Annotated, Literal
+
+from pydantic import BaseModel, Field
+
+from .client import Paged, PagedResponse, ServiceBase
+from .types import OrgId, ProjectId, RoleId
+
+
+class Project(BaseModel):
+    id: ProjectId
+    org_id: OrgId
+    name: str | None = None
+
+
+class CreateProjectRequest(BaseModel):
+    id_prefix: str | None = None
+    name: str | None = None
+
+
+class CreateProjectResponse(BaseModel):
+    project: Project
+
+
+class Grant(BaseModel):
+    id: str
+    project_id: ProjectId
+    role_id: RoleId
+    principal: str
+    conditions: dict | None = None
+
+
+class OrgRolePrincipalConditions(BaseModel):
+    type: Literal["org_role"] = "org_role"
+    org_id: OrgId
+    role: str
+
+
+class OrgUserPrincipalConditions(BaseModel):
+    type: Literal["org_user"] = "org_user"
+    org_id: OrgId
+    user_id: str
+
+
+class WorkloadPrincipalConditions(BaseModel):
+    type: Literal["workload"] = "workload"
+    workload_id: str
+
+
+class GitHubConditions(BaseModel):
+    environment: str | None = None
+    ref: str | None = None
+    ref_type: str | None = None
+    sha: str | None = None
+    repository: str | None = None
+    repository_owner: str | None = None
+    repository_visibility: str | None = None
+    repository_id: str | None = None
+    repository_owner_id: str | None = None
+    run_id: str | None = None
+    run_number: str | None = None
+    run_attempt: str | None = None
+    runner_environment: str | None = None
+    actor_id: str | None = None
+    actor: str | None = None
+    workflow: str | None = None
+    head_ref: str | None = None
+    base_ref: str | None = None
+    job_workflow_ref: str | None = None
+    event_name: str | None = None
+
+
+class GitHubPrincipalConditions(BaseModel):
+    type: Literal["github"] = "github"
+    org: str
+    repo: str
+    conditions: GitHubConditions | None = None
+
+
+class ModalConditions(BaseModel):
+    app_id: str | None = None
+    app_name: str | None = None
+    function_id: str | None = None
+    function_name: str | None = None
+    container_id: str | None = None
+
+
+class ModalPrincipalConditions(BaseModel):
+    type: Literal["modal"] = "modal"
+
+    # 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).
+    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
+
+    conditions: ModalConditions | None = None
+
+
+class GCPPrincipalConditions(BaseModel):
+    type: Literal["gcp"] = "gcp"
+    service_account: str
+    unique_id: str
+
+
+PrincipalConditions = Annotated[
+    OrgRolePrincipalConditions
+    | OrgUserPrincipalConditions
+    | WorkloadPrincipalConditions
+    | GitHubPrincipalConditions
+    | ModalPrincipalConditions
+    | GCPPrincipalConditions,
+    Field(discriminator="type"),
+]
+
+
+class GrantRoleRequest(BaseModel):
+    role_id: RoleId
+    principal: PrincipalConditions
+
+
+class GrantRoleResponse(BaseModel):
+    grant: Grant
+
+
+class TableResource(BaseModel):
+    id: str
+    project_id: ProjectId
+    dataset: str
+    table: str
+
+
+class TextIndexResource(BaseModel):
+    id: str
+    project_id: ProjectId
+    name: str
+
+
+class ProjectService(ServiceBase):
+    """Service for project operations."""
+
+    def create(self, request: CreateProjectRequest) -> CreateProjectResponse:
+        """Create a new project."""
+        return self.client.post("/v1/projects", request, CreateProjectResponse)
+
+    def list(self) -> Paged[Project]:
+        """List projects."""
+        return self.client.paged("/v1/projects", PagedResponse[Project])
+
+    def list_tables(
+        self, project_id: ProjectId, dataset: str | None = None, table: str | None = None
+    ) -> Paged[TableResource]:
+        """List tables in a project."""
+        params = {}
+        if dataset:
+            params["dataset"] = dataset
+        if table:
+            params["table"] = table
+        return self.client.paged(f"/v1/projects/{project_id}/tables", PagedResponse[TableResource], params=params)
+
+    def list_text_indexes(self, project_id: ProjectId, name: str | None = None) -> Paged[TextIndexResource]:
+        """List text indexes in a project."""
+        params = {}
+        if name:
+            params["name"] = name
+        return self.client.paged(
+            f"/v1/projects/{project_id}/text-indexes", PagedResponse[TextIndexResource], params=params
+        )
+
+    def get(self, project_id: ProjectId) -> Project:
+        """Get a project."""
+        return self.client.get(f"/v1/projects/{project_id}", Project)
+
+    def grant_role(self, project_id: ProjectId, request: GrantRoleRequest) -> GrantRoleResponse:
+        """Grant a role to a principal."""
+        return self.client.post(f"/v1/projects/{project_id}/grants", request, GrantRoleResponse)
+
+    def list_grants(
+        self,
+        project_id: ProjectId,
+        principal: str | None = None,
+    ) -> Paged[Grant]:
+        """List active project grants."""
+        params = {}
+        if principal:
+            params["principal"] = principal
+        return self.client.paged(f"/v1/projects/{project_id}/grants", PagedResponse[Grant], params=params)
+
+    def get_grant(self, grant_id: str) -> Grant:
+        """Get a grant."""
+        return self.client.get(f"/v1/grants/{grant_id}", Grant)
+
+    def revoke_grant(self, grant_id: str) -> None:
+        """Revoke a grant."""
+        return self.client.delete(f"/v1/grants/{grant_id}", None)

spiral/api/telemetry.py

@@ -0,0 +1,19 @@
+from pydantic import BaseModel
+
+from .client import ServiceBase
+
+
+class IssueExportTokenRequest(BaseModel):
+    pass
+
+
+class IssueExportTokenResponse(BaseModel):
+    token: str
+
+
+class TelemetryService(ServiceBase):
+    """Service for telemetry operations."""
+
+    def issue_export_token(self) -> IssueExportTokenResponse:
+        """Issue telemetry export token."""
+        return self.client.put("/v1/telemetry/token", IssueExportTokenRequest(), IssueExportTokenResponse)

spiral/api/types.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+from pydantic import AfterValidator, StringConstraints
+
+
+def _validate_root_uri(uri: str) -> str:
+    if uri.endswith("/"):
+        raise ValueError("Root URI must not end with a slash.")
+    return uri
+
+
+UserId = str
+OrgId = str
+ProjectId = str
+RoleId = str
+
+RootUri = Annotated[str, AfterValidator(_validate_root_uri)]
+DatasetName = Annotated[str, StringConstraints(max_length=128, pattern=r"^[a-zA-Z_][a-zA-Z0-9_-]+$")]
+TableName = Annotated[str, StringConstraints(max_length=128, pattern=r"^[a-zA-Z_][a-zA-Z0-9_-]*$")]
+IndexName = Annotated[str, StringConstraints(max_length=128, pattern=r"^[a-zA-Z_][a-zA-Z0-9_-]*$")]

spiral/api/workloads.py

@@ -0,0 +1,52 @@
+from pydantic import BaseModel
+
+from .client import Paged, PagedResponse, ServiceBase
+from .types import ProjectId
+
+
+class Workload(BaseModel):
+    id: str
+    project_id: ProjectId
+    name: str | None = None
+
+
+class CreateWorkloadRequest(BaseModel):
+    name: str | None = None
+
+
+class CreateWorkloadResponse(BaseModel):
+    workload: Workload
+
+
+class IssueWorkloadCredentialsResponse(BaseModel):
+    client_id: str
+    client_secret: str
+    revoked_client_id: str | None = None
+
+
+class WorkloadService(ServiceBase):
+    """Service for workload operations."""
+
+    def create(self, project_id: ProjectId, request: CreateWorkloadRequest) -> CreateWorkloadResponse:
+        """Create a new workload."""
+        return self.client.post(f"/v1/projects/{project_id}/workloads", request, CreateWorkloadResponse)
+
+    def list(self, project_id: ProjectId) -> Paged[Workload]:
+        """List active project workloads."""
+        return self.client.paged(f"/projects/{project_id}/workloads", PagedResponse[Workload])
+
+    def get(self, workload_id: str) -> Workload:
+        """Get a workload."""
+        return self.client.get(f"/v1/workloads/{workload_id}", Workload)
+
+    def deactivate(self, workload_id: str) -> None:
+        """De-activate a workload."""
+        return self.client.delete(f"/v1/workloads/{workload_id}", None)
+
+    def issue_credentials(self, workload_id: str) -> IssueWorkloadCredentialsResponse:
+        """Issue workload credentials."""
+        return self.client.post(f"/v1/workloads/{workload_id}/credentials", None, IssueWorkloadCredentialsResponse)
+
+    def revoke_credentials(self, client_id: str) -> None:
+        """Revoke workload credentials."""
+        return self.client.delete(f"/v1/credentials/{client_id}", None)