pyspiral 0.8.9__cp311-abi3-macosx_11_0_arm64.whl

spiral/cli/workloads.py

@@ -0,0 +1,86 @@
+from typing import Annotated
+
+import pyperclip
+import questionary
+from questionary import Choice
+from typer import Argument, Option
+
+from spiral.api.workloads import (
+    CreateWorkloadRequest,
+    CreateWorkloadResponse,
+    IssueWorkloadCredentialsResponse,
+    Workload,
+)
+from spiral.cli import CONSOLE, ERR_CONSOLE, AsyncTyper, printer, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@app.command(help="List workloads.")
+def ls(
+    project: ProjectArg,
+):
+    workloads = list(state.spiral.api.workloads.list(project))
+    CONSOLE.print(printer.table_of_models(Workload, workloads, fields=["id", "project_id", "name"]))
+
+
+@app.command(help="Create a new workload.")
+def create(
+    project: ProjectArg,
+    name: Annotated[str | None, Option(help="Friendly name for the workload.")] = None,
+):
+    res: CreateWorkloadResponse = state.spiral.api.workloads.create(project, CreateWorkloadRequest(name=name))
+    CONSOLE.print(f"{res.workload.id}")
+
+
+@app.command(help="Deactivate a workload. Removes all associated credentials.")
+def deactivate(
+    workload_id: Annotated[str, Argument(help="Workload ID.")],
+):
+    state.spiral.api.workloads.deactivate(workload_id)
+    CONSOLE.print(f"Deactivated workload {workload_id}")
+
+
+@app.command(help="Issue new workflow credentials.")
+def issue_creds(
+    workload_id: Annotated[str, Argument(help="Workload ID.")],
+    skip_prompt: Annotated[bool, Option(help="Skip prompt and print secret to console.")] = False,
+):
+    res: IssueWorkloadCredentialsResponse = state.spiral.api.workloads.issue_credentials(workload_id)
+
+    if skip_prompt:
+        CONSOLE.print(f"[green]SPIRAL_CLIENT_ID[/green] {res.client_id}")
+        CONSOLE.print(f"[green]SPIRAL_CLIENT_SECRET[/green] {res.client_secret}")
+    else:
+        while True:
+            choice = questionary.select(
+                "What would you like to do with the secret? You will not be able to see this secret again!",
+                choices=[
+                    Choice(title="Copy to clipboard", value=1),
+                    Choice(title="Print to console", value=2),
+                    Choice(title="Exit", value=3),
+                ],
+            ).ask()
+
+            if choice == 1:
+                pyperclip.copy(res.client_secret)
+                CONSOLE.print("[green]Secret copied to clipboard![/green]")
+                break
+            elif choice == 2:
+                CONSOLE.print(f"[green]SPIRAL_CLIENT_SECRET[/green] {res.client_secret}")
+                break
+            elif choice == 3:
+                break
+            else:
+                ERR_CONSOLE.print("Invalid choice. Please try again.")
+
+        CONSOLE.print(f"[green]SPIRAL_CLIENT_ID[/green] {res.client_id}")
+
+
+@app.command(help="Revoke workflow credentials.")
+def revoke_creds(
+    client_id: Annotated[str, Argument(help="Client ID to revoke.")],
+):
+    state.spiral.api.workloads.revoke_credentials(client_id)
+    CONSOLE.print(f"Revoked credentials for client ID {client_id}")

spiral/client.py

@@ -0,0 +1,279 @@
+import os
+from datetime import datetime, timedelta
+from typing import TYPE_CHECKING
+
+import jwt
+import pyarrow as pa
+
+from spiral.api import SpiralAPI
+from spiral.api.projects import CreateProjectRequest, CreateProjectResponse
+from spiral.core.authn import Authn
+from spiral.core.client import Internal, Shard
+from spiral.core.client import Spiral as CoreSpiral
+from spiral.core.config import ClientSettings
+from spiral.datetime_ import timestamp_micros
+from spiral.expressions import ExprLike
+from spiral.scan import Scan
+
+if TYPE_CHECKING:
+    from spiral.iceberg import Iceberg
+    from spiral.key_space_index import KeySpaceIndex
+    from spiral.project import Project
+    from spiral.table import Table
+    from spiral.text_index import TextIndex
+
+
+class Spiral:
+    """Main client for interacting with the Spiral data platform.
+
+    Configuration is loaded with the following priority (highest to lowest):
+    1. Explicit parameters.
+    2. Environment variables (`SPIRAL__*`)
+    3. Config file (`~/.spiral.toml`)
+    4. Default values (production URLs)
+
+    Examples:
+
+    ```python
+    import spiral
+    # Default configuration
+    sp = spiral.Spiral()
+
+    # With config overrides
+    sp = spiral.Spiral(overrides={"limits.concurrency": "16"})
+    ```
+
+    Args:
+        config: Custom ClientSettings object. Defaults to global settings.
+        overrides: Configuration overrides using dot notation,
+            see the [Client Configuration](https://docs.spiraldb.com/config) page for a full list.
+    """
+
+    def __init__(
+        self,
+        *,
+        config: ClientSettings | None = None,
+        overrides: dict[str, str] | None = None,
+    ):
+        if bool(os.environ.get("SPIRAL__DEV", None)):
+            overrides = overrides or {}
+            overrides["server.url"] = "http://localhost:4279"
+            overrides["spfs.url"] = "http://localhost:4295"
+
+        self._overrides = overrides
+        self._config = config
+        self._org = None
+        self._core = None
+        self._api = None
+
+    @property
+    def config(self) -> ClientSettings:
+        """Returns the client's configuration"""
+        return self.core.config()
+
+    @property
+    def authn(self) -> Authn:
+        """Get the authentication handler for this client."""
+        return self.core.authn()
+
+    @property
+    def api(self) -> SpiralAPI:
+        if self._api is None:
+            self._api = SpiralAPI(self.config.server_url, self.authn)
+        return self._api
+
+    @property
+    def core(self) -> CoreSpiral:
+        if self._core is None:
+            self._core = CoreSpiral(
+                config=self._config,
+                overrides=self._overrides,
+            )
+
+        return self._core
+
+    @property
+    def internal(self) -> Internal:
+        return self.core.internal(format=self.config.file_format)
+
+    @property
+    def organization(self) -> str:
+        if self._org is None:
+            token = self.authn.token()
+            if token is None:
+                raise ValueError("Authentication failed.")
+            token_payload = jwt.decode(token.expose_secret(), options={"verify_signature": False})
+            if "org_id" not in token_payload:
+                raise ValueError("Please create an organization.")
+            self._org = token_payload["org_id"]
+        return self._org
+
+    def list_projects(self) -> list["Project"]:
+        """List project IDs."""
+        from .project import Project
+
+        return [Project(self, project_id=p.id, name=p.name) for p in self.api.projects.list()]
+
+    def create_project(
+        self,
+        id_prefix: str | None = None,
+        *,
+        name: str | None = None,
+    ) -> "Project":
+        """Create a project in the current, or given, organization."""
+        from .project import Project
+
+        res: CreateProjectResponse = self.api.projects.create(CreateProjectRequest(id_prefix=id_prefix, name=name))
+        return Project(self, res.project.id, name=res.project.name)
+
+    def project(self, project_id: str) -> "Project":
+        """Open an existing project."""
+        from spiral.project import Project
+
+        # We avoid an API call since we'd just be fetching a human-readable name. Seems a waste in most cases.
+        return Project(self, project_id=project_id, name=project_id)
+
+    def table(self, table_id: str) -> "Table":
+        """Open a table using an ID."""
+        from spiral.table import Table
+
+        return Table(self, self.core.table(table_id))
+
+    def text_index(self, index_id: str) -> "TextIndex":
+        """Open a text index using an ID."""
+        from spiral.text_index import TextIndex
+
+        return TextIndex(self.core.text_index(index_id))
+
+    def key_space_index(self, index_id: str) -> "KeySpaceIndex":
+        """Open a key space index using an ID."""
+        from spiral.key_space_index import KeySpaceIndex
+
+        return KeySpaceIndex(self.core.key_space_index(index_id))
+
+    def scan(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        shard: Shard | None = None,
+    ) -> Scan:
+        """Starts a read transaction on the Spiral.
+
+        Args:
+            projections: a set of expressions that return struct arrays.
+            where: a query expression to apply to the data.
+            asof: execute the scan on the version of the table as of the given timestamp.
+            shard: if provided, opens the scan only for the given shard.
+                While shards can be provided when executing the scan, providing a shard here
+                optimizes the scan planning phase and can significantly reduce metadata download.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return Scan(
+            self,
+            self.core.scan(projection.__expr__, filter=where.__expr__ if where else None, asof=asof, shard=shard),
+        )
+
+    # TODO(marko): This should be query, and search should be query + scan.
+    def search(
+        self,
+        top_k: int,
+        *rank_by: ExprLike,
+        filters: ExprLike | None = None,
+        freshness_window: timedelta | None = None,
+    ) -> pa.RecordBatchReader:
+        """Queries the index with the given rank by and filters clauses. Returns a stream of scored keys.
+
+        Args:
+            top_k: The number of top results to return.
+            rank_by: Rank by expressions are combined for scoring.
+                See `se.text.find` and `se.text.boost` for scoring expressions.
+            filters: The `filters` expression is used to filter the results.
+                It must return a boolean value and use only conjunctions (ANDs). Expressions in filters
+                statement are considered either a `must` or `must_not` clause in search terminology.
+            freshness_window: If provided, the index will not be refreshed if its freshness does not exceed this window.
+        """
+        from spiral import expressions as se
+
+        if not rank_by:
+            raise ValueError("At least one rank by expression is required.")
+        rank_by = se.or_(*rank_by)
+        if filters is not None:
+            filters = se.lift(filters)
+
+        if freshness_window is None:
+            freshness_window = timedelta(seconds=0)
+        freshness_window_s = int(freshness_window.total_seconds())
+
+        return self.core.search(
+            top_k=top_k,
+            rank_by=rank_by.__expr__,
+            filters=filters.__expr__ if filters else None,
+            freshness_window_s=freshness_window_s,
+        )
+
+    def resume_scan(self, state_json: str) -> Scan:
+        """Resumes a previously started scan using its scan state.
+
+        Args:
+            state_json: The scan state returned by a previous scan.
+        """
+        from spiral.core.table import ScanState
+
+        state = ScanState.from_json(state_json)
+        return Scan(self, self.core.load_scan(state))
+
+    def compute_shards(
+        self,
+        max_batch_size: int,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        stream: bool = False,
+    ) -> list[Shard]:
+        """Computes shards over the given projections and filter.
+
+        Args:
+            max_batch_size: The maximum number of rows per shard.
+            projections: a set of expressions that return struct arrays.
+            where: a query expression to apply to the data.
+            asof: execute the scan on the version of the table as of the given timestamp.
+            stream: if true, builds shards in a streaming fashion, suitable for very large tables.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        if not projections:
+            raise ValueError("At least one projection is required.")
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return self.core.compute_shards(
+            max_batch_size, projection.__expr__, where.__expr__ if where else None, asof=asof, stream=stream
+        )
+
+    @property
+    def iceberg(self) -> "Iceberg":
+        """
+        Apache Iceberg is a powerful open-source table format designed for high-performance data lakes.
+        Iceberg brings reliability, scalability, and advanced features like time travel, schema evolution,
+        and ACID transactions to your warehouse.
+        """
+        from spiral.iceberg import Iceberg
+
+        return Iceberg(self)

spiral/core/_tools/__init__.pyi

@@ -0,0 +1,5 @@
+from ..table.spec import Schema
+
+def pretty_key(key: bytes, schema: Schema) -> str:
+    """Represent a key in a human-readable way."""
+    ...

spiral/core/authn/__init__.pyi

@@ -0,0 +1,21 @@
+from spiral.api.types import OrgId
+
+class Token:
+    def __init__(self, value: str): ...
+    def expose_secret(self) -> str: ...
+
+class Authn:
+    def token(self) -> Token | None: ...
+
+class DeviceCodeAuth:
+    @staticmethod
+    def default() -> DeviceCodeAuth:
+        """Return the static device code instance."""
+        ...
+    def authenticate(self, force: bool = False, org_id: OrgId | None = None) -> Token:
+        """Authenticate using device code flow."""
+        ...
+
+    def logout(self) -> None:
+        """Logout from the device authentication session."""
+        ...

spiral/core/client/__init__.pyi

@@ -0,0 +1,270 @@
+from enum import Enum
+from typing import Any
+
+import pyarrow as pa
+from spiral.api.types import DatasetName, IndexName, ProjectId, RootUri, TableId, TableName
+from spiral.core.authn import Authn
+from spiral.core.config import ClientSettings
+from spiral.core.table import ColumnGroupState, KeyRange, KeySpaceState, Scan, ScanState, Snapshot, Table, Transaction
+from spiral.core.table.spec import ColumnGroup, Schema
+from spiral.expressions import Expr
+
+class Spiral:
+    """A client for Spiral database"""
+    def __init__(
+        self,
+        config: ClientSettings | None = None,
+        overrides: dict[str, str] | None = None,
+    ):
+        """Initialize the Spiral client.
+
+        Args:
+            config: Client configuration, defaults to the global config.
+            overrides: Configuration overrides using dot notation,
+                see the [Client Configuration](/python-client) page for a full list.
+        """
+        ...
+
+    def authn(self) -> Authn:
+        """Get the current authentication context."""
+        ...
+
+    def scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        shard: Shard | None = None,
+        key_columns: KeyColumns | None = None,
+    ) -> Scan:
+        """Construct a table scan."""
+        ...
+
+    def load_scan(self, plan_state: ScanState) -> Scan:
+        """Load a scan from a serialized scan state."""
+        ...
+
+    def transaction(self, table: Table, *, partition_max_bytes: int | None = None) -> Transaction:
+        """Being a table transaction."""
+        ...
+
+    def search(
+        self,
+        top_k: int,
+        rank_by: Expr,
+        *,
+        filters: Expr | None = None,
+        freshness_window_s: int | None = None,
+    ) -> pa.RecordBatchReader:
+        """Search an index.
+
+        Searching an index returns a stream of record batches that match table's key schema + float score column.
+        """
+        ...
+
+    def table(self, table_id: str) -> Table:
+        """Get a table."""
+        ...
+
+    def create_table(
+        self,
+        project_id: ProjectId,
+        dataset: DatasetName,
+        table: TableName,
+        key_schema: Schema,
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> Table:
+        """Create a new table in the specified project."""
+        ...
+
+    def move_table(
+        self,
+        table_id: TableId,
+        new_dataset: DatasetName,
+    ):
+        """Move a table to a dataset in the same project."""
+        ...
+
+    def rename_table(
+        self,
+        table_id: TableId,
+        new_table: TableName,
+    ):
+        """Rename a table."""
+        ...
+
+    def drop_table(self, table_id: TableId):
+        """Drop a table."""
+        ...
+
+    def text_index(self, index_id: str) -> TextIndex:
+        """Get a text index."""
+        ...
+
+    def create_text_index(
+        self,
+        project_id: ProjectId,
+        name: IndexName,
+        projection: Expr,
+        filter: Expr | None = None,
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> TextIndex:
+        """Create a new index in the specified project."""
+        ...
+
+    def key_space_index(self, index_id: str) -> KeySpaceIndex:
+        """Get a key space index."""
+        ...
+
+    def create_key_space_index(
+        self,
+        project_id: ProjectId,
+        name: IndexName,
+        granularity: int,
+        projection: Expr,
+        filter: Expr | None = None,
+        *,
+        root_uri: RootUri | None = None,
+        exist_ok: bool = False,
+    ) -> KeySpaceIndex:
+        """Create a new key space index in the specified project."""
+        ...
+
+    def compute_shards(
+        self,
+        max_batch_size: int,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        stream: bool = False,
+    ) -> list[Shard]:
+        """Constructs shards for a given projection (and filter).
+
+        Useful for distributing work.
+        """
+        ...
+
+    def internal(self, *, format: str | None = None) -> Internal:
+        """Internal client APIs. It can change without notice."""
+        ...
+
+    def config(self) -> ClientSettings:
+        """Client-side configuration."""
+        ...
+
+class KeyColumns(Enum):
+    IfProjected = 0
+    Included = 1
+    Only = 2
+
+class TextIndex:
+    id: str
+
+class KeySpaceIndex:
+    id: str
+    table_id: str
+    granularity: int
+    projection: Expr
+    filter: Expr
+    asof: int
+
+class Shard:
+    """A shard representing a partition of data.
+
+    Attributes:
+        key_range: The key range for this shard.
+        cardinality: The number of rows in this shard, if known.
+    """
+
+    key_range: KeyRange
+    cardinality: int | None
+
+    def __init__(self, key_range: KeyRange, cardinality: int | None): ...
+    def __getnewargs__(self) -> tuple[KeyRange, int | None]: ...
+    def union(self, other: KeyRange) -> KeyRange: ...
+    def __or__(self, other):
+        """Combine two shards into one that covers both key ranges.
+
+        The cardinality of the resulting shard is set to None.
+        """
+        ...
+
+class ShuffleConfig:
+    """Configuration for within-shard sample shuffling.
+
+    This controls how samples are shuffled within a buffer, separate from
+    which shards to read (which is specified as a parameter to the scan).
+
+    Attributes:
+        buffer_size: Size of the buffer pool for shuffling samples.
+        seed: Random seed for reproducibility. If None, uses OS randomness.
+    """
+
+    buffer_size: int
+    seed: int | None
+
+    def __init__(
+        self,
+        buffer_size: int,
+        *,
+        seed: int | None = None,
+    ): ...
+
+class Internal:
+    def flush_wal(self, table: Table) -> None:
+        """
+        Flush the write-ahead log of the table.
+        """
+        ...
+    def update_text_index(self, index: TextIndex, snapshot: Snapshot) -> None:
+        """
+        Index table changes up to the given snapshot.
+        """
+        ...
+    def update_key_space_index(self, index: KeySpaceIndex, snapshot: Snapshot) -> None:
+        """
+        Index table changes up to the given snapshot.
+        """
+        ...
+    def key_space_state(self, snapshot: Snapshot) -> KeySpaceState:
+        """
+        The key space state for the table.
+        """
+        ...
+    def column_group_state(
+        self, snapshot: Snapshot, key_space_state: KeySpaceState, column_group: ColumnGroup
+    ) -> ColumnGroupState:
+        """
+        The state the column group of the table.
+        """
+        ...
+    def column_groups_states(self, snapshot: Snapshot, key_space_state: KeySpaceState) -> list[ColumnGroupState]:
+        """
+        The state of each column group of the table.
+        """
+        ...
+    def key_space_index_shards(self, index: KeySpaceIndex) -> list[Shard]:
+        """
+        Compute the scan shards from a key space index.
+        """
+        ...
+    def prepare_shard(
+        self,
+        output_path: str,
+        scan: Scan,
+        shard: Shard,
+        row_block_size: int = 8192,
+    ) -> None:
+        """
+        Prepare a shard locally. Used for `SpiralStream` integration with `streaming` which requires on-disk shards.
+        """
+        ...
+    def metrics(self) -> dict[str, Any]: ...
+
+def flush_telemetry() -> None:
+    """Flush telemetry data to the configured exporter."""
+    ...

spiral/core/config/__init__.pyi

@@ -0,0 +1,35 @@
+class ClientSettings:
+    """Client configuration loaded from ~/.spiral.toml and environment variables."""
+
+    @staticmethod
+    def load() -> ClientSettings:
+        """Load ClientSettings from ~/.spiral.toml and environment variables.
+
+        Configuration priority (highest to lowest):
+        1. Environment variables (SPIRAL__*)
+        2. Config file (~/.spiral.toml)
+        3. Default values
+        """
+        ...
+
+    @property
+    def server_url(self) -> str:
+        """The Spiral API endpoint URL."""
+        ...
+
+    @property
+    def spfs_url(self) -> str:
+        """The SpFS endpoint URL."""
+        ...
+
+    @property
+    def file_format(self) -> str:
+        """File format for table storage (vortex or parquet)."""
+        ...
+
+    def to_json(self) -> str:
+        """Serialize to a JSON string"""
+        ...
+    @staticmethod
+    def from_json(json: str) -> ClientSettings:
+        """Deserialize from a JSON-formatted string"""