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

spiral/protogen/substrait/extensions/__init__.py

@@ -0,0 +1,115 @@
+# Generated by the protocol buffer compiler.  DO NOT EDIT!
+# sources: substrait/extensions/extensions.proto
+# plugin: python-betterproto
+# This file has been @generated
+
+from dataclasses import dataclass
+from typing import List
+
+import betterproto
+import betterproto.lib.google.protobuf as betterproto_lib_google_protobuf
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionUri(betterproto.Message):
+    extension_uri_anchor: int = betterproto.uint32_field(1)
+    """
+    A surrogate key used in the context of a single plan used to reference the
+     URI associated with an extension.
+    """
+
+    uri: str = betterproto.string_field(2)
+    """
+    The URI where this extension YAML can be retrieved. This is the "namespace"
+     of this extension.
+    """
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclaration(betterproto.Message):
+    """
+    Describes a mapping between a specific extension entity and the uri where
+     that extension can be found.
+    """
+
+    extension_type: "SimpleExtensionDeclarationExtensionType" = (
+        betterproto.message_field(1, group="mapping_type")
+    )
+    extension_type_variation: "SimpleExtensionDeclarationExtensionTypeVariation" = (
+        betterproto.message_field(2, group="mapping_type")
+    )
+    extension_function: "SimpleExtensionDeclarationExtensionFunction" = (
+        betterproto.message_field(3, group="mapping_type")
+    )
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionType(betterproto.Message):
+    """Describes a Type"""
+
+    extension_uri_reference: int = betterproto.uint32_field(1)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    type_anchor: int = betterproto.uint32_field(2)
+    """
+    A surrogate key used in the context of a single plan to reference a
+     specific extension type
+    """
+
+    name: str = betterproto.string_field(3)
+    """the name of the type in the defined extension YAML."""
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionTypeVariation(betterproto.Message):
+    extension_uri_reference: int = betterproto.uint32_field(1)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    type_variation_anchor: int = betterproto.uint32_field(2)
+    """
+    A surrogate key used in the context of a single plan to reference a
+     specific type variation
+    """
+
+    name: str = betterproto.string_field(3)
+    """the name of the type in the defined extension YAML."""
+
+
+@dataclass(eq=False, repr=False)
+class SimpleExtensionDeclarationExtensionFunction(betterproto.Message):
+    extension_uri_reference: int = betterproto.uint32_field(1)
+    """
+    references the extension_uri_anchor defined for a specific extension URI.
+    """
+
+    function_anchor: int = betterproto.uint32_field(2)
+    """
+    A surrogate key used in the context of a single plan to reference a
+     specific function
+    """
+
+    name: str = betterproto.string_field(3)
+    """A function signature compound name"""
+
+
+@dataclass(eq=False, repr=False)
+class AdvancedExtension(betterproto.Message):
+    """
+    A generic object that can be used to embed additional extension information
+     into the serialized substrait plan.
+    """
+
+    optimization: List[
+        "betterproto_lib_google_protobuf.Any"
+    ] = betterproto.message_field(1)
+    """
+    An optimization is helpful information that don't influence semantics. May
+     be ignored by a consumer.
+    """
+
+    enhancement: "betterproto_lib_google_protobuf.Any" = betterproto.message_field(2)
+    """An enhancement alter semantics. Cannot be ignored by a consumer."""

spiral/server.py

@@ -0,0 +1,17 @@
+import socket
+import time
+
+
+def wait_for_port(port: int, host: str = "localhost", timeout: float = 5.0):
+    """Wait until a port starts accepting TCP connections."""
+    start_time = time.time()
+    while True:
+        try:
+            with socket.create_connection((host, port), timeout=timeout):
+                break
+        except OSError as ex:
+            time.sleep(0.01)
+            if time.time() - start_time >= timeout:
+                raise TimeoutError(
+                    f"Waited too long for the port {port} on host {host} to start accepting connections."
+                ) from ex

spiral/settings.py

@@ -1,54 +1,35 @@
 import functools
-import hashlib
 import os
-from collections.abc import Callable, Mapping
 from pathlib import Path
-from typing import Any
+from typing import Annotated
 
-import jwt
 import typer
-from pydantic import BaseModel, Field
+from pydantic import Field, ValidatorFunctionWrapHandler, WrapValidator
 from pydantic_settings import (
     BaseSettings,
     InitSettingsSource,
     PydanticBaseSettingsSource,
     SettingsConfigDict,
-    TomlConfigSettingsSource,
 )
 
 from spiral.api import SpiralAPI
-from spiral.authn.authn import (
-    DeviceAuthProvider,
-    EnvironmentAuthn,
-    FallbackAuthn,
-    TokenAuthn,
-    TokenExchangeProvider,
-)
-from spiral.authn.device import DeviceAuth
-from spiral.authn.github_ import GitHubActionsProvider
-from spiral.authn.modal_ import ModalProvider
+from spiral.core.client import Authn, DeviceCodeAuth, Token
 
 DEV = "PYTEST_VERSION" in os.environ or bool(os.environ.get("SPIRAL_DEV", None))
-FILE_FORMAT = os.environ.get("SPIRAL_FILE_FORMAT", "parquet")
+CI = "GITHUB_ACTIONS" in os.environ
 
 APP_DIR = Path(typer.get_app_dir("pyspiral"))
 LOG_DIR = APP_DIR / "logs"
-CONFIG_FILE = APP_DIR / "config.toml"
-
 
-class AuthSettings(BaseSettings):
-    model_config = SettingsConfigDict(frozen=True)
 
-    domain: str = "https://device.spiraldb.com"
-    client_id: str = "client_01J1CRS967RFQY7JSE8XZ44ATR" if DEV else "client_01J1CRS9MGF103FKTKZDCNJVM3"
+def validate_token(v, handler: ValidatorFunctionWrapHandler):
+    if isinstance(v, str):
+        return Token(v)
+    else:
+        raise ValueError("Token value must be a string")
 
 
-class AccessToken(BaseModel):
-    token: str
-
-    @functools.cached_property
-    def payload(self) -> dict:
-        return jwt.decode(self.token, options={"verify_signature": False})
+TokenType = Annotated[Token, WrapValidator(validate_token)]
 
 
 class SpiralDBSettings(BaseSettings):
@@ -57,38 +38,35 @@ class SpiralDBSettings(BaseSettings):
     host: str = "localhost" if DEV else "api.spiraldb.com"
     port: int = 4279 if DEV else 443
     ssl: bool = not DEV
-    auth: AuthSettings = Field(default_factory=AuthSettings)
-    token: str | None = None
+    token: TokenType | None = None
 
     @property
     def uri(self) -> str:
         return f"{'https' if self.ssl else 'http'}://{self.host}:{self.port}"
 
-    @property
-    def uri_scandal(self) -> str:
-        # TODO(marko): Scandal will be a different service. For now, gRPC API is hosted on the SpiralDB service.
-        return f"{'grpc+tls' if self.ssl else 'grpc'}://{self.host}:{self.port}"
 
-    @property
-    def uri_iceberg(self) -> str:
-        return self.uri + "/iceberg"
+class SpfsSettings(BaseSettings):
+    model_config = SettingsConfigDict(frozen=True)
 
-    def device_auth(self) -> DeviceAuth:
-        auth_file = (
-            APP_DIR / hashlib.md5(f"{self.auth.domain}/{self.auth.client_id}".encode()).hexdigest() / "auth.json"
-        )
-        return DeviceAuth(auth_file=auth_file, domain=self.auth.domain, client_id=self.auth.client_id)
+    host: str = "localhost" if DEV else "spfs.spiraldb.dev"
+    port: int = 4295 if DEV else 443
+    ssl: bool = not DEV
+
+    @property
+    def uri(self) -> str:
+        return f"{'https' if self.ssl else 'http'}://{self.host}:{self.port}"
 
 
 class Settings(BaseSettings):
     model_config = SettingsConfigDict(
-        toml_file=CONFIG_FILE,
         env_nested_delimiter="__",
         env_prefix="SPIRAL__",
         frozen=True,
     )
 
     spiraldb: SpiralDBSettings = Field(default_factory=SpiralDBSettings)
+    spfs: SpfsSettings = Field(default_factory=SpfsSettings)
+    file_format: str = Field(default="vortex")
 
     @functools.cached_property
     def api(self) -> "SpiralAPI":
@@ -99,17 +77,12 @@ class Settings(BaseSettings):
     @functools.cached_property
     def authn(self):
         if self.spiraldb.token:
-            return TokenAuthn(self.spiraldb.token)
-
-        return FallbackAuthn(
-            [
-                GitHubActionsProvider(),
-                ModalProvider(),
-                # TODO(marko): Github and Modal should also be behind token exchanged.
-                TokenExchangeProvider(EnvironmentAuthn(), self.spiraldb.uri),
-                DeviceAuthProvider(self.spiraldb.device_auth()),
-            ]
-        )
+            return Authn.from_token(self.spiraldb.token)
+        return Authn.from_fallback()
+
+    @functools.cached_property
+    def device_code_auth(self) -> DeviceCodeAuth:
+        return DeviceCodeAuth.default()
 
     @classmethod
     def settings_customise_sources(
@@ -120,42 +93,7 @@ class Settings(BaseSettings):
         init_settings: InitSettingsSource,
         **kwargs,
     ) -> tuple[PydanticBaseSettingsSource, ...]:
-        return env_settings, dotenv_settings, TomlConfigSettingsSource(settings_cls), init_settings
-
-
-class _LazyDict(Mapping):
-    def __init__(self, lazy_values: dict[str, Any | Callable[[], Any]]):
-        self._lazy_values = lazy_values
-
-    def _dict(self):
-        return {k: v() if callable(v) else v for k, v in self._lazy_values.items()}
-
-    def __getitem__(self, key, /):
-        return self._lazy_values[key]()
-
-    def get(self, key, /, default=None):
-        return self._lazy_values.get(key, lambda: default)()
-
-    def items(self):
-        return self._dict().items()
-
-    def keys(self):
-        return self._lazy_values.keys()
-
-    def values(self):
-        return self._dict().values()
-
-    def __contains__(self, key, /):
-        return key in self._lazy_values
-
-    def __eq__(self, other, /):
-        return False
-
-    def __iter__(self):
-        return iter(self._dict())
-
-    def __len__(self):
-        return len(self._lazy_values)
+        return env_settings, dotenv_settings, init_settings
 
 
 @functools.cache

spiral/substrait_.py

@@ -3,7 +3,7 @@ import pyarrow as pa
 
 import spiral.expressions as se
 from spiral.expressions.base import Expr
-from spiral.proto.substrait import (
+from spiral.protogen.substrait import (
     Expression,
     ExpressionFieldReference,
     ExpressionLiteral,
@@ -17,7 +17,7 @@ from spiral.proto.substrait import (
     ExpressionScalarFunction,
     ExtendedExpression,
 )
-from spiral.proto.substrait.extensions import (
+from spiral.protogen.substrait.extensions import (
     SimpleExtensionDeclaration,
     SimpleExtensionDeclarationExtensionFunction,
     SimpleExtensionDeclarationExtensionType,
@@ -43,7 +43,7 @@ class SubstraitConverter:
     def convert(self, buffer: pa.Buffer) -> Expr:
         """Convert a Substrait Extended Expression into a Spiral expression."""
 
-        expr: ExtendedExpression = ExtendedExpression().parse(buffer.to_pybytes())
+        expr: ExtendedExpression = ExtendedExpression().parse(buffer)
         assert len(expr.referred_expr) == 1, "Only one expression is supported"
 
         # Parse the extension URIs from the plan.
@@ -255,12 +255,16 @@ class SubstraitConverter:
 
                 scope = se.getitem(scope, field_name)
                 scope_type = scope_type.field(ref.field).type
-                return self._expr_direct_reference(scope, scope_type, ref.child) if ref.child else scope
+                if ref.is_set("child"):
+                    return self._expr_direct_reference(scope, scope_type, ref.child)
+                return scope
             case "list_element", ref:
                 ref: ExpressionReferenceSegmentListElement
                 scope = se.getitem(scope, ref.offset)
                 scope_type = scope_type.field(ref.field).type
-                return self._expr_direct_reference(scope, scope_type, ref.child) if ref.child else scope
+                if ref.is_set("child"):
+                    return self._expr_direct_reference(scope, scope_type, ref.child)
+                return scope
             case "", ref:
                 # Because Proto... we hit this case when we recurse into a child node and it's actually "None".
                 return scope

spiral/tables/__init__.py

@@ -0,0 +1,12 @@
+from spiral import _lib
+from spiral.tables.client import Tables
+from spiral.tables.maintenance import Maintenance
+from spiral.tables.scan import Scan
+from spiral.tables.snapshot import Snapshot
+from spiral.tables.table import Table
+from spiral.tables.transaction import Transaction
+
+# Eagerly import the Spiral library
+assert _lib, "Spiral library"
+
+__all__ = ["Tables", "Table", "Snapshot", "Scan", "Transaction", "Maintenance"]

spiral/tables/client.py

@@ -0,0 +1,130 @@
+from datetime import datetime
+from typing import Any
+
+import pyarrow as pa
+
+from spiral.api import SpiralAPI
+from spiral.api.projects import TableResource
+from spiral.core.client import Spiral as CoreSpiral
+from spiral.core.table.spec import Schema
+from spiral.datetime_ import timestamp_micros
+from spiral.expressions import ExprLike
+from spiral.tables.scan import Scan
+from spiral.tables.table import Table
+from spiral.types_ import Uri
+
+
+class Tables:
+    """
+    Spiral Tables a powerful and flexible way for storing, analyzing,
+    and querying massive and/or multimodal datasets.
+
+    The data model will feel familiar to users of SQL- or DataFrame-style systems,
+    yet is designed to be more flexible, more powerful, and more useful in the context
+    of modern data processing. Tables are stored and queried directly from object storage.
+    """
+
+    def __init__(self, api: SpiralAPI, spiral: CoreSpiral, *, project_id: str | None = None):
+        self._api = api
+        self._spiral = spiral
+        self._project_id = project_id
+
+    def table(self, identifier: str) -> Table:
+        """Open a table with a `dataset.table` identifier, or `table` name using the `default` dataset."""
+        project_id, dataset, table = self._parse_identifier(identifier)
+        if project_id is None:
+            raise ValueError("Must provide a fully qualified table identifier.")
+
+        res = list(self._api.project.list_tables(project_id, dataset=dataset, table=table))
+        if len(res) == 0:
+            raise ValueError(f"Table not found: {project_id}.{dataset}.{table}")
+
+        res = res[0]
+        return Table(self, self._spiral.get_table(res.id), identifier=f"{res.project_id}.{res.dataset}.{res.table}")
+
+    def list_tables(self) -> list[TableResource]:
+        project_id = self._project_id
+        if project_id is None:
+            raise ValueError("Must provide a project ID to list tables.")
+        return list(self._api.project.list_tables(project_id))
+
+    def create_table(
+        self,
+        identifier: str,
+        *,
+        key_schema: pa.Schema | Any,
+        root_uri: Uri | None = None,
+        exist_ok: bool = False,
+    ) -> Table:
+        """Create a new table in the project.
+
+        Args:
+            identifier: The table identifier, in the form `project.dataset.table`, `dataset.table` or `table`.
+            key_schema: The schema of the table's keys.
+            root_uri: The root URI for the table.
+            exist_ok: If True, do not raise an error if the table already exists.
+        """
+        project_id, dataset, table = self._parse_identifier(identifier)
+        if project_id is None:
+            raise ValueError("Must provide a fully qualified table identifier.")
+
+        if not isinstance(key_schema, pa.Schema):
+            key_schema = pa.schema(key_schema)
+        key_schema = Schema.from_arrow(key_schema)
+
+        core_table = self._spiral.create_table(
+            project_id,
+            dataset=dataset,
+            table=table,
+            key_schema=key_schema,
+            root_uri=root_uri,
+            exist_ok=exist_ok,
+        )
+
+        return Table(self, core_table, identifier=f"{project_id}.{dataset}.{table}")
+
+    def _parse_identifier(self, identifier: str) -> tuple[str | None, str, str]:
+        parts = identifier.split(".")
+        if len(parts) == 1:
+            return self._project_id, "default", parts[0]
+        elif len(parts) == 2:
+            return self._project_id, parts[0], parts[1]
+        elif len(parts) == 3:
+            return parts[0], parts[1], parts[2]
+        else:
+            raise ValueError(f"Invalid table identifier: {identifier}")
+
+    def scan(
+        self,
+        *projections: ExprLike,
+        where: ExprLike | None = None,
+        asof: datetime | int | None = None,
+        exclude_keys: bool = False,
+    ) -> 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: only data written before the given timestamp will be returned, caveats around compaction.
+            exclude_keys: whether to exclude the key columns in the scan result, defaults to False.
+                Note that if a projection includes a key column, it will be included in the result.
+        """
+        from spiral import expressions as se
+
+        if isinstance(asof, datetime):
+            asof = timestamp_micros(asof)
+
+        # Combine all projections into a single struct.
+        projection = se.merge(*projections)
+        if where is not None:
+            where = se.lift(where)
+
+        return Scan(
+            self._spiral.open_table_scan(
+                projection.__expr__,
+                filter=where.__expr__ if where else None,
+                asof=asof,
+                exclude_keys=exclude_keys,
+            ),
+        )