pyspiral 0.7.18__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/cli/__init__.py

@@ -0,0 +1,88 @@
+import asyncio
+import functools
+import inspect
+from collections.abc import Callable
+from typing import IO, Generic, ParamSpec, TypeVar, override
+
+import typer
+from click import ClickException
+from grpclib import GRPCError
+from httpx import HTTPStatusError
+from rich.console import Console
+
+P = ParamSpec("P")
+T = TypeVar("T")
+
+CONSOLE = Console()
+ERR_CONSOLE = Console(stderr=True, style="red")
+
+
+class AsyncTyper(typer.Typer, Generic[P]):
+    """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,
+        )
+
+    @override
+    def callback(self, *args, **kwargs) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        decorator = super().callback(*args, **kwargs)
+        for wrapper in (_wrap_exceptions, _maybe_run_async):
+            decorator = functools.partial(wrapper, decorator)
+        return decorator
+
+    @override
+    def command(self, *args, **kwargs) -> Callable[[Callable[P, T]], Callable[P, T]]:
+        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:
+        ERR_CONSOLE.print(f"Error: {self.format_message()}")
+
+
+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,14 @@
+from spiral.api.types import OrgId
+from spiral.cli import CONSOLE, AsyncTyper, state
+
+app = AsyncTyper()
+
+
+@app.command()
+def sync(
+    org_id: OrgId | None = None,
+):
+    state.spiral.api._admin.sync_orgs()
+
+    for membership in state.spiral.api._admin.sync_memberships(org_id):
+        CONSOLE.print(membership)

spiral/cli/app.py

@@ -0,0 +1,108 @@
+import logging
+import os
+from importlib import metadata
+from logging.handlers import RotatingFileHandler
+from typing import Annotated
+
+import typer
+
+from spiral.cli import (
+    AsyncTyper,
+    admin,
+    console,
+    fs,
+    iceberg,
+    key_spaces,
+    login,
+    orgs,
+    projects,
+    state,
+    tables,
+    telemetry,
+    text,
+    workloads,
+)
+from spiral.settings import LOG_DIR, PACKAGE_NAME
+
+app = AsyncTyper(name="spiral")
+
+
+def version_callback(ctx: typer.Context, value: bool):
+    """
+    Display the version of the Spiral CLI.
+    """
+    # True when generating completion, we can just return
+    if ctx.resilient_parsing:
+        return
+
+    if value:
+        ver = metadata.version(PACKAGE_NAME)
+        print(f"spiral {ver}")
+        raise typer.Exit()
+
+
+def verbose_callback(ctx: typer.Context, value: bool):
+    """
+    Use more verbose output.
+    """
+    # True when generating completion, we can just return
+    if ctx.resilient_parsing:
+        return
+
+    if value:
+        logging.getLogger().setLevel(level=logging.INFO)
+
+
+@app.callback(invoke_without_command=True)
+def _callback(
+    ctx: typer.Context,
+    version: Annotated[
+        bool | None,
+        typer.Option("--version", callback=version_callback, help=version_callback.__doc__, is_eager=True),
+    ] = None,
+    verbose: Annotated[
+        bool | None, typer.Option("--verbose", callback=verbose_callback, help=verbose_callback.__doc__)
+    ] = None,
+) -> None:
+    # Reload the spiral client to support testing under different env vars
+    from spiral import Spiral
+    from spiral.settings import settings
+
+    config = settings()
+    state.spiral = Spiral(config=config)
+
+
+app.add_typer(orgs.app, name="orgs")
+app.add_typer(projects.app, name="projects")
+app.add_typer(fs.app, name="fs")
+app.add_typer(tables.app, name="tables")
+app.add_typer(key_spaces.app, name="ks")
+app.add_typer(text.app, name="text")
+app.add_typer(telemetry.app, name="telemetry")
+app.add_typer(iceberg.app, name="iceberg")
+app.command("login")(login.command)
+app.command("console")(console.command)
+
+
+# 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(admin.app, name="admin", hidden=True)
+    app.add_typer(workloads.app, name="workloads", hidden=True)
+    app.command("whoami", hidden=True)(login.whoami)
+    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,76 @@
+from typing import Literal
+
+import questionary
+from typer import Option
+
+from spiral.api.filesystems import (
+    BuiltinFileSystem,
+    GCSFileSystem,
+    S3FileSystem,
+    UpstreamFileSystem,
+)
+from spiral.cli import CONSOLE, 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.spiral.api.file_system.get_file_system(project)
+    CONSOLE.print(file_system)
+
+
+def ask_provider():
+    res = state.spiral.api.file_system.list_providers()
+    return questionary.select("Select a file system provider", choices=res).ask()
+
+
+@app.command(help="Update a project's default file system.")
+def update(
+    project: ProjectArg,
+    type_: Literal["builtin", "s3", "gcs", "upstream"] = Option(None, "--type", help="Type of the file system."),
+    provider: str = Option(None, help="Provider, when using `builtin` type."),
+    endpoint: str = Option(None, help="Endpoint, when using `s3` type."),
+    region: str = Option(
+        None, help="Region, when using `s3` or `gcs` type (defaults to `auto` for `s3` when `endpoint` is set)."
+    ),
+    bucket: str = Option(None, help="Bucket, when using `s3` or `gcs` type."),
+    role_arn: str = Option(None, help="Role ARN to assume, when using `s3` type."),
+):
+    if type_ == "builtin":
+        provider = provider or ask_provider()
+        file_system = BuiltinFileSystem(provider=provider)
+
+    elif type_ == "upstream":
+        upstream_project = ask_project(title="Select a project to use as file system.")
+        file_system = UpstreamFileSystem(project_id=upstream_project)
+
+    elif type_ == "s3":
+        if role_arn is None:
+            raise ValueError("--role-arn is required for S3 provider.")
+        if not role_arn.startswith("arn:aws:iam::") or ":role/" not in role_arn:
+            raise ValueError("Invalid role ARN format. Expected `arn:aws:iam::<account>:role/<role_name>`")
+        if bucket is None:
+            raise ValueError("--bucket is required for S3 provider.")
+        region = region or ("auto" if endpoint else None)
+        file_system = S3FileSystem(bucket=bucket, role_arn=role_arn, region=region)
+        if endpoint:
+            file_system.endpoint = endpoint
+
+    elif type_ == "gcs":
+        if region is None or bucket is None:
+            raise ValueError("--region and --bucket is required for GCS provider.")
+        file_system = GCSFileSystem(bucket=bucket, region=region)
+
+    else:
+        raise ValueError(f"Unknown file system type: {type_}")
+
+    fs = state.spiral.api.file_system.update_file_system(project, file_system)
+    CONSOLE.print(fs)
+
+
+@app.command(help="Lists the available built-in file system providers.")
+def list_providers():
+    for provider in state.spiral.api.file_system.list_providers():
+        CONSOLE.print(provider)

spiral/cli/iceberg.py

@@ -0,0 +1,97 @@
+import sys
+from typing import Annotated
+
+import rich
+import typer
+from typer import Argument
+
+from spiral.cli import CONSOLE, ERR_CONSOLE, AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Apache Iceberg Catalog.")
+
+
+@app.command(help="List namespaces.")
+def namespaces(
+    project: ProjectArg,
+    namespace: Annotated[str | None, Argument(help="List only namespaces under this namespace.")] = None,
+):
+    """List Iceberg namespaces."""
+    import pyiceberg.exceptions
+
+    catalog = state.spiral.iceberg.catalog()
+
+    if namespace is None:
+        try:
+            namespaces = catalog.list_namespaces(project)
+        except pyiceberg.exceptions.ForbiddenError:
+            ERR_CONSOLE.print(
+                f"The project, {repr(project)}, does not exist or you lack the "
+                f"`iceberg:view` permission to list namespaces in it.",
+            )
+            raise typer.Exit(code=1)
+    else:
+        try:
+            namespaces = catalog.list_namespaces((project, namespace))
+        except pyiceberg.exceptions.ForbiddenError:
+            ERR_CONSOLE.print(
+                f"The namespace, {repr(project)}.{repr(namespace)}, does not exist or you lack the "
+                f"`iceberg:view` permission to list namespaces in it.",
+            )
+            raise typer.Exit(code=1)
+
+    table = CONSOLE.table.Table("Namespace ID", title="Iceberg namespaces")
+    for ns in namespaces:
+        table.add_row(".".join(ns))
+    CONSOLE.print(table)
+
+
+@app.command(help="List tables.")
+def tables(
+    project: ProjectArg,
+    namespace: Annotated[str | None, Argument(help="Show only tables in the given namespace.")] = None,
+):
+    import pyiceberg.exceptions
+
+    catalog = state.spiral.iceberg.catalog()
+
+    try:
+        if namespace is None:
+            tables = catalog.list_tables(project)
+        else:
+            tables = catalog.list_tables((project, namespace))
+    except pyiceberg.exceptions.ForbiddenError:
+        ERR_CONSOLE.print(
+            f"The namespace, {repr(project)}.{repr(namespace)}, does not exist or you lack the "
+            f"`iceberg:view` permission to list tables in it.",
+        )
+        raise typer.Exit(code=1)
+
+    rich_table = rich.table.Table("table id", title="Iceberg tables")
+    for table in tables:
+        rich_table.add_row(".".join(table))
+    CONSOLE.print(rich_table)
+
+
+@app.command(help="Show the table schema.")
+def schema(
+    project: ProjectArg,
+    namespace: Annotated[str, Argument(help="Table namespace.")],
+    table: Annotated[str, Argument(help="Table name.")],
+):
+    import pyiceberg.exceptions
+
+    catalog = state.spiral.iceberg.catalog()
+
+    try:
+        tbl = catalog.load_table((project, namespace, table))
+    except pyiceberg.exceptions.NoSuchTableError:
+        ERR_CONSOLE.print(f"No table {repr(table)} found in {repr(project)}.{repr(namespace)}", file=sys.stderr)
+        raise typer.Exit(code=1)
+
+    rich_table = rich.table.Table(
+        "Field ID", "Field name", "Type", "Required", "Doc", title=f"{project}.{namespace}.{table}"
+    )
+    for col in tbl.schema().columns:
+        rich_table.add_row(str(col.field_id), col.name, str(col.field_type), str(col.required), col.doc)
+    CONSOLE.print(rich_table)

spiral/cli/key_spaces.py

@@ -0,0 +1,103 @@
+from typing import Annotated
+
+import questionary
+import rich
+import typer
+from questionary import Choice
+from typer import Option
+
+from spiral.api.key_space_indexes import SyncIndexRequest
+from spiral.api.projects import KeySpaceIndexResource
+from spiral.api.types import IndexId
+from spiral.api.workers import ResourceClass
+from spiral.cli import CONSOLE, AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Key Space Indexes.")
+
+
+def ask_index(project_id, title="Select an index"):
+    indexes: list[KeySpaceIndexResource] = list(state.spiral.project(project_id).list_key_space_indexes())
+
+    if not indexes:
+        CONSOLE.print("[red]No indexes found[/red]")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        title,
+        choices=[Choice(title=index.name, value=index.id) for index in sorted(indexes, key=lambda t: (t.name, t.id))],
+    ).ask()
+
+
+def get_index_id(
+    project: ProjectArg,
+    name: Annotated[str | None, Option(help="Index name.")] = None,
+) -> IndexId:
+    if name is None:
+        return ask_index(project)
+
+    indexes: list[KeySpaceIndexResource] = list(state.spiral.project(project).list_key_space_indexes())
+    for index in indexes:
+        if index.name == name:
+            return index.id
+    raise ValueError(f"Index not found: {name}")
+
+
+@app.command(help="List indexes.")
+def ls(
+    project: ProjectArg,
+):
+    """List indexes."""
+    indexes = state.spiral.project(project).list_key_space_indexes()
+
+    rich_table = rich.table.Table("id", "name", title="Key Space Indexes")
+    for index in indexes:
+        rich_table.add_row(index.id, index.name)
+    CONSOLE.print(rich_table)
+
+
+@app.command(help="Show index partitions.")
+def show(
+    project: ProjectArg,
+    name: Annotated[str | None, Option(help="Index name.")] = None,
+):
+    """Show index partitions."""
+    index_id = get_index_id(project, name)
+    index = state.spiral.key_space_index(index_id)
+    shards = state.spiral.internal.key_space_index_shards(index.core)
+
+    rich_table = rich.table.Table("Begin", "End", "Cardinality", title=f"Index {index.name} Partitions")
+    for partition in shards:
+        rich_table.add_row(
+            # TODO(marko): This isn't really pretty...
+            repr(partition.key_range.begin),
+            repr(partition.key_range.end),
+            str(partition.cardinality),
+        )
+    CONSOLE.print(rich_table)
+
+
+@app.command(help="Trigger a sync job for an index.")
+def sync(
+    project: ProjectArg,
+    name: Annotated[str | None, Option(help="Index name.")] = None,
+    resources: Annotated[ResourceClass, Option(help="Resources to use for the sync job.")] = ResourceClass.SMALL,
+):
+    """Trigger a sync job."""
+    index_id = get_index_id(project, name)
+    response = state.spiral.api.key_space_indexes.sync_index(index_id, SyncIndexRequest(resources=resources))
+    CONSOLE.print(f"Triggered sync job {response.worker_id} for index {index_id}.")
+
+
+# TODO(marko): This will be removed.
+@app.command(help="Run a sync and wait for it to complete.")
+def sync_local(
+    project: ProjectArg,
+    name: Annotated[str | None, Option(help="Index name.")] = None,
+):
+    """Run a sync and wait for it to complete."""
+    index_id = get_index_id(project, name)
+    index = state.spiral.key_space_index(index_id)
+    snapshot = state.spiral.table(index.table_id).snapshot()
+    state.spiral.internal.update_key_space_index(index.core, snapshot.core)
+    CONSOLE.print(f"Index {index.name} is up to date as-of {snapshot.asof}.")

spiral/cli/login.py

@@ -0,0 +1,25 @@
+import jwt
+
+from spiral.cli import CONSOLE, state
+from spiral.core.authn import DeviceCodeAuth
+
+
+def command(org_id: str | None = None, force: bool = False, show_token: bool = False):
+    token = DeviceCodeAuth.default().authenticate(force=force, org_id=org_id)
+    CONSOLE.print("Successfully logged in.")
+    if show_token:
+        CONSOLE.print(token.expose_secret(), soft_wrap=True)
+
+
+def whoami():
+    """Display the current user's information."""
+    payload = jwt.decode(state.spiral.authn.token().expose_secret(), options={"verify_signature": False})
+
+    if "org_id" in payload:
+        CONSOLE.print(f"{payload['org_id']}")
+    CONSOLE.print(f"{payload['sub']}")
+
+
+def logout():
+    DeviceCodeAuth.default().logout()
+    CONSOLE.print("Logged out.")