pyspiral 0.1.0__cp310-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

spiral/api/tables.py

@@ -0,0 +1,94 @@
+from typing import Annotated
+
+from pydantic import (
+    AfterValidator,
+    BaseModel,
+    StringConstraints,
+)
+
+from . import ArrowSchema, Paged, PagedRequest, PagedResponse, ProjectId, ServiceBase
+
+
+def _validate_root_uri(uri: str) -> str:
+    if uri.endswith("/"):
+        raise ValueError("Root URI must not end with a slash.")
+    return uri
+
+
+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_-]+$")]
+
+
+class TableMetadata(BaseModel):
+    key_schema: ArrowSchema
+    root_uri: RootUri
+    spfs_mount_id: str | None = None
+
+    # TODO(marko): Randomize this on creation of metadata.
+    # Column group salt is used to compute column group IDs.
+    # It's used to ensure that column group IDs are unique
+    # across different tables, even if paths are the same.
+    # It's never modified.
+    column_group_salt: int = 0
+
+
+class Table(BaseModel):
+    id: str
+    project_id: ProjectId
+    dataset: DatasetName
+    table: TableName
+    metadata: TableMetadata
+
+
+class CreateTable:
+    class Request(BaseModel):
+        project_id: ProjectId
+        dataset: DatasetName
+        table: TableName
+        key_schema: ArrowSchema
+        root_uri: RootUri | None = None
+        exist_ok: bool = False
+
+    class Response(BaseModel):
+        table: Table
+
+
+class FindTable:
+    class Request(BaseModel):
+        project_id: ProjectId
+        dataset: DatasetName = None
+        table: TableName = None
+
+    class Response(BaseModel):
+        table: Table | None
+
+
+class GetTable:
+    class Request(BaseModel):
+        id: str
+
+    class Response(BaseModel):
+        table: Table
+
+
+class ListTables:
+    class Request(PagedRequest):
+        project_id: ProjectId
+        dataset: DatasetName | None = None
+
+    class Response(PagedResponse[Table]): ...
+
+
+class TableService(ServiceBase):
+    def create(self, req: CreateTable.Request) -> CreateTable.Response:
+        return self.client.post("/table/create", req, CreateTable.Response)
+
+    def find(self, req: FindTable.Request) -> FindTable.Response:
+        return self.client.put("/table/find", req, FindTable.Response)
+
+    def get(self, req: GetTable.Request) -> GetTable.Response:
+        return self.client.put(f"/table/{req.id}", GetTable.Response)
+
+    def list(self, req: ListTables.Request) -> Paged[Table]:
+        return self.client.paged("/table/list", req, ListTables.Response)

spiral/api/tokens.py

@@ -0,0 +1,56 @@
+from pydantic import BaseModel
+
+from . import Paged, PagedRequest, PagedResponse, ServiceBase
+
+
+class Token(BaseModel):
+    id: str
+    project_id: str
+    on_behalf_of: str
+
+
+class ExchangeToken:
+    class Request(BaseModel): ...
+
+    class Response(BaseModel):
+        token: str
+
+
+class IssueToken:
+    class Request(BaseModel): ...
+
+    class Response(BaseModel):
+        token: Token
+        token_secret: str
+
+
+class RevokeToken:
+    class Request(BaseModel):
+        token_id: str
+
+    class Response(BaseModel):
+        token: Token
+
+
+class ListTokens:
+    class Request(PagedRequest):
+        project_id: str
+        on_behalf_of: str | None = None
+
+    class Response(PagedResponse[Token]): ...
+
+
+class TokenService(ServiceBase):
+    def exchange(self) -> ExchangeToken.Response:
+        """Exchange a basic / identity token to a short-lived Spiral token."""
+        return self.client.post("/token/exchange", ExchangeToken.Request(), ExchangeToken.Response)
+
+    def issue(self) -> IssueToken.Response:
+        """Issue an API token on behalf of a principal."""
+        return self.client.post("/token/issue", IssueToken.Request(), IssueToken.Response)
+
+    def revoke(self, request: RevokeToken.Request) -> RevokeToken.Response:
+        return self.client.put("/token/revoke", request, RevokeToken.Response)
+
+    def list(self, request: ListTokens.Request) -> Paged[Token]:
+        return self.client.paged("/token/list", request, ListTokens.Response)

spiral/api/workloads.py

@@ -0,0 +1,45 @@
+from pydantic import BaseModel, Field
+
+from . import Paged, PagedRequest, PagedResponse, ProjectId, ServiceBase
+
+
+class Workload(BaseModel):
+    id: str
+    project_id: ProjectId
+    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 Response(BaseModel):
+        workload: Workload
+
+
+class IssueToken:
+    class Request(BaseModel):
+        workload_id: str
+
+    class Response(BaseModel):
+        token_id: str
+        token_secret: str
+
+
+class ListWorkloads:
+    class Request(PagedRequest):
+        project_id: str
+
+    class Response(PagedResponse[Workload]): ...
+
+
+class WorkloadService(ServiceBase):
+    def create(self, request: CreateWorkload.Request) -> CreateWorkload.Response:
+        return self.client.post("/workload/create", request, CreateWorkload.Response)
+
+    def issue_token(self, request: IssueToken.Request) -> IssueToken.Response:
+        return self.client.post("/workload/issue-token", request, IssueToken.Response)
+
+    def list(self, request: ListWorkloads.Request) -> Paged[Workload]:
+        return self.client.paged("/workload/list", request, ListWorkloads.Response)

spiral/arrow.py

@@ -0,0 +1,209 @@
+from collections import defaultdict
+from collections.abc import Callable, Iterable
+from functools import reduce
+from typing import TypeVar
+
+import numpy as np
+import pyarrow as pa
+from pyarrow import compute as pc
+
+T = TypeVar("T")
+
+
+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 dict_to_table(data) -> pa.Table:
+    return pa.Table.from_struct_array(dict_to_struct_array(data))
+
+
+def dict_to_struct_array(data, propagate_nulls: bool = False) -> pa.StructArray:
+    """Convert a nested dictionary of arrays to a table with nested structs."""
+    if isinstance(data, pa.ChunkedArray):
+        return data.combine_chunks()
+    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(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")

spiral/authn/authn.py

@@ -0,0 +1,89 @@
+import base64
+import logging
+import os
+
+from spiral.api import Authn, SpiralAPI
+
+ENV_TOKEN_ID = "SPIRAL_TOKEN_ID"
+ENV_TOKEN_SECRET = "SPIRAL_TOKEN_SECRET"
+
+log = logging.getLogger(__name__)
+
+
+class FallbackAuthn(Authn):
+    """Credential provider that tries multiple providers in order."""
+
+    def __init__(self, providers: list[Authn]):
+        self._providers = providers
+
+    def token(self) -> str | None:
+        for provider in self._providers:
+            token = provider.token()
+            if token is not None:
+                return token
+        return None
+
+
+class TokenAuthn(Authn):
+    """Credential provider that returns a fixed token."""
+
+    def __init__(self, token: str):
+        self._token = token
+
+    def token(self) -> str:
+        return self._token
+
+
+class EnvironmentAuthn(Authn):
+    """Credential provider that returns a basic token from the environment.
+
+    NOTE: Returns basic token. Must be exchanged.
+    """
+
+    def token(self) -> str | None:
+        if ENV_TOKEN_ID not in os.environ:
+            return None
+        if ENV_TOKEN_SECRET not in os.environ:
+            raise ValueError(f"{ENV_TOKEN_SECRET} is missing.")
+
+        token_id = os.environ[ENV_TOKEN_ID]
+        token_secret = os.environ[ENV_TOKEN_SECRET]
+        basic_token = base64.b64encode(f"{token_id}:{token_secret}".encode()).decode("utf-8")
+
+        return basic_token
+
+
+class DeviceAuthProvider(Authn):
+    """Auth provider that uses the device flow to authenticate a Spiral user."""
+
+    def __init__(self, device_auth):
+        # NOTE(ngates): device_auth: spiral.auth.device_code.DeviceAuth
+        #  We don't type it to satisfy our import linter
+        self._device_auth = device_auth
+
+    def token(self) -> str | None:
+        # TODO(ngates): only run this if we're in a notebook, CLI, or otherwise on the user's machine.
+        return self._device_auth.authenticate().access_token
+
+
+class TokenExchangeProvider(Authn):
+    """Auth provider that exchanges a basic token for a Spiral token."""
+
+    def __init__(self, authn: Authn, base_url: str):
+        self._authn = authn
+        self._token_service = SpiralAPI(authn, base_url).token
+
+        self._sp_token = None
+
+    def token(self) -> str | None:
+        if self._sp_token is not None:
+            return self._sp_token
+
+        # Don't try to exchange if token is not discovered.
+        if self._authn.token() is None:
+            return None
+
+        log.debug("Exchanging token")
+        self._sp_token = self._token_service.exchange().token
+
+        return self._sp_token

spiral/authn/device.py

@@ -0,0 +1,206 @@
+import logging
+import sys
+import textwrap
+import time
+import webbrowser
+from pathlib import Path
+
+import httpx
+import jwt
+from pydantic import BaseModel
+
+log = logging.getLogger(__name__)
+
+
+class TokensModel(BaseModel):
+    access_token: str
+    refresh_token: str
+
+    @property
+    def organization_id(self) -> str | None:
+        return self.unverified_access_token().get("org_id")
+
+    def unverified_access_token(self):
+        return jwt.decode(self.access_token, options={"verify_signature": False})
+
+
+class AuthModel(BaseModel):
+    tokens: TokensModel | None = None
+
+
+class DeviceAuth:
+    def __init__(
+        self,
+        auth_file: Path,
+        domain: str,
+        client_id: str,
+        http: httpx.Client = None,
+    ):
+        self._auth_file = auth_file
+        self._domain = domain
+        self._client_id = client_id
+        self._http = http or httpx.Client()
+
+        if self._auth_file.exists():
+            with self._auth_file.open("r") as f:
+                self._auth = AuthModel.model_validate_json(f.read())
+        else:
+            self._auth = AuthModel()
+
+        self._default_scope = ["email", "profile"]
+
+    def is_authenticated(self) -> bool:
+        """Check if the user is authenticated."""
+        tokens = self._auth.tokens
+        if tokens is None:
+            return False
+
+        # Give ourselves a 30-second buffer before the token expires.
+        return tokens.unverified_access_token()["exp"] - 30 > time.time()
+
+    def authenticate(self, force: bool = False, refresh: bool = False, organization_id: str = None) -> TokensModel:
+        """Blocking call to authenticate the user.
+
+        Triggers a device code flow and polls for the user to login.
+        """
+        if force:
+            return self._device_code(organization_id)
+
+        if refresh:
+            if self._auth.tokens is None:
+                raise ValueError("No tokens to refresh.")
+            tokens = self._refresh(self._auth.tokens, organization_id)
+            if not tokens:
+                raise ValueError("Failed to refresh token.")
+            return tokens
+
+        # Check for mis-matched organization.
+        if organization_id is not None:
+            tokens = self._auth.tokens
+            if tokens is not None and tokens.unverified_access_token().get("org_id") != organization_id:
+                tokens = self._refresh(self._auth.tokens, organization_id)
+                if tokens is None:
+                    return self._device_code(organization_id)
+
+        if self.is_authenticated():
+            return self._auth.tokens
+
+        # Try to refresh.
+        tokens = self._auth.tokens
+        if tokens is not None:
+            tokens = self._refresh(tokens)
+            if tokens is not None:
+                return tokens
+
+        # Otherwise, we kick off the device code flow.
+        return self._device_code(organization_id)
+
+    def logout(self):
+        self._remove_tokens()
+
+    def _device_code(self, organization_id: str | None):
+        scope = " ".join(self._default_scope)
+        res = self._http.post(
+            f"{self._domain}/auth/device/code",
+            data={
+                "client_id": self._client_id,
+                "scope": scope,
+                "organization_id": organization_id,
+            },
+        )
+        res = res.raise_for_status().json()
+        device_code = res["device_code"]
+        user_code = res["user_code"]
+        expires_at = res["expires_in"] + time.time()
+        interval = res["interval"]
+        verification_uri_complete = res["verification_uri_complete"]
+
+        # We need to detect if the user is running in a terminal, in Jupyter, etc.
+        # For now, we'll try to open the browser.
+        sys.stderr.write(
+            textwrap.dedent(
+                f"""
+                Please login here: {verification_uri_complete}
+                Your code is {user_code}.
+            """
+            )
+        )
+
+        # Try to open the browser (this also works if the Jupiter notebook is running on the user's machine).
+        opened = webbrowser.open(verification_uri_complete)
+
+        # If we have a server-side Jupyter notebook, we can try to open with client-side JavaScript.
+        if not opened and _in_notebook():
+            from IPython.display import Javascript, display
+
+            display(Javascript(f'window.open("{verification_uri_complete}");'))
+
+        # In the meantime, we need to poll for the user to login.
+        while True:
+            if time.time() > expires_at:
+                raise TimeoutError("Login timed out.")
+            time.sleep(interval)
+            res = self._http.post(
+                f"{self._domain}/auth/token",
+                data={
+                    "client_id": self._client_id,
+                    "device_code": device_code,
+                    "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+                },
+            )
+            if not res.is_success:
+                continue
+
+            tokens = TokensModel(
+                access_token=res.json()["access_token"],
+                refresh_token=res.json()["refresh_token"],
+            )
+            self._save_tokens(tokens)
+            return self._auth.tokens
+
+    def _refresh(self, tokens: TokensModel, organization_id: str = None) -> TokensModel | None:
+        """Attempt to use the refresh token."""
+        log.debug("Refreshing token %s", self._client_id)
+
+        res = self._http.post(
+            f"{self._domain}/auth/refresh",
+            data={
+                "client_id": self._client_id,
+                "grant_type": "refresh_token",
+                "refresh_token": tokens.refresh_token,
+                "organization_id": organization_id,
+            },
+        )
+        if not res.is_success:
+            print("Failed to refresh token", res.status_code, res.text)
+            return None
+
+        tokens = TokensModel(
+            access_token=res.json()["access_token"],
+            refresh_token=res.json()["refresh_token"],
+        )
+        self._save_tokens(tokens)
+        return tokens
+
+    def _save_tokens(self, tokens: TokensModel):
+        self._auth = self._auth.model_copy(update={"tokens": tokens})
+        self._auth_file.parent.mkdir(parents=True, exist_ok=True)
+        with self._auth_file.open("w") as f:
+            f.write(self._auth.model_dump_json(exclude_defaults=True))
+
+    def _remove_tokens(self):
+        self._auth_file.unlink(missing_ok=True)
+        self._auth = self._auth.model_copy(update={"tokens": None})
+
+
+def _in_notebook():
+    try:
+        from IPython import get_ipython
+
+        if "IPKernelApp" not in get_ipython().config:  # pragma: no cover
+            return False
+    except ImportError:
+        return False
+    except AttributeError:
+        return False
+    return True