pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/arrow_.py

@@ -0,0 +1,221 @@
+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 struct_array(fields: list[tuple[str, bool, pa.Array]], /, mask: list[bool] | None = None) -> pa.StructArray:
+    return pa.StructArray.from_arrays(
+        arrays=[x[2] for x in fields],
+        fields=[pa.field(x[0], type=x[2].type, nullable=x[1]) for x in fields],
+        mask=pa.array(mask) if mask else mask,
+    )
+
+
+def table(fields: list[tuple[str, bool, pa.Array]], /) -> pa.Table:
+    return pa.Table.from_struct_array(struct_array(fields))
+
+
+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/cli/__init__.py

@@ -0,0 +1,79 @@
+import asyncio
+import functools
+import inspect
+from typing import IO
+
+import rich
+import typer
+from click import ClickException
+from grpclib import GRPCError
+from httpx import HTTPStatusError
+
+
+class AsyncTyper(typer.Typer):
+    """Wrapper to allow async functions to be used as commands.
+
+    We also pre-bake some configuration.
+
+    Per https://github.com/tiangolo/typer/issues/88#issuecomment-1732469681
+    """
+
+    def __init__(self, **kwargs):
+        super().__init__(
+            no_args_is_help=True,
+            pretty_exceptions_enable=False,
+            **kwargs,
+        )
+
+    def callback(self, *args, **kwargs):
+        decorator = super().callback(*args, **kwargs)
+        for wrapper in (_wrap_exceptions, _maybe_run_async):
+            decorator = functools.partial(wrapper, decorator)
+        return decorator
+
+    def command(self, *args, **kwargs):
+        decorator = super().command(*args, **kwargs)
+        for wrapper in (_wrap_exceptions, _maybe_run_async):
+            decorator = functools.partial(wrapper, decorator)
+        return decorator
+
+
+class _ClickGRPCException(ClickException):
+    def __init__(self, err: GRPCError):
+        super().__init__(err.message or "GRPCError message was None.")
+        self.err = err
+        self.exit_code = 1
+
+    def format_message(self) -> str:
+        if self.err.details:
+            return f"{self.message}: {self.err.details}"
+        return self.message
+
+    def show(self, file: IO[str] | None = None) -> None:
+        rich.print(f"Error: {self.format_message()}", file=file)
+
+
+def _maybe_run_async(decorator, f):
+    if inspect.iscoroutinefunction(f):
+
+        @functools.wraps(f)
+        def runner(*args, **kwargs):
+            return asyncio.run(f(*args, **kwargs))
+
+        decorator(runner)
+    else:
+        decorator(f)
+    return f
+
+
+def _wrap_exceptions(decorator, f):
+    @functools.wraps(f)
+    def runner(*args, **kwargs):
+        try:
+            return f(*args, **kwargs)
+        except HTTPStatusError as e:
+            raise ClickException(str(e))
+        except GRPCError as e:
+            raise _ClickGRPCException(e)
+
+    return decorator(runner)

spiral/cli/__main__.py

@@ -0,0 +1,4 @@
+from spiral.cli.app import main
+
+if __name__ == "__main__":
+    main()

spiral/cli/admin.py

@@ -0,0 +1,16 @@
+from rich import print
+
+from spiral.api.types import OrgId
+from spiral.cli import AsyncTyper, state
+
+app = AsyncTyper()
+
+
+@app.command()
+def sync(
+    org_id: OrgId | None = None,
+):
+    state.settings.api._admin.sync_orgs()
+
+    for membership in state.settings.api._admin.sync_memberships(org_id):
+        print(membership)

spiral/cli/app.py

@@ -0,0 +1,65 @@
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+from spiral.cli import (
+    AsyncTyper,
+    admin,
+    console,
+    fs,
+    iceberg,
+    indexes,
+    login,
+    orgs,
+    projects,
+    state,
+    tables,
+    telemetry,
+    workloads,
+)
+from spiral.settings import LOG_DIR, Settings
+
+app = AsyncTyper(name="spiral")
+
+
+@app.callback()
+def _callback(verbose: bool = False):
+    if verbose:
+        logging.getLogger().setLevel(level=logging.INFO)
+
+    # Load the settings (we reload in the callback to support testing under different env vars)
+    state.settings = Settings()
+
+
+app.add_typer(fs.app, name="fs")
+app.add_typer(orgs.app, name="orgs")
+app.add_typer(projects.app, name="projects")
+app.add_typer(iceberg.app, name="iceberg")
+app.add_typer(tables.app, name="tables")
+app.add_typer(indexes.app, name="indexes")
+app.add_typer(telemetry.app, name="telemetry")
+app.command("console")(console.command)
+app.command("login")(login.command)
+app.command("whoami")(login.whoami)
+
+# Register unless we're building docs. Because Typer docs command does not skip hidden commands...
+if not bool(os.environ.get("SPIRAL_DOCS", False)):
+    app.add_typer(workloads.app, name="workloads", hidden=True)
+    app.add_typer(admin.app, name="admin", hidden=True)
+    app.command("logout", hidden=True)(login.logout)
+
+
+def main():
+    # Setup rotating CLI logging.
+    # NOTE(ngates): we should do the same for the Spiral client? Maybe move this logic elsewhere?
+    LOG_DIR.mkdir(parents=True, exist_ok=True)
+    logging.basicConfig(
+        level=logging.DEBUG,
+        handlers=[RotatingFileHandler(LOG_DIR / "cli.log", maxBytes=2**20, backupCount=10)],
+    )
+
+    app()
+
+
+if __name__ == "__main__":
+    main()

spiral/cli/console.py

@@ -0,0 +1,95 @@
+import os
+import subprocess
+
+from spiral import Spiral
+from spiral.adbc import ADBCFlightServer, SpiralADBCServer
+from spiral.server import wait_for_port
+
+
+def command():
+    """Launch a SQL console to query Spiral tables."""
+
+    # To avoid taking a dependency on Harlequin, we install it on-demand using
+    # either uvx or pipx.
+    harlequin_args = _uvx()
+    if harlequin_args is None:
+        harlequin_args = _pipx()
+    if harlequin_args is None:
+        raise ValueError("Please install pipx to continue\n\tSee https://github.com/pypa/pipx")
+
+    # Set up a pipe to send the server port to the child process.
+    r, w = os.pipe()
+
+    pid = os.fork()
+    if pid == 0:  # In the child
+        os.close(w)
+        port = int.from_bytes(os.read(r, 4), "big")
+
+        # Wait for the server to be up.
+        wait_for_port(port)
+
+        os.execv(
+            harlequin_args[0],
+            harlequin_args
+            + [
+                "-a",
+                "adbc",
+                "--driver-type",
+                "flightsql",
+                f"grpc://localhost:{port}",
+            ],
+        )
+    else:
+        os.close(r)
+
+        # I can't get the Flight server to stop writing to stdout. So we need to spawn a new process I think and
+        # then hope we can kill it?
+        fd = os.open("/dev/null", os.O_WRONLY)
+        os.dup2(fd, 1)
+        os.dup2(fd, 2)
+
+        # In the parent, we launch the Flight SQL server and send the port to the child
+        server = ADBCFlightServer(SpiralADBCServer(Spiral()))
+        os.write(w, server.port.to_bytes(4, "big"))
+
+        # Then wait for the console app to exit
+        os.waitpid(pid, 0)
+
+
+def _pipx() -> list[str] | None:
+    """Run harlequin via pipx."""
+    res = subprocess.run(["which", "pipx"], stdout=subprocess.PIPE)
+    if res.returncode != 0:
+        return None
+        # raise ValueError("Please install pipx to continue\n\tSee https://github.com/pypa/pipx")
+    pipx = res.stdout.strip()
+
+    return [
+        pipx,
+        "run",
+        "--pip-args",
+        "adbc_driver_flightsql",
+        "--pip-args",
+        # for now, we pin rich
+        "rich<=13.9.1",
+        "harlequin[adbc]",
+    ]
+
+
+def _uvx() -> list[str] | None:
+    """Run harlequin via uvx."""
+    res = subprocess.run(["which", "uvx"], stdout=subprocess.PIPE)
+    if res.returncode != 0:
+        return None
+    uvx = res.stdout.strip()
+
+    return [
+        uvx,
+        "--with",
+        "adbc_driver_flightsql",
+        "--with",
+        "rich<=13.9.1",
+        "--from",
+        "harlequin[adbc]",
+        "harlequin",
+    ]

spiral/cli/fs.py

@@ -0,0 +1,112 @@
+from typing import Annotated
+
+import questionary
+import rich
+from pydantic import SecretStr
+from typer import Option
+
+from spiral.api.filesystems import (
+    AWSSecretAccessKey,
+    BuiltinFileSystem,
+    GCPServiceAccount,
+    UpdateGCSFileSystem,
+    UpdateS3FileSystem,
+    UpstreamFileSystem,
+)
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import ProjectArg, ask_project
+
+app = AsyncTyper(short_help="File Systems.")
+
+
+@app.command(help="Show the file system configured for project.")
+def show(project: ProjectArg):
+    file_system = state.settings.api.file_system.get_file_system(project)
+    match file_system:
+        case BuiltinFileSystem(provider=provider):
+            rich.print(f"provider: {provider}")
+        case _:
+            rich.print(file_system)
+
+
+def ask_provider():
+    res = state.settings.api.file_system.list_providers()
+    return questionary.select("Select a file system provider", choices=res).ask()
+
+
+BuiltinProviderOpt = Annotated[
+    str,
+    Option(help="Built-in provider to use for the file system.", show_default=False, default_factory=ask_provider),
+]
+
+
+@app.command(help="Update a project's default file system.")
+def update(
+    project: ProjectArg,
+    builtin: bool = Option(False, help="Use a built-in file system provider."),
+    upstream: bool = Option(
+        False, help="Use another project as default file system. Only if another project is an external provider."
+    ),
+    s3: bool = Option(False, help="Use S3 compatible provider."),
+    gcs: bool = Option(False, help="Use GCS provider."),
+    provider: str = Option(None, help="Built-in provider to use for the file system."),
+    endpoint: str = Option(None, help="Endpoint for S3 provider."),
+    region: str = Option(None, help="Region for S3 or GCS provider. Required for GCS."),
+    bucket: str = Option(None, help="Bucket name for S3 or GCS provider."),
+    directory: str = Option(None, help="Directory for S3 or GCS provider."),
+    access_key_id: str = Option(None, help="Access key ID for S3 provider. Required for S3."),
+    secret_access_key: str = Option(None, help="Secret access key for S3 provider. Required for S3."),
+    credentials_path: str = Option(
+        None, help="Path to service account credentials file for GCS provider. Required for GCS."
+    ),
+):
+    if not any([builtin, s3, gcs, upstream]):
+        raise ValueError("Must specify one of --builtin, --upstream, --s3, or --gcs.")
+
+    if builtin:
+        provider = provider or ask_provider()
+        file_system = BuiltinFileSystem(provider=provider)
+
+    elif upstream:
+        upstream_project = ask_project(title="Select a project to use as file system.")
+        file_system = UpstreamFileSystem(project_id=upstream_project)
+
+    elif s3:
+        if access_key_id is None or secret_access_key is None:
+            raise ValueError("--access-key-id and --secret-access-key are required for S3 provider.")
+        credentials = AWSSecretAccessKey(access_key_id=access_key_id, secret_access_key=secret_access_key)
+
+        if bucket is None:
+            raise ValueError("--bucket is required for S3 provider.")
+        file_system = UpdateS3FileSystem(bucket=bucket, credentials=credentials)
+        if endpoint:
+            file_system.endpoint = endpoint
+        if region:
+            file_system.region = region
+        if directory:
+            file_system.directory = directory
+
+    elif gcs:
+        if credentials_path is None:
+            raise ValueError("--credentials-path is required for GCS provider.")
+        with open(credentials_path) as f:
+            service_account = f.read()
+        credentials = GCPServiceAccount(credentials=SecretStr(service_account))
+
+        if region is None or bucket is None:
+            raise ValueError("--region and --bucket is required for GCS provider.")
+        file_system = UpdateGCSFileSystem(bucket=bucket, region=region, credentials=credentials)
+        if directory:
+            file_system.directory = directory
+
+    else:
+        raise ValueError("Must specify either --s3 or --gcs.")
+
+    res = state.settings.api.file_system.update_file_system(project, file_system)
+    rich.print(res.file_system)
+
+
+@app.command(help="Lists the available built-in file system providers.")
+def list_providers():
+    for provider in state.settings.api.file_system.list_providers():
+        rich.print(provider)

spiral/cli/iceberg/__init__.py

@@ -0,0 +1,7 @@
+from spiral.cli import AsyncTyper
+
+from . import namespaces, tables
+
+app = AsyncTyper(short_help="Apache Iceberg Catalog.")
+app.add_typer(tables.app, name="tables")
+app.add_typer(namespaces.app, name="namespaces")

spiral/cli/iceberg/namespaces.py

@@ -0,0 +1,47 @@
+import sys
+from typing import Annotated
+
+import pyiceberg.exceptions
+import rich
+import typer
+from typer import Argument
+
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Apache Iceberg Namespaces.")
+
+
+@app.command(help="List namespaces.")
+def ls(
+    project: ProjectArg,
+    namespace: Annotated[str | None, Argument(help="List only namespaces under this namespace.")] = None,
+):
+    """List Iceberg namespaces."""
+    catalog = state.spiral.iceberg.catalog()
+
+    if namespace is None:
+        try:
+            namespaces = catalog.list_namespaces(project)
+        except pyiceberg.exceptions.ForbiddenError:
+            print(
+                f"The project, {repr(project)}, does not exist or you lack the "
+                f"`iceberg:view` permission to list namespaces in it.",
+                file=sys.stderr,
+            )
+            raise typer.Exit(code=1)
+    else:
+        try:
+            namespaces = catalog.list_namespaces((project, namespace))
+        except pyiceberg.exceptions.ForbiddenError:
+            print(
+                f"The namespace, {repr(project)}.{repr(namespace)}, does not exist or you lack the "
+                f"`iceberg:view` permission to list namespaces in it.",
+                file=sys.stderr,
+            )
+            raise typer.Exit(code=1)
+
+    table = rich.table.Table("Namespace ID", title="Iceberg namespaces")
+    for ns in namespaces:
+        table.add_row(".".join(ns))
+    rich.print(table)