pyspiral 0.6.6__cp312-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl

spiral/api/projects.py

@@ -0,0 +1,219 @@
+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 GcpServiceAccountPrincipalConditions(BaseModel):
+    type: Literal["gcp"] = "gcp"
+    service_account: str
+    unique_id: str
+
+
+class AwsAssumedRolePrincipalConditions(BaseModel):
+    type: Literal["aws"] = "aws"
+    account_id: str
+    role_name: str
+
+
+PrincipalConditions = Annotated[
+    OrgRolePrincipalConditions
+    | OrgUserPrincipalConditions
+    | WorkloadPrincipalConditions
+    | GitHubPrincipalConditions
+    | ModalPrincipalConditions
+    | GcpServiceAccountPrincipalConditions
+    | AwsAssumedRolePrincipalConditions,
+    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 KeySpaceIndexResource(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 list_key_space_indexes(self, project_id: ProjectId, name: str | None = None) -> Paged[KeySpaceIndexResource]:
+        """List key space indexes in a project."""
+        params = {}
+        if name:
+            params["name"] = name
+        return self.client.paged(
+            f"/v1/projects/{project_id}/key-space-indexes", PagedResponse[KeySpaceIndexResource], 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):
+        """Revoke a grant."""
+        return self.client.delete(f"/v1/grants/{grant_id}", type[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/text_indexes.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel
+
+from .client import Paged, PagedResponse, ServiceBase
+from .types import IndexId, ProjectId, WorkerId
+from .workers import CPU, GcpRegion, Memory, ResourceClass
+
+
+class TextSearchWorker(BaseModel):
+    worker_id: WorkerId
+    project_id: ProjectId
+    index_id: IndexId
+    url: str | None
+
+
+class CreateWorkerRequest(BaseModel):
+    cpu: CPU
+    memory: Memory
+    region: GcpRegion
+
+
+class CreateWorkerResponse(BaseModel):
+    worker_id: WorkerId
+
+
+class SyncIndexRequest(BaseModel):
+    """Request to sync a text index."""
+
+    resources: ResourceClass
+
+
+class SyncIndexResponse(BaseModel):
+    worker_id: WorkerId
+
+
+class TextIndexesService(ServiceBase):
+    """Service for text index operations."""
+
+    def create_worker(self, index_id: IndexId, request: CreateWorkerRequest) -> CreateWorkerResponse:
+        """Create a new search worker."""
+        return self.client.post(f"/v1/text-indexes/{index_id}/workers", request, CreateWorkerResponse)
+
+    def list_workers(self, index_id: IndexId) -> Paged[WorkerId]:
+        """List text index workers for the given index."""
+        return self.client.paged(f"/v1/text-indexes/{index_id}/workers", PagedResponse[WorkerId])
+
+    def get_worker(self, worker_id: WorkerId) -> TextSearchWorker:
+        """Get a text index worker."""
+        return self.client.get(f"/v1/text-index-workers/{worker_id}", TextSearchWorker)
+
+    def shutdown_worker(self, worker_id: WorkerId) -> None:
+        """Shutdown a text index worker."""
+        return self.client.delete(f"/v1/text-index-workers/{worker_id}", type[None])
+
+    def sync_index(self, index_id: IndexId, request: SyncIndexRequest) -> SyncIndexResponse:
+        """Start a job to sync an index."""
+        return self.client.post(f"/v1/text-indexes/{index_id}/sync", request, SyncIndexResponse)

spiral/api/types.py

@@ -0,0 +1,22 @@
+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
+IndexId = str
+WorkerId = 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/workers.py

@@ -0,0 +1,40 @@
+from enum import Enum, IntEnum
+
+
+class CPU(IntEnum):
+    ONE = 1
+    TWO = 2
+    FOUR = 4
+    EIGHT = 8
+
+    def __str__(self):
+        return str(self.value)
+
+
+class Memory(str, Enum):
+    MB_512 = "512Mi"
+    GB_1 = "1Gi"
+    GB_2 = "2Gi"
+    GB_4 = "4Gi"
+    GB_8 = "8Gi"
+
+    def __str__(self):
+        return self.value
+
+
+class GcpRegion(str, Enum):
+    US_EAST4 = "us-east4"
+    EUROPE_WEST4 = "europe-west4"
+
+    def __str__(self):
+        return self.value
+
+
+class ResourceClass(str, Enum):
+    """Resource class for text index sync."""
+
+    SMALL = "small"
+    LARGE = "large"
+
+    def __str__(self):
+        return self.value

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)

spiral/arrow_.py

@@ -0,0 +1,216 @@
+from collections import defaultdict
+from collections.abc import Callable, Iterable
+from functools import reduce
+
+import numpy as np
+import pyarrow as pa
+from pyarrow import compute as pc
+
+
+def arange(*args, **kwargs) -> pa.Array:
+    return pa.array(np.arange(*args, **kwargs), type=pa.int32())
+
+
+def zip_tables(tables: Iterable[pa.Table]) -> pa.Table:
+    data = []
+    names = []
+    for table in tables:
+        data.extend(table.columns)
+        names.extend(table.column_names)
+    return pa.Table.from_arrays(data, names=names)
+
+
+def merge_arrays(*arrays: pa.StructArray) -> pa.StructArray:
+    """Recursively merge arrays into nested struct arrays."""
+    if len(arrays) == 1:
+        return arrays[0]
+
+    nstructs = sum(pa.types.is_struct(a.type) for a in arrays)
+    if nstructs == 0:
+        # Then we have conflicting arrays and we choose the last.
+        return arrays[-1]
+
+    if nstructs != len(arrays):
+        raise ValueError("Cannot merge structs with non-structs.")
+
+    data = defaultdict(list)
+    for array in arrays:
+        if isinstance(array, pa.ChunkedArray):
+            array = array.combine_chunks()
+        for field in array.type:
+            data[field.name].append(array.field(field.name))
+
+    return pa.StructArray.from_arrays([merge_arrays(*v) for v in data.values()], names=list(data.keys()))
+
+
+def merge_scalars(*scalars: pa.StructScalar) -> pa.StructScalar:
+    """Recursively merge scalars into nested struct scalars."""
+    if len(scalars) == 1:
+        return scalars[0]
+
+    nstructs = sum(pa.types.is_struct(a.type) for a in scalars)
+    if nstructs == 0:
+        # Then we have conflicting scalars and we choose the last.
+        return scalars[-1]
+
+    if nstructs != len(scalars):
+        raise ValueError("Cannot merge scalars with non-scalars.")
+
+    data = defaultdict(list)
+    for scalar in scalars:
+        for field in scalar.type:
+            data[field.name].append(scalar[field.name])
+
+    return pa.scalar({k: merge_scalars(*v) for k, v in data.items()})
+
+
+def null_table(schema: pa.Schema, length: int = 0) -> pa.Table:
+    # We add an extra nulls column to ensure the length is correctly applied.
+    return pa.table(
+        [pa.nulls(length, type=field.type) for field in schema] + [pa.nulls(length)],
+        schema=pa.schema(list(schema) + [pa.field("__", type=pa.null())]),
+    ).drop(["__"])
+
+
+def coalesce_all(table: pa.Table) -> pa.Table:
+    """Coalesce all columns that share the same name."""
+    columns: dict[str, list[pa.Array]] = defaultdict(list)
+    for i, col in enumerate(table.column_names):
+        columns[col].append(table[i])
+
+    data = []
+    names = []
+    for col, arrays in columns.items():
+        names.append(col)
+        if len(arrays) == 1:
+            data.append(arrays[0])
+        else:
+            data.append(pc.coalesce(*arrays))
+
+    return pa.Table.from_arrays(data, names=names)
+
+
+def join(left: pa.Table, right: pa.Table, keys: list[str], join_type: str) -> pa.Table:
+    """Arrow's builtin join doesn't support struct columns. So we join ourselves and zip them in."""
+    # TODO(ngates): if join_type == inner, we may have better luck performing two index_in operations since this
+    #   also preserves sort order.
+    lhs = left.select(keys).add_column(0, "__lhs", arange(len(left)))
+    rhs = right.select(keys).add_column(0, "__rhs", arange(len(right)))
+    joined = lhs.join(rhs, keys=keys, join_type=join_type).sort_by([(k, "ascending") for k in keys])
+    return zip_tables(
+        [joined.select(keys), left.take(joined["__lhs"]).drop(keys), right.take(joined["__rhs"]).drop(keys)]
+    )
+
+
+def nest_structs(array: pa.StructArray | pa.StructScalar | dict) -> dict:
+    """Turn a struct-like value with dot-separated column names into a nested dictionary."""
+    data = {}
+
+    if isinstance(array, pa.StructArray | pa.StructScalar):
+        array = {f.name: field(array, f.name) for f in array.type}
+
+    for name in array.keys():
+        if "." not in name:
+            data[name] = array[name]
+            continue
+
+        parts = name.split(".")
+        child_data = data
+        for part in parts[:-1]:
+            if part not in child_data:
+                child_data[part] = {}
+            child_data = child_data[part]
+        child_data[parts[-1]] = array[name]
+
+    return data
+
+
+def flatten_struct_table(table: pa.Table, separator=".") -> pa.Table:
+    """Turn a nested struct table into a flat table with dot-separated names."""
+    data = []
+    names = []
+
+    def _unfold(array: pa.Array, prefix: str):
+        if pa.types.is_struct(array.type):
+            if isinstance(array, pa.ChunkedArray):
+                array = array.combine_chunks()
+            for f in array.type:
+                _unfold(field(array, f.name), f"{prefix}{separator}{f.name}")
+        else:
+            data.append(array)
+            names.append(prefix)
+
+    for col in table.column_names:
+        _unfold(table[col], col)
+
+    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))
+
+
+def dict_to_struct_array(data: dict | pa.StructArray, propagate_nulls: bool = False) -> pa.StructArray:
+    """Convert a nested dictionary of arrays to a table with nested structs."""
+    if isinstance(data, pa.Array):
+        return data
+    arrays = [dict_to_struct_array(value) for value in data.values()]
+    return pa.StructArray.from_arrays(
+        arrays,
+        names=list(data.keys()),
+        mask=reduce(pc.and_, [pc.is_null(array) for array in arrays]) if propagate_nulls else None,
+    )
+
+
+def struct_array_to_dict[T](array: pa.StructArray, array_fn: Callable[[pa.Array], T] = lambda a: a) -> dict | T:
+    """Convert a struct array to a nested dictionary."""
+    if not pa.types.is_struct(array.type):
+        return array_fn(array)
+    if isinstance(array, pa.ChunkedArray):
+        array = array.combine_chunks()
+    return {field.name: struct_array_to_dict(array.field(i), array_fn=array_fn) for i, field in enumerate(array.type)}
+
+
+def table_to_struct_array(table: pa.Table) -> pa.StructArray:
+    if not table.num_rows:
+        return pa.array([], type=pa.struct(table.schema))
+    array = table.to_struct_array()
+    if isinstance(array, pa.ChunkedArray):
+        array = array.combine_chunks()
+    return array
+
+
+def table_from_struct_array(array: pa.StructArray | pa.ChunkedArray):
+    if len(array) == 0:
+        return null_table(pa.schema(array.type))
+    return pa.Table.from_struct_array(array)
+
+
+def field(value: pa.StructArray | pa.StructScalar, name: str) -> pa.Array | pa.Scalar:
+    """Get a field from a struct-like value."""
+    if isinstance(value, pa.StructScalar):
+        return value[name]
+    return value.field(name)
+
+
+def concat_tables(tables: list[pa.Table]) -> pa.Table:
+    """
+    Concatenate pyarrow.Table objects, filling "missing" data with appropriate null arrays
+    and casting arrays to the most common denominator type that fits all fields.
+    """
+    if len(tables) == 1:
+        return tables[0]
+    else:
+        return pa.concat_tables(tables, promote_options="permissive")