pyspiral 0.1.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/cli/state.py

@@ -0,0 +1,3 @@
+from spiral.settings import Settings
+
+settings: Settings = Settings()

spiral/cli/table.py

@@ -0,0 +1,20 @@
+from typing import Annotated
+
+import rich
+from typer import Option
+
+from spiral.api.tables import ListTables, Table
+from spiral.cli import AsyncTyper, OptionalStr, printer, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@app.command(help="List tables.")
+def ls(
+    project: ProjectArg,
+    dataset: Annotated[OptionalStr, Option(help="Filter by dataset name.")] = None,
+):
+    """List tables."""
+    tables = list(state.settings.api.table.list(ListTables.Request(project_id=project, dataset=dataset)))
+    rich.print(printer.table_of_models(Table, tables, fields=["id", "project_id", "dataset", "table"]))

spiral/cli/token.py

@@ -0,0 +1,27 @@
+from typing import Annotated
+
+import rich
+from typer import Argument, Option
+
+from spiral.api.tokens import ListTokens, RevokeToken, Token
+from spiral.cli import AsyncTyper, OptionalStr, printer, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@app.command(help="List tokens.")
+def ls(
+    project: ProjectArg,
+    on_behalf_of: Annotated[OptionalStr, Option(help="Filter by on behalf of.")] = None,
+):
+    tokens = list(state.settings.api.token.list(ListTokens.Request(project_id=project, on_behalf_of=on_behalf_of)))
+    rich.print(printer.table_of_models(Token, tokens, fields=["id", "project_id", "on_behalf_of"]))
+
+
+@app.command(help="Revoke a token.")
+def revoke(token_id: Annotated[str, Argument(help="Token ID.")]):
+    res = state.settings.api.token.revoke(RevokeToken.Request(token_id=token_id))
+    rich.print(
+        f"Revoked token {res.token.id} for project {res.token.project_id} acting on behalf of {res.token.on_behalf_of}"
+    )

spiral/cli/types.py

@@ -0,0 +1,53 @@
+from typing import Annotated
+
+import questionary
+import rich
+import typer
+from questionary import Choice
+from typer import Argument
+
+from spiral.api import OrganizationId, ProjectId
+from spiral.cli import state
+
+
+def _project_default():
+    projects = list(state.settings.api.project.list())
+
+    if not projects:
+        rich.print("[red]No projects found[/red]")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        "Select a project",
+        choices=[
+            Choice(title=f"{project.id} - {project.name}" if project.name else project.id, value=project.id)
+            for project in projects
+        ],
+    ).ask()
+
+
+ProjectArg = Annotated[ProjectId, Argument(help="Project ID", show_default=False, default_factory=_project_default)]
+
+
+def _org_default():
+    memberships = list(state.settings.api.organization.list_user_memberships())
+
+    if not memberships:
+        rich.print("[red]No organizations found[/red]")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        "Select an organization",
+        choices=[
+            Choice(
+                title=f"{m.organization.id} - {m.organization.name}" if m.organization.name else m.organization.id,
+                value=m.organization.id,
+            )
+            for m in memberships
+        ],
+    ).ask()
+
+
+OrganizationArg = Annotated[
+    OrganizationId, Argument(help="Organization ID", show_default=False, default_factory=_org_default)
+]

spiral/cli/workload.py

@@ -0,0 +1,59 @@
+from typing import Annotated
+
+import pyperclip
+import questionary
+import rich
+from questionary import Choice
+from typer import Argument, Option
+
+from spiral.api.workloads import CreateWorkload, IssueToken, ListWorkloads, Workload
+from spiral.cli import AsyncTyper, OptionalStr, printer, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@app.command(help="Create a new workload.")
+def create(
+    project: ProjectArg,
+    name: Annotated[OptionalStr, Option(help="Friendly name for the workload.")] = None,
+):
+    res = state.settings.api.workload.create(CreateWorkload.Request(project_id=project, name=name))
+    rich.print(f"Created workload {res.workload.id}")
+
+
+@app.command(help="List workloads.")
+def ls(
+    project: ProjectArg,
+):
+    workloads = list(state.settings.api.workload.list(ListWorkloads.Request(project_id=project)))
+    rich.print(printer.table_of_models(Workload, workloads, fields=["id", "project_id", "name"]))
+
+
+@app.command(help="Issue a token.")
+def token(workload_id: Annotated[str, Argument(help="Workload ID.")]):
+    res = state.settings.api.workload.issue_token(IssueToken.Request(workload_id=workload_id))
+
+    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.token_secret)
+            rich.print("[green]Secret copied to clipboard![/green]")
+            break
+        elif choice == 2:
+            rich.print(f"[green]Token Secret:[/green] {res.token_secret}")
+            break
+        elif choice == 3:
+            break
+        else:
+            rich.print("[red]Invalid choice. Please try again.[/red]")
+
+    rich.print(f"[green]Token ID:[/green] {res.token_id}")

spiral/config.py

@@ -0,0 +1,26 @@
+import os
+
+from pydantic_settings import BaseSettings, SettingsConfigDict
+
+FILE_FORMAT = os.environ.get("SPIRAL_FILE_FORMAT", "parquet")
+
+
+class Config(BaseSettings):
+    model_config = SettingsConfigDict(
+        env_nested_delimiter="__",
+        env_prefix="SPIRAL_CORE__",
+        frozen=True,
+    )
+
+    partition_file_min_size: int = 256 * 1024 * 1024  # 256MB
+    flush_wal_on_write: bool = False
+
+    # TODO(marko): Support config. Unused after migration to Rust.
+    # #: Defaults to ThreadPoolExecutor's default (based on os.cpu_count().
+    # scan_num_threads: int | None = 61  # 61 is used by Golang and Tokio, for some reason...
+    #
+    # #: The duration of WAL that is preserved to allow for txn conflict resolution.
+    # transaction_window: int = 0 if DEV else timedelta(days=1).total_seconds()
+    #
+    # #: Truncation length of string statistics.
+    # string_truncation_length: int = 1024

spiral/core/core/__init__.pyi

@@ -0,0 +1,53 @@
+from typing import Any
+
+import pyarrow as pa
+from spiral.core.manifests import FragmentManifest
+from spiral.core.metastore import PyMetastore
+from spiral.core.spec import ColumnGroup, ColumnGroupMetadata, KeyRange, Schema, WriteAheadLog
+from spiral.expressions import Expr
+
+class Table:
+    def __init__(self, metastore: PyMetastore): ...
+
+    id: str
+    root_uri: str
+    key_schema: Schema
+    metastore: PyMetastore
+
+    def get_wal(self, *, asof: int | None) -> WriteAheadLog: ...
+    def get_column_group_metadata(self, column_group: ColumnGroup, *, asof: int | None) -> ColumnGroupMetadata: ...
+    def list_column_groups(self, *, asof: int | None) -> list[ColumnGroup] | None: ...
+    def get_schema(self, *, asof: int | None) -> Schema: ...
+
+class TableScan:
+    def __init__(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        exclude_keys: bool = False,
+        aux_schema: pa.Schema | None = None,
+    ) -> TableScan: ...
+    def key_schema(self) -> Schema: ...
+    def schema(self) -> Schema: ...
+    def is_empty(self) -> bool: ...
+    def split(self) -> list[KeyRange]: ...
+    def table_ids(self) -> list[str]: ...
+    def column_groups(self) -> list[ColumnGroup]: ...
+    def to_record_batches(self, aux_table: pa.Table | pa.RecordBatch | None = None) -> pa.RecordBatchReader: ...
+    def column_group_scan(self, column_group: ColumnGroup) -> ColumnGroupScan: ...
+    def key_space_scan(self, table_id: str) -> KeySpaceScan: ...
+    def metrics(self) -> dict[str, Any]: ...
+
+class KeySpaceScan:
+    manifest: FragmentManifest
+
+    def key_schema(self) -> Schema: ...
+
+class ColumnGroupScan:
+    manifest: FragmentManifest
+
+    def schema(self) -> Schema: ...
+
+def write(table: Table, expr: Expr, format: str = "parquet", *, partition_size: int | None = None): ...
+def flush_wal(table: Table, manifest_format: str = "parquet"): ...

spiral/core/manifests/__init__.pyi

@@ -0,0 +1,53 @@
+import pyarrow as pa
+from spiral.core.spec import FileFormat, FragmentLevel, KeyExtent, KeyMap, KeyRange, KeySpan
+from spiral.types_ import Timestamp
+
+class FragmentManifest:
+    def __len__(self): ...
+    def __getitem__(self, idx: int): ...
+    def to_arrow(self) -> pa.RecordBatchReader: ...
+    @staticmethod
+    def compute_schema(data_schema: pa.Schema) -> pa.Schema: ...
+    @staticmethod
+    def from_fragment(fragment_file: FragmentFile) -> FragmentManifest: ...
+    @staticmethod
+    def from_arrow(reader: pa.RecordBatchReader) -> FragmentManifest: ...
+    @staticmethod
+    def empty() -> FragmentManifest: ...
+
+class FragmentFile:
+    id: str
+    committed_at: Timestamp | None
+    compacted_at: Timestamp | None
+    format: FileFormat
+    format_metadata: bytes | None
+    size_bytes: int
+    # NOTE: Empty for keyspace file.
+    column_ids: list[str]
+    fs_id: str
+    fs_level: FragmentLevel
+    ks_id: str
+    key_span: KeySpan
+    key_extent: KeyExtent
+    key_map: KeyMap | None
+
+    def __init__(
+        self,
+        *,
+        id: str,
+        committed_at: Timestamp | None,
+        compacted_at: Timestamp | None,
+        format: FileFormat,
+        format_metadata: bytes,
+        size_bytes: int,
+        column_ids: list[str],
+        fs_id: str,
+        fs_level: FragmentLevel,
+        ks_id: str,
+        key_span: KeySpan,
+        key_extent: KeyExtent,
+        key_map: KeyMap | None,
+        stats: pa.StructArray,
+    ): ...
+    @property
+    def key_range(self) -> KeyRange: ...

spiral/core/metastore/__init__.pyi

@@ -0,0 +1,91 @@
+"""The SpiralDB metastore API."""
+
+from collections.abc import Callable
+
+from spiral.core.spec import ColumnGroup, ColumnGroupMetadata, FileFormat, LogEntry, Schema, WriteAheadLog
+from spiral.types_ import Timestamp, Uri
+from spiraldb.proto.spiral.table import ManifestHandle
+
+class FileHandle:
+    def __init__(self, *, uri: str, format: FileFormat, spfs_token: str | None): ...
+
+    uri: str
+    format: FileFormat
+    spfs_token: str | None
+
+class FileRef:
+    def __init__(self, *, id: str, file_type: FileType, file_format: FileFormat): ...
+
+    id: str
+    file_type: FileType
+    file_format: FileFormat
+
+    def resolve_uri(self, root_uri: str) -> str:
+        """Resolves the file reference URI given the root URI."""
+
+class FileType:
+    FragmentFile: FileType
+    FragmentManifest: FileType
+    ReferenceFile: FileType
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+
+class PyMetastore:
+    """Rust implementation of the metastore API."""
+
+    @property
+    def table_id(self) -> str: ...
+    @property
+    def root_uri(self) -> Uri: ...
+    @property
+    def key_schema(self) -> Schema: ...
+    def get_wal(self) -> WriteAheadLog:
+        """Return the log for the table."""
+        ...
+
+    def append_wal(self, prev_last_modified_at: Timestamp, entries: list[LogEntry]) -> WriteAheadLog:
+        """Append additional entries into the write-ahead log given the previous write-ahead log timestamp.
+
+        The given entries should have a timestamp of zero and will be assigned an actual timestamp by the server.
+
+        This API is designed to support both a trivial compare-and-swap on the WAL, and also to support more advanced
+        conflict resolution within the metastore.
+        """
+        ...
+
+    def update_wal(
+        self,
+        prev_ks_manifest_handle_id: str,
+        truncate_ts_max: Timestamp | None = None,
+        new_ks_manifest_handle: ManifestHandle | None = None,
+    ) -> WriteAheadLog:
+        """Update the write-ahead log atomically.
+
+        Supports WAL truncation and manifest handle updates necessary for flushing.
+        """
+        ...
+
+    def get_column_group_metadata(self, column_group: ColumnGroup) -> ColumnGroupMetadata:
+        """Return the metadata for column group."""
+        ...
+
+    def update_column_group_metadata(
+        self, prev_last_modified_at: Timestamp, column_group_metadata: ColumnGroupMetadata
+    ) -> ColumnGroupMetadata:
+        """Update the column group metadata to the metastore given the previous metadata timestamp."""
+        ...
+
+    def list_column_groups(self) -> tuple[list[ColumnGroup], Timestamp]:
+        """List all column groups in the table, or None if no index is available."""
+        ...
+
+    @staticmethod
+    def http(
+        table_id: str, root_uri: str, key_schema: Schema, base_url: str, token_provider: Callable[[], str]
+    ) -> PyMetastore:
+        """Construct a PyMetastore backed by an HTTP metastore service."""
+
+    @staticmethod
+    def test(table_id: str, root_uri: str, key_schema: Schema) -> PyMetastore:
+        """Construct a PyMetastore backed by an in-memory mock metastore service."""

spiral/core/spec/__init__.pyi

@@ -0,0 +1,257 @@
+"""Type definitions for the spiral.core.spec module shipped as part of the native library."""
+
+import pyarrow as pa
+
+class ColumnGroup:
+    def __init__(self, path: list[str]): ...
+    @property
+    def table_id(self) -> str: ...
+    @property
+    def path(self) -> list[str]: ...
+    def identifier(self, salt: int) -> str:
+        """Return the column group identifier based on the given salt."""
+
+    @staticmethod
+    def from_str(path: str) -> ColumnGroup: ...
+
+class ColumnGroupMetadata:
+    def __init__(
+        self,
+        *,
+        column_group: ColumnGroup,
+        manifest_handle: ManifestHandle | None,
+        last_modified_at: int,
+        schema_versions: list[VersionedSchema] | None,
+        immutable_schema: bool,
+        schema_salt: int,
+    ): ...
+
+    column_group: ColumnGroup
+    manifest_handle: ManifestHandle | None
+    last_modified_at: int
+    schema_versions: list[VersionedSchema]
+    immutable_schema: bool
+    schema_salt: int
+
+    def latest_schema(self) -> VersionedSchema:
+        """Returns the latest schema of the column group."""
+        ...
+
+    def asof(self, asof: int) -> ColumnGroupMetadata:
+        """Returns the metadata as of a given timestamp. Currently just filtering versioned schemas."""
+        ...
+
+    def apply_wal(self, wal: WriteAheadLog) -> ColumnGroupMetadata:
+        """Applies the given WAL to the metadata."""
+
+    def __bytes__(self):
+        """Serializes the ColumnGroupMetadata to a protobuf buffer."""
+
+    @staticmethod
+    def from_proto(buffer: bytes) -> ColumnGroupMetadata:
+        """Deserializes a ColumnGroupMetadata from a protobuf buffer."""
+        ...
+
+class LogEntry:
+    ts: int
+    operation: KeySpaceWriteOp | FragmentSetWriteOp | ConfigurationOp | SchemaEvolutionOp | SchemaBreakOp
+
+    def column_group(self) -> ColumnGroup | None:
+        """Returns the column group of the entry if it is associated with one."""
+
+    def replace_timestamp(self, ts: int) -> LogEntry:
+        """Returns a copy of the entry with the timestamp replaced."""
+
+    @staticmethod
+    def schema_break(*, column_group: ColumnGroup, removed_column_names: list[str]) -> LogEntry: ...
+    @staticmethod
+    def schema_evolution(*, column_group: ColumnGroup, new_schema: Schema) -> LogEntry: ...
+    @staticmethod
+    def ks_write(
+        *,
+        ks_id: str,
+        manifest_handle: ManifestHandle,
+    ) -> LogEntry: ...
+    @staticmethod
+    def fs_write(
+        *,
+        column_group: ColumnGroup,
+        fs_id: str,
+        fs_level: FragmentLevel,
+        manifest_handle: ManifestHandle,
+        key_span: KeySpan,
+        key_extent: KeyExtent,
+        column_ids: list[str],
+    ) -> LogEntry: ...
+
+class FileFormat:
+    def __init__(self, value: int): ...
+
+    Parquet: FileFormat
+    Protobuf: FileFormat
+    BinaryArray: FileFormat
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+        ...
+
+class FragmentLevel:
+    L0: FragmentLevel
+    L1: FragmentLevel
+
+    def __int__(self) -> int:
+        """Returns the protobuf enum int value."""
+        ...
+
+class Key:
+    def __init__(self, key: bytes): ...
+
+    key: bytes
+
+    def __add__(self, other: Key) -> Key:
+        """Concatenates two keys.
+
+        TODO(ngates): remove this function. It should not be necessary to concatenate keys."""
+
+    def __bytes__(self): ...
+    def step(self) -> Key:
+        """Returns the next key in the key space."""
+
+    @staticmethod
+    def min() -> Key: ...
+    @staticmethod
+    def max() -> Key: ...
+    @staticmethod
+    def from_array_tuple(array_tuple: tuple[pa.Array]) -> Key: ...
+
+class KeyExtent:
+    """An inclusive range of keys."""
+
+    def __init__(self, *, min: Key, max: Key): ...
+
+    min: Key
+    max: Key
+
+    def to_range(self) -> KeyRange:
+        """Turn this inclusive key extent into an exclusive key range."""
+
+    def union(self, key_extent: KeyExtent) -> KeyExtent: ...
+    def __or__(self, other: KeyExtent) -> KeyExtent: ...
+    def intersection(self, key_extent: KeyExtent) -> KeyExtent | None: ...
+    def __and__(self, other: KeyExtent) -> KeyExtent | None: ...
+    def contains(self, item: Key) -> bool: ...
+    def __contains__(self, item: Key) -> bool: ...
+
+class KeyRange:
+    """A right-exclusive range of keys."""
+
+    def __init__(self, *, begin: Key, end: Key): ...
+
+    begin: Key
+    end: Key
+
+    def union(self, other: KeyRange) -> KeyRange: ...
+    def __or__(self, other: KeyRange) -> KeyRange: ...
+    def intersection(self, key_extent: KeyRange) -> KeyRange | None: ...
+    def __and__(self, other: KeyRange) -> KeyRange | None: ...
+    def contains(self, item: Key) -> bool: ...
+    def __contains__(self, item: Key) -> bool: ...
+    def is_disjoint(self, key_range: KeyRange) -> bool:
+        return self.end <= key_range.begin or self.begin >= key_range.end
+
+    @staticmethod
+    def beginning_with(begin: Key) -> KeyRange: ...
+    @staticmethod
+    def ending_with(end: Key) -> KeyRange: ...
+    @staticmethod
+    def full() -> KeyRange: ...
+
+class KeySpan:
+    """An exclusive range of keys as indexed by their position in a key space."""
+
+    def __init__(self, *, begin: int, end: int): ...
+
+    begin: int
+    end: int
+
+    def __len__(self) -> int: ...
+    def shift(self, offset: int) -> KeySpan: ...
+    def union(self, other: KeySpan) -> KeySpan: ...
+    def __or__(self, other: KeySpan) -> KeySpan: ...
+
+class KeyMap:
+    """Displacement map."""
+
+class ManifestHandle:
+    def __init__(self, id: str, format: FileFormat, file_size: int, spfs_format_metadata: bytes | None): ...
+
+    id: str
+    format: FileFormat
+    file_size: int
+    spfs_format_metadata: bytes | None
+
+class Schema:
+    def to_arrow(self) -> pa.Schema:
+        """Returns the Arrow schema."""
+        ...
+
+    @staticmethod
+    def from_arrow(arrow: pa.Schema) -> Schema:
+        """Creates a Schema from an Arrow schema."""
+        ...
+
+class VersionedSchema:
+    ts: int
+    schema: Schema
+    column_ids: list[str]
+
+class KeySpaceWriteOp:
+    ks_id: str
+    manifest_handle: ManifestHandle
+
+class FragmentSetWriteOp:
+    column_group: ColumnGroup
+    fs_id: str
+    fs_level: FragmentLevel
+    manifest_handle: ManifestHandle
+    key_span: KeySpan
+    key_extent: KeyExtent
+    column_ids: list[str]
+
+class ConfigurationOp:
+    column_group: ColumnGroup
+
+class SchemaEvolutionOp:
+    column_group: ColumnGroup
+
+class SchemaBreakOp:
+    column_group: ColumnGroup
+
+class WriteAheadLog:
+    def __init__(
+        self,
+        *,
+        entries: list[LogEntry] | None = None,
+        truncated_up_to: int = 0,
+    ): ...
+
+    entries: list[LogEntry]
+    truncated_up_to: int
+    ks_manifest_handle: ManifestHandle | None
+
+    @property
+    def last_modified_at(self) -> int:
+        """Returns the timestamp of the last modification of the log."""
+
+    def filter(
+        self, asof: int | None = None, since: int | None = None, column_group: ColumnGroup | None = None
+    ) -> WriteAheadLog:
+        """Filters the WAL to entries by the given parameters."""
+
+    def __bytes__(self):
+        """Serializes the ColumnGroupMetadata to a protobuf buffer."""
+
+    @staticmethod
+    def from_proto(buffer: bytes) -> WriteAheadLog:
+        """Deserializes a WriteAheadLog from a protobuf buffer."""
+        ...