pyspiral 0.2.5__cp310-abi3-macosx_11_0_arm64.whl → 0.4.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/api/projects.py

@@ -1,160 +1,197 @@
-from enum import Enum
-from typing import Annotated, Literal, Union
+from typing import Annotated, Literal
 
 from pydantic import BaseModel, Field
 
-from . import OrganizationId, Paged, PagedRequest, PagedResponse, ProjectId, RoleId, ServiceBase
-from .organizations import OrganizationRole
+from .client import Paged, PagedResponse, ServiceBase
+from .types import OrgId, ProjectId, RoleId
 
 
 class Project(BaseModel):
     id: ProjectId
-    organization_id: OrganizationId
+    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 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")
+class OrgRolePrincipalConditions(BaseModel):
+    type: Literal["org_role"] = "org_role"
+    org_id: OrgId
+    role: str
 
-        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 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):
-    def get(self, request: GetProject.Request) -> GetProject.Response:
-        return self.client.put("/project/get", request, GetProject.Response)
+    """Service for project operations."""
 
-    def create(self, request: CreateProject.Request) -> CreateProject.Response:
-        return self.client.post("/project/create", request, CreateProject.Response)
+    def create(self, request: CreateProjectRequest) -> CreateProjectResponse:
+        """Create a new project."""
+        return self.client.post("/v1/projects", request, CreateProjectResponse)
 
     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)
+        """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 grant_role(self, request: GrantRole.Request) -> GrantRole.Response:
-        return self.client.post("/project/grant-role", request, GrantRole.Response)
+    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

@@ -1,6 +1,7 @@
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 
-from . import Paged, PagedRequest, PagedResponse, ProjectId, ServiceBase
+from .client import Paged, PagedResponse, ServiceBase
+from .types import ProjectId
 
 
 class Workload(BaseModel):
@@ -9,37 +10,43 @@ class Workload(BaseModel):
     name: str | None = None
 
 
-class CreateWorkload:
-    class Request(BaseModel):
-        project_id: str
-        name: str | None = Field(default=None, description="Optional human-readable name for the workload")
+class CreateWorkloadRequest(BaseModel):
+    name: str | None = None
+
 
-    class Response(BaseModel):
-        workload: Workload
+class CreateWorkloadResponse(BaseModel):
+    workload: Workload
 
 
-class IssueToken:
-    class Request(BaseModel):
-        workload_id: str
+class IssueWorkloadCredentialsResponse(BaseModel):
+    client_id: str
+    client_secret: str
+    revoked_client_id: str | None = None
 
-    class Response(BaseModel):
-        token_id: str
-        token_secret: str
 
+class WorkloadService(ServiceBase):
+    """Service for workload operations."""
 
-class ListWorkloads:
-    class Request(PagedRequest):
-        project_id: str
+    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)
 
-    class Response(PagedResponse[Workload]): ...
+    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)
 
-class WorkloadService(ServiceBase):
-    def create(self, request: CreateWorkload.Request) -> CreateWorkload.Response:
-        return self.client.post("/workload/create", request, CreateWorkload.Response)
+    def deactivate(self, workload_id: str) -> None:
+        """De-activate a workload."""
+        return self.client.delete(f"/v1/workloads/{workload_id}", None)
 
-    def issue_token(self, request: IssueToken.Request) -> IssueToken.Response:
-        return self.client.post("/workload/issue-token", request, IssueToken.Response)
+    def issue_credentials(self, workload_id: str) -> IssueWorkloadCredentialsResponse:
+        """Issue workload credentials."""
+        return self.client.post(f"/v1/workloads/{workload_id}/credentials", None, IssueWorkloadCredentialsResponse)
 
-    def list(self, request: ListWorkloads.Request) -> Paged[Workload]:
-        return self.client.paged("/workload/list", request, ListWorkloads.Response)
+    def revoke_credentials(self, client_id: str) -> None:
+        """Revoke workload credentials."""
+        return self.client.delete(f"/v1/credentials/{client_id}", None)

spiral/{arrow.py → arrow_.py}

@@ -149,6 +149,18 @@ def flatten_struct_table(table: pa.Table, separator=".") -> pa.Table:
     return pa.Table.from_arrays(data, names=names)
 
 
+def struct_array(fields: list[tuple[str, bool, pa.Array]], /, mask: list[bool] | None = None) -> pa.StructArray:
+    return pa.StructArray.from_arrays(
+        arrays=[x[2] for x in fields],
+        fields=[pa.field(x[0], type=x[2].type, nullable=x[1]) for x in fields],
+        mask=pa.array(mask) if mask else mask,
+    )
+
+
+def table(fields: list[tuple[str, bool, pa.Array]], /) -> pa.Table:
+    return pa.Table.from_struct_array(struct_array(fields))
+
+
 def dict_to_table(data) -> pa.Table:
     return pa.Table.from_struct_array(dict_to_struct_array(data))
 

spiral/cli/__init__.py

@@ -1,7 +1,7 @@
 import asyncio
 import functools
 import inspect
-from typing import IO, Optional
+from typing import IO
 
 import rich
 import typer
@@ -9,9 +9,6 @@ from click import ClickException
 from grpclib import GRPCError
 from httpx import HTTPStatusError
 
-# We need to use Optional[str] since Typer doesn't support str | None.
-OptionalStr = Optional[str]  # noqa: UP007
-
 
 class AsyncTyper(typer.Typer):
     """Wrapper to allow async functions to be used as commands.
@@ -43,7 +40,7 @@ class AsyncTyper(typer.Typer):
 
 class _ClickGRPCException(ClickException):
     def __init__(self, err: GRPCError):
-        super().__init__(err.message)
+        super().__init__(err.message or "GRPCError message was None.")
         self.err = err
         self.exit_code = 1
 

spiral/cli/admin.py

@@ -1,21 +1,16 @@
 from rich import print
 
-from spiral.api.admin import SyncMemberships, SyncOrgs
+from spiral.api.types import OrgId
 from spiral.cli import AsyncTyper, state
 
 app = AsyncTyper()
 
 
 @app.command()
-def sync(orgs: bool = False, memberships: bool = False):
-    run_all = True
-    if any([orgs, memberships]):
-        run_all = False
+def sync(
+    org_id: OrgId | None = None,
+):
+    state.settings.api._admin.sync_orgs()
 
-    if run_all or orgs:
-        for org_id in state.settings.api._admin.sync_orgs(SyncOrgs.Request()):
-            print(org_id)
-
-    if run_all or memberships:
-        for membership in state.settings.api._admin.sync_memberships(SyncMemberships.Request()):
-            print(membership)
+    for membership in state.settings.api._admin.sync_memberships(org_id):
+        print(membership)

spiral/cli/app.py

@@ -2,7 +2,21 @@ import logging
 import os
 from logging.handlers import RotatingFileHandler
 
-from spiral.cli import AsyncTyper, admin, console, fs, login, org, project, state, table, token, workload
+from spiral.cli import (
+    AsyncTyper,
+    admin,
+    console,
+    fs,
+    iceberg,
+    indexes,
+    login,
+    orgs,
+    projects,
+    state,
+    tables,
+    telemetry,
+    workloads,
+)
 from spiral.settings import LOG_DIR, Settings
 
 app = AsyncTyper(name="spiral")
@@ -18,16 +32,19 @@ def _callback(verbose: bool = False):
 
 
 app.add_typer(fs.app, name="fs")
-app.add_typer(org.app, name="org")
-app.add_typer(project.app, name="project")
-app.add_typer(table.app, name="table")
-app.add_typer(workload.app, name="workload")
-app.add_typer(token.app, name="token")
+app.add_typer(orgs.app, name="orgs")
+app.add_typer(projects.app, name="projects")
+app.add_typer(iceberg.app, name="iceberg")
+app.add_typer(tables.app, name="tables")
+app.add_typer(indexes.app, name="indexes")
+app.add_typer(telemetry.app, name="telemetry")
 app.command("console")(console.command)
 app.command("login")(login.command)
+app.command("whoami")(login.whoami)
 
 # Register unless we're building docs. Because Typer docs command does not skip hidden commands...
 if not bool(os.environ.get("SPIRAL_DOCS", False)):
+    app.add_typer(workloads.app, name="workloads", hidden=True)
     app.add_typer(admin.app, name="admin", hidden=True)
     app.command("logout", hidden=True)(login.logout)
 

spiral/cli/console.py

@@ -3,7 +3,7 @@ import subprocess
 
 from spiral import Spiral
 from spiral.adbc import ADBCFlightServer, SpiralADBCServer
-from spiral.api import wait_for_port
+from spiral.server import wait_for_port
 
 
 def command():