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

spiral/cli/telemetry.py

@@ -0,0 +1,18 @@
+import pyperclip
+import rich
+
+from spiral.api.telemetry import IssueExportTokenResponse
+from spiral.cli import AsyncTyper, state
+
+app = AsyncTyper(short_help="Client-side telemetry.")
+
+
+@app.command(help="Issue new telemetry export token.")
+def export():
+    res: IssueExportTokenResponse = state.settings.api.telemetry.issue_export_token()
+
+    command = f"export SPIRAL_OTEL_TOKEN={res.token}"
+    pyperclip.copy(command)
+
+    rich.print("Export command copied to clipboard! Paste and run to set [green]SPIRAL_OTEL_TOKEN[/green].")
+    rich.print("[dim]Token is valid for 1h.[/dim]")

spiral/cli/types.py

@@ -6,11 +6,11 @@ import typer
 from questionary import Choice
 from typer import Argument
 
-from spiral.api import OrganizationId, ProjectId
+from spiral.api.types import OrgId, ProjectId
 from spiral.cli import state
 
 
-def _project_default():
+def ask_project(title="Select a project"):
     projects = list(state.settings.api.project.list())
 
     if not projects:
@@ -18,7 +18,7 @@ def _project_default():
         raise typer.Exit(1)
 
     return questionary.select(
-        "Select a project",
+        title,
         choices=[
             Choice(title=f"{project.id} - {project.name}" if project.name else project.id, value=project.id)
             for project in projects
@@ -26,11 +26,11 @@ def _project_default():
     ).ask()
 
 
-ProjectArg = Annotated[ProjectId, Argument(help="Project ID", show_default=False, default_factory=_project_default)]
+ProjectArg = Annotated[ProjectId, Argument(help="Project ID", show_default=False, default_factory=ask_project)]
 
 
 def _org_default():
-    memberships = list(state.settings.api.organization.list_user_memberships())
+    memberships = list(state.settings.api.organization.list_memberships())
 
     if not memberships:
         rich.print("[red]No organizations found[/red]")
@@ -40,14 +40,12 @@ def _org_default():
         "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,
+                title=f"{m.org.id} - {m.org.name}" if m.org.name else m.org.id,
+                value=m.org.id,
             )
             for m in memberships
         ],
     ).ask()
 
 
-OrganizationArg = Annotated[
-    OrganizationId, Argument(help="Organization ID", show_default=False, default_factory=_org_default)
-]
+OrganizationArg = Annotated[OrgId, Argument(help="Organization ID", show_default=False, default_factory=_org_default)]

spiral/cli/{workload.py → workloads.py}

@@ -6,8 +6,8 @@ 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.api.workloads import CreateWorkloadRequest, IssueWorkloadCredentialsResponse, Workload
+from spiral.cli import AsyncTyper, printer, state
 from spiral.cli.types import ProjectArg
 
 app = AsyncTyper()
@@ -16,9 +16,9 @@ app = AsyncTyper()
 @app.command(help="Create a new workload.")
 def create(
     project: ProjectArg,
-    name: Annotated[OptionalStr, Option(help="Friendly name for the workload.")] = None,
+    name: Annotated[str | None, Option(help="Friendly name for the workload.")] = None,
 ):
-    res = state.settings.api.workload.create(CreateWorkload.Request(project_id=project, name=name))
+    res = state.settings.api.workload.create(project, CreateWorkloadRequest(name=name))
     rich.print(f"Created workload {res.workload.id}")
 
 
@@ -26,13 +26,13 @@ def create(
 def ls(
     project: ProjectArg,
 ):
-    workloads = list(state.settings.api.workload.list(ListWorkloads.Request(project_id=project)))
+    workloads = list(state.settings.api.workload.list(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))
+@app.command(help="Issue new workflow credentials.")
+def issue_credentials(workload_id: Annotated[str, Argument(help="Workload ID.")]):
+    res: IssueWorkloadCredentialsResponse = state.settings.api.workload.issue_credentials(workload_id)
 
     while True:
         choice = questionary.select(
@@ -45,15 +45,15 @@ def token(workload_id: Annotated[str, Argument(help="Workload ID.")]):
         ).ask()
 
         if choice == 1:
-            pyperclip.copy(res.token_secret)
+            pyperclip.copy(res.client_secret)
             rich.print("[green]Secret copied to clipboard![/green]")
             break
         elif choice == 2:
-            rich.print(f"[green]Token Secret:[/green] {res.token_secret}")
+            rich.print(f"[green]Token Secret:[/green] {res.client_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}")
+    rich.print(f"[green]Token ID:[/green] {res.client_id}")

spiral/{catalog.py → client.py}

@@ -3,23 +3,22 @@ from typing import TYPE_CHECKING
 import jwt
 
 from spiral.api import SpiralAPI
-from spiral.api.projects import CreateProject
+from spiral.api.projects import CreateProjectRequest, CreateProjectResponse
+from spiral.core.client import Spiral as CoreSpiral
 from spiral.settings import Settings, settings
 
 if TYPE_CHECKING:
-    from pyiceberg import catalog
-
+    from spiral.iceberg import Iceberg
     from spiral.project import Project
-    from spiral.table import Table
-
-_default = object()
 
 
 class Spiral:
     def __init__(self, config: Settings | None = None):
         self._config = config or settings()
         self._api = self._config.api
-
+        self._core = CoreSpiral(
+            api_url=self._config.spiraldb.uri, spfs_url=self._config.spfs.uri, authn=self._config.authn
+        )
         self._org = None
 
     @property
@@ -33,7 +32,10 @@ class Spiral:
     @property
     def organization(self) -> str:
         if self._org is None:
-            token_payload = jwt.decode(self._config.authn.token(), options={"verify_signature": False})
+            token = self._config.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"]
@@ -45,49 +47,33 @@ class Spiral:
 
         return [Project(self, id=p.id, name=p.name) for p in self.api.project.list()]
 
-    def list_project_ids(self) -> list[str]:
-        """List project IDs."""
-        return [p.id for p in self.list_projects()]
-
     def create_project(
         self,
         id_prefix: str | None = None,
         *,
-        org: str | None = None,
         name: str | None = None,
     ) -> "Project":
         """Create a project in the current, or given, organization."""
         from .project import Project
 
-        org = org or self.organization
-        res = self.api.project.create(CreateProject.Request(organization_id=org, id_prefix=id_prefix, name=name))
+        res: CreateProjectResponse = self.api.project.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 .project import 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, id=project_id, name=project_id)
 
-    def table(self, identifier: str) -> "Table":
-        """Open a table with a "project.dataset.table" identifier."""
-        parts = identifier.split(".")
-        if len(parts) != 3:
-            raise ValueError(f"Invalid table identifier: {identifier}")
-        project_id, dataset, table = parts
-
-        return self.project(project_id).table(f"{dataset}.{table}")
-
-    def iceberg_catalog(self) -> "catalog.Catalog":
-        """Open the Iceberg catalog."""
-        from pyiceberg.catalog import load_catalog
-
-        return load_catalog(
-            "default",
-            **{
-                "type": "rest",
-                "uri": self._config.spiraldb.uri_iceberg,
-                "token": self._config.authn.token(),
-            },
-        )
+    @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/client/__init__.pyi

@@ -0,0 +1,117 @@
+from spiral.api.types import DatasetName, IndexName, OrgId, ProjectId, RootUri, TableName
+from spiral.core.index import SearchScan, TextIndex
+from spiral.core.table import Table, TableMaintenance, TableScan, TableSnapshot, TableTransaction
+from spiral.core.table.spec import Schema
+from spiral.expressions import Expr
+
+class Token:
+    def __init__(self, value: str): ...
+    def expose_secret(self) -> str: ...
+
+class Authn:
+    @staticmethod
+    def from_token(token: Token) -> Authn: ...
+    @staticmethod
+    def from_fallback() -> Authn: ...
+    @staticmethod
+    def from_device() -> 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."""
+        ...
+
+class Spiral:
+    """A client for Spiral database"""
+    def __init__(
+        self,
+        api_url: str | None = None,
+        spfs_url: str | None = None,
+        authn: Authn | None = None,
+    ):
+        """Initialize the Spiral client."""
+        ...
+    def authn(self) -> Authn:
+        """Get the current authentication context."""
+        ...
+    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 get_table(self, table_id: str) -> Table:
+        """Get and open table."""
+
+    def open_table(self, table_id: str, key_schema: Schema, root_uri: RootUri) -> Table:
+        """Open a table. This does not make any network calls."""
+        ...
+
+    def open_table_scan(
+        self,
+        projection: Expr,
+        filter: Expr | None = None,
+        asof: int | None = None,
+        exclude_keys: bool = False,
+    ) -> TableScan:
+        """Construct a table scan."""
+        ...
+
+    def open_transaction(self, table: Table, format: str | None = None) -> TableTransaction:
+        """Being transaction."""
+        ...
+
+    def open_maintenance(self, table: Table, format: str | None = None) -> TableMaintenance:
+        """Access maintenance operations for a table."""
+        ...
+    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 get_text_index(self, index_id: str) -> TextIndex:
+        """Get a text-based index."""
+        ...
+
+    def open_search_scan(
+        self,
+        rank_by: Expr,
+        top_k: int,
+        # NOTE(marko): Required for now.
+        freshness_window_s: int,
+        *,
+        filter: Expr | None = None,
+    ) -> SearchScan:
+        """Query an index."""
+        ...
+
+    def _sync_snapshot(self, index_id: str, snapshot: TableSnapshot) -> None:
+        """Synchronize an index with a table snapshot.
+
+        IMPORTANT: This is only exposed for testing purposes and should not be used.
+        """
+        ...

spiral/core/index/__init__.pyi

@@ -0,0 +1,15 @@
+import pyarrow as pa
+
+class IndexStatus:
+    status: str
+    staleness_s: int | None
+    # An extent of keys that are indexed.
+    # key_extent: KeyExtent | None
+
+class TextIndex:
+    id: str
+
+    def status(self) -> IndexStatus: ...
+
+class SearchScan:
+    def to_record_batches(self) -> pa.RecordBatchReader: ...

spiral/core/{core → table}/__init__.pyi

@@ -1,11 +1,36 @@
 from typing import Any, Literal
 
 import pyarrow as pa
-from spiral.core.manifests import FragmentManifest
-from spiral.core.metastore import PyMetastore
-from spiral.core.spec import ColumnGroup, ColumnGroupMetadata, FileFormat, KeyRange, Schema, WriteAheadLog
 from spiral.expressions import Expr
 
+from .manifests import FragmentManifest
+from .metastore import PyMetastore
+from .spec import ColumnGroup, Key, Schema, WriteAheadLog
+
+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 Table:
     def __init__(self, metastore: PyMetastore): ...
 
@@ -16,18 +41,16 @@ class Table:
 
     def get_wal(self, *, asof: int | None) -> WriteAheadLog: ...
     def get_schema(self, *, asof: int | None) -> Schema: ...
-    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_snapshot(self, *, asof: int | None) -> TableSnapshot: ...
+
+class TableSnapshot:
+    """A snapshot of a table at a specific point in time."""
+
+    asof: int
+    table: Table
+    wal: WriteAheadLog
 
 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: ...
@@ -39,6 +62,12 @@ class TableScan:
         key_table: pa.Table | pa.RecordBatch | None = None,
         batch_readahead: int | None = None,
     ) -> pa.RecordBatchReader: ...
+    def to_shuffled_record_batches(
+        self,
+        batch_readahead: int | None = None,
+        shuffle_buffer_size: int | None = None,
+        shuffle_pool_num_rows: int | 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]: ...
@@ -54,16 +83,14 @@ class ColumnGroupScan:
     def schema(self) -> Schema: ...
 
 class TableTransaction:
-    def __init__(self, metastore: PyMetastore, format: FileFormat): ...
-    @property
-    def status(self) -> str: ...
+    status: str
+
     def write(self, expr: Expr, *, partition_size_bytes: int | None = None): ...
     def commit(self): ...
     def abort(self): ...
     def metrics(self) -> dict[str, Any]: ...
 
 class TableMaintenance:
-    def __init__(self, metastore: PyMetastore, format: FileFormat): ...
     def flush_wal(self): ...
     def compact_key_space(
         self,

spiral/core/{manifests → table/manifests}/__init__.pyi

@@ -1,5 +1,6 @@
 import pyarrow as pa
-from spiral.core.spec import FileFormat, FragmentLevel, KeyExtent, KeyMap, KeyRange, KeySpan
+from spiral.core.table import KeyRange
+from spiral.core.table.spec import FileFormat, FragmentLevel, KeyExtent, KeySpan
 from spiral.types_ import Timestamp
 
 class FragmentManifest:
@@ -22,32 +23,13 @@ class FragmentFile:
     format: FileFormat
     format_metadata: bytes | None
     size_bytes: int
-    # NOTE: Empty for keyspace file.
     column_ids: list[str]
-    fs_id: str
-    fs_level: FragmentLevel
+    level: FragmentLevel
+    # NOTE: Empty for key space files.
+    column_ids: list[str]
     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/table/metastore/__init__.pyi

@@ -0,0 +1,62 @@
+"""The SpiralDB metastore API."""
+
+from spiral.core.client import Authn
+from spiral.core.table.spec import ColumnGroup, ColumnGroupMetadata, FileFormat, KeySpaceMetadata, Schema, WriteAheadLog
+from spiral.types_ import Uri
+
+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 get_key_space_metadata(self) -> KeySpaceMetadata:
+        """Return the metadata for the key space."""
+        ...
+    def get_column_group_metadata(self, column_group: ColumnGroup) -> ColumnGroupMetadata:
+        """Return the metadata for a column group."""
+        ...
+
+    @staticmethod
+    def http(
+        table_id: str,
+        root_uri: str,
+        key_schema: Schema,
+        base_url: str,
+        authn: Authn,
+    ) -> 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."""