pyspiral 0.6.8__cp312-abi3-manylinux_2_28_x86_64.whl

spiral/cli/printer.py

@@ -0,0 +1,53 @@
+import json
+from collections.abc import Iterable
+from typing import TypeVar
+
+import betterproto2
+from pydantic import BaseModel
+from rich.console import ConsoleRenderable, Group
+from rich.padding import Padding
+from rich.pretty import Pretty
+from rich.table import Table
+
+T = TypeVar("T", bound=betterproto2.Message)
+M = TypeVar("M", bound=BaseModel)
+
+
+def table_of_models(cls: type[M], messages: Iterable[M], fields: list[str] = None, title: str = None) -> Table:
+    """Centralized logic for printing tables of Pydantic models."""
+    cols = fields or cls.model_fields.keys()
+    table = Table(*cols, title=title or f"{cls.__name__}s")
+    for msg in messages:
+        table.add_row(*[_renderable(msg, col) for col in cols])
+    return table
+
+
+def _renderable(msg, col):
+    attr = getattr(msg, col, "")
+    if isinstance(attr, dict):
+        return json.dumps(attr)
+    return attr
+
+
+def table_of_protos(cls: type[T], messages: Iterable[T], fields: list[str] = None, title: str = None) -> Table:
+    """Centralized logic for printing tables of proto messages.
+
+    TODO(ngates): add a CLI switch to emit JSON results instead of tables.
+    """
+    cols = fields or cls()._betterproto2.sorted_field_names
+    table = Table(*cols, title=title or f"{cls.__name__}s")
+    for msg in messages:
+        table.add_row(*[getattr(msg, col, "") for col in cols])
+    return table
+
+
+def proto[T](message: T, title: str = None, fields: list[str] = None) -> ConsoleRenderable:
+    """Centralized logic for printing a single proto message."""
+    value = Pretty({k: v for k, v in message.to_dict().items() if not fields or k in fields})
+    if title:
+        return Group(
+            f"[bold]{title}[/bold]",
+            Padding.indent(value, level=2),
+        )
+    else:
+        return value

spiral/cli/projects.py

@@ -0,0 +1,147 @@
+from typing import Annotated, Literal
+
+import typer
+from typer import Option
+
+from spiral.api.organizations import OrgRole
+from spiral.api.projects import (
+    AwsAssumedRolePrincipalConditions,
+    CreateProjectRequest,
+    GcpServiceAccountPrincipalConditions,
+    GitHubConditions,
+    GitHubPrincipalConditions,
+    Grant,
+    GrantRoleRequest,
+    ModalConditions,
+    ModalPrincipalConditions,
+    OrgRolePrincipalConditions,
+    OrgUserPrincipalConditions,
+    Project,
+    WorkloadPrincipalConditions,
+)
+from spiral.cli import CONSOLE, AsyncTyper, printer, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Projects and grants.")
+
+
+@app.command(help="List projects.")
+def ls():
+    projects = list(state.settings.api.project.list())
+    CONSOLE.print(printer.table_of_models(Project, projects))
+
+
+@app.command(help="Create a new project.")
+def create(
+    id_prefix: Annotated[
+        str | None, Option(help="An optional ID prefix to which a random number will be appended.")
+    ] = None,
+    name: Annotated[str | None, Option(help="Friendly name for the project.")] = None,
+):
+    res = state.settings.api.project.create(CreateProjectRequest(id_prefix=id_prefix, name=name))
+    CONSOLE.print(f"Created project {res.project.id}")
+
+
+@app.command(help="Grant a role on a project to a principal.")
+def grant(
+    project: ProjectArg,
+    role: Annotated[Literal["viewer", "editor", "admin"], Option(help="Project role to grant.")],
+    org_id: Annotated[
+        str | None, Option(help="Pass an organization ID to grant a role to an organization user(s).")
+    ] = None,
+    org_user: Annotated[
+        str | None, Option(help="Pass a user ID when using --org-id to grant a role to grant a role to a user.")
+    ] = None,
+    org_role: Annotated[
+        Literal["owner", "member", "guest"] | None,
+        Option(help="Pass an org role when using --org-id to grant a role to all users with that role."),
+    ] = None,
+    workload_id: Annotated[str | None, Option(help="Pass a workload ID to grant a role to a workload.")] = None,
+    github: Annotated[
+        str | None, Option(help="Pass an `{org}/{repo}` string to grant a role to a job running in GitHub Actions.")
+    ] = None,
+    modal: Annotated[
+        str | None,
+        Option(help="Pass a `{workspace_id}/{env_name}` string to grant a role to a job running in Modal environment."),
+    ] = None,
+    gcp_service_account: Annotated[
+        str | None,
+        Option(help="Pass a `{service_account_email}/{unique_id}` to grant a role to a GCP service account."),
+    ] = None,
+    aws_iam_role: Annotated[
+        str | None,
+        Option(help="Pass a `{account_id}/{role_name}` to grant a Spiral role to an AWS IAM Role."),
+    ] = None,
+    conditions: list[str] | None = Option(
+        default=None,
+        help="`{key}={value}` token conditions to apply to the grant",
+    ),
+):
+    # Check mutual exclusion
+    if sum(int(bool(opt)) for opt in {org_id, workload_id, github, modal, gcp_service_account, aws_iam_role}) != 1:
+        raise typer.BadParameter(
+            "Only one of [--org-id, --workload-id, --github, --modal, --gcp-service-account, --aws-iam-role] "
+            "may be specified."
+        )
+
+    if github:
+        org, repo = github.split("/", 1)
+        github_conditions = None
+        if conditions is not None:
+            github_conditions = GitHubConditions()
+            for k, v in dict(c.split("=", 1) for c in conditions).items():
+                github_conditions = github_conditions.model_copy(update={k: v})
+        principal = GitHubPrincipalConditions(org=org, repo=repo, conditions=github_conditions)
+
+    elif modal:
+        workspace_id, environment_name = modal.split("/", 1)
+        modal_conditions = None
+        if conditions is not None:
+            modal_conditions = ModalConditions()
+            for k, v in dict(c.split("=", 1) for c in conditions).items():
+                modal_conditions = modal_conditions.model_copy(update={k: v})
+        principal = ModalPrincipalConditions(
+            workspace_id=workspace_id, environment_name=environment_name, conditions=modal_conditions
+        )
+
+    elif org_id:
+        # Check mutual exclusion
+        if sum(int(bool(opt)) for opt in {org_user, org_role}) != 1:
+            raise typer.BadParameter("Only one of --org-user or --org-role may be specified.")
+
+        if org_user is not None:
+            principal = OrgUserPrincipalConditions(org_id=org_id, user_id=org_user)
+        elif org_role is not None:
+            principal = OrgRolePrincipalConditions(org_id=org_id, role=OrgRole(org_role))
+        else:
+            raise typer.BadParameter("One of --org-user or --org-role must be specified with --org-id.")
+
+    elif workload_id:
+        principal = WorkloadPrincipalConditions(workload_id=workload_id)
+
+    elif gcp_service_account:
+        service_account, unique_id = gcp_service_account.split("/", 1)
+        principal = GcpServiceAccountPrincipalConditions(service_account=service_account, unique_id=unique_id)
+
+    elif aws_iam_role:
+        account_id, role_name = aws_iam_role.split("/", 1)
+        principal = AwsAssumedRolePrincipalConditions(account_id=account_id, role_name=role_name)
+
+    else:
+        raise ValueError("Invalid grant principal")
+
+    state.settings.api.project.grant_role(
+        project,
+        GrantRoleRequest(
+            role_id=role,
+            principal=principal,
+        ),
+    )
+
+    CONSOLE.print(f"Granted role {role} on project {project}")
+
+
+@app.command(help="List project grants.")
+def grants(project: ProjectArg):
+    project_grants = list(state.settings.api.project.list_grants(project))
+    CONSOLE.print(printer.table_of_models(Grant, project_grants, title="Project Grants"))

spiral/cli/state.py

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

spiral/cli/tables.py

@@ -0,0 +1,174 @@
+import datetime
+from typing import Annotated, Literal
+
+import questionary
+import rich
+import rich.table
+import typer
+from questionary import Choice
+from typer import Argument, Option
+
+from spiral import Spiral
+from spiral.api.projects import TableResource
+from spiral.cli import CONSOLE, ERR_CONSOLE, AsyncTyper, state
+from spiral.cli.types import ProjectArg
+from spiral.debug.manifests import display_manifests
+from spiral.table import Table
+
+app = AsyncTyper(short_help="Spiral Tables.")
+
+
+def ask_table(project_id: str, title: str = "Select a table") -> str:
+    tables: list[TableResource] = list(state.spiral.project(project_id).list_tables())
+
+    if not tables:
+        ERR_CONSOLE.print("No tables found")
+        raise typer.Exit(1)
+
+    return questionary.select(  # pyright: ignore[reportAny]
+        title,
+        choices=[
+            Choice(title=f"{table.dataset}.{table.table}", value=f"{table.dataset}.{table.table}")
+            for table in sorted(tables, key=lambda t: (t.dataset, t.table))
+        ],
+    ).ask()
+
+
+def get_table(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+) -> tuple[str, Table]:
+    if table is None:
+        identifier = ask_table(project)
+    else:
+        identifier = table
+        if dataset is not None:
+            identifier = f"{dataset}.{table}"
+    return identifier, state.spiral.project(project).table(identifier)
+
+
+@app.command(help="List tables.")
+def ls(
+    project: ProjectArg,
+):
+    tables = Spiral().project(project).list_tables()
+
+    rich_table = rich.table.Table("id", "dataset", "name", title="Spiral tables")
+    for table in tables:
+        rich_table.add_row(table.id, table.dataset, table.table)
+    CONSOLE.print(rich_table)
+
+
+@app.command(help="Show the table key schema.")
+def key_schema(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+):
+    _, t = get_table(project, table, dataset)
+    CONSOLE.print(t.key_schema)
+
+
+@app.command(help="Compute the full table schema.")
+def schema(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+):
+    _, t = get_table(project, table, dataset)
+    CONSOLE.print(t.schema())
+
+
+@app.command(help="Fetch Write-Ahead-Log.")
+def wal(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+):
+    _, t = get_table(project, table, dataset)
+    wal_ = t.core.get_wal(asof=None)
+    # Don't use CONSOLE.print here so that it can be piped.
+    print(wal_)
+
+
+@app.command(help="Flush Write-Ahead-Log.")
+def flush(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+    keep: Annotated[
+        Literal["1h", "2h", "4h"] | None,
+        Option(help="Duration string that indicates how much WAL to keep. Defaults to 24h."),
+    ] = None,
+    full: Annotated[bool, Option(help="Flush full Write-Ahead-Log.")] = False,
+):
+    # TODO(marko): Use some human-readable duration parsing library.
+    duration = None
+    if keep is not None:
+        if full:
+            raise ValueError("Cannot specify both --keep and --full")
+        match keep:
+            case "1h":
+                duration = datetime.timedelta(hours=1)
+            case "2h":
+                duration = datetime.timedelta(hours=2)
+            case "4h":
+                duration = datetime.timedelta(hours=4)
+            case _:
+                raise ValueError(f"Invalid duration string: {keep}")
+
+    if full:
+        # Warn and wait for confirmation.
+        ERR_CONSOLE.print("[bold yellow]Warning: All currently open transaction will fail to commit.[/bold yellow]")
+        if not questionary.confirm("Are you sure you want to continue?", default=False).ask():  # pyright: ignore[reportAny]
+            ERR_CONSOLE.print("Aborting.")
+            raise typer.Exit(1)
+
+        duration = datetime.timedelta(hours=0)
+
+    keep_latest_s = int(duration.total_seconds()) if duration is not None else None
+
+    identifier, t = get_table(project, table, dataset)
+    state.spiral._ops().flush_wal(t.core, keep_latest_s=keep_latest_s)  # pyright: ignore[reportPrivateUsage]
+    CONSOLE.print(f"Flushed WAL for table {identifier} in project {project}.")
+
+
+@app.command(help="Display all manifests.")
+def manifests(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+):
+    _, t = get_table(project, table, dataset)
+    s = t.snapshot()
+
+    key_space_state = state.spiral._ops().key_space_state(s.core)  # pyright: ignore[reportPrivateUsage]
+    key_space_manifest = key_space_state.manifest
+
+    column_groups_states = state.spiral._ops().column_groups_states(s.core, key_space_state)  # pyright: ignore[reportPrivateUsage]
+    display_manifests(key_space_manifest, [(x.column_group, x.manifest) for x in column_groups_states])
+
+
+@app.command(help="Visualize the scan of a given column group.")
+def debug_scan(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+    column_group: Annotated[str, Argument(help="Dot-separated column group path.")] = ".",
+):
+    _, t = get_table(project, table, dataset)
+    scan = state.spiral.scan(t[column_group] if column_group != "." else t)
+    scan._debug()  # pyright: ignore[reportPrivateUsage]
+
+
+@app.command(help="Display the manifests for a scan of a given column group.")
+def dump_scan(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+    column_group: Annotated[str, Argument(help="Dot-separated column group path.")] = ".",
+):
+    _, t = get_table(project, table, dataset)
+    scan = state.spiral.scan(t[column_group] if column_group != "." else t)
+    scan._dump_manifests()  # pyright: ignore[reportPrivateUsage]

spiral/cli/telemetry.py

@@ -0,0 +1,17 @@
+import pyperclip
+
+from spiral.api.telemetry import IssueExportTokenResponse
+from spiral.cli import CONSOLE, 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)
+
+    CONSOLE.print("Export command copied to clipboard! Paste and run to set [green]SPIRAL_OTEL_TOKEN[/green].")
+    CONSOLE.print("[dim]Token is valid for 1h.[/dim]")

spiral/cli/text.py

@@ -0,0 +1,115 @@
+from typing import Annotated
+
+import questionary
+import rich
+import typer
+from questionary import Choice
+from typer import Option
+
+from spiral.api.projects import TextIndexResource
+from spiral.api.text_indexes import CreateWorkerRequest, SyncIndexRequest
+from spiral.api.types import IndexId
+from spiral.api.workers import CPU, GcpRegion, Memory, ResourceClass
+from spiral.cli import CONSOLE, ERR_CONSOLE, AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Text Indexes.")
+
+
+def ask_index(project_id, title="Select an index"):
+    indexes: list[TextIndexResource] = list(state.spiral.project(project_id).list_text_indexes())
+
+    if not indexes:
+        ERR_CONSOLE.print("No indexes found")
+        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[TextIndexResource] = list(state.spiral.project(project).list_text_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_text_indexes()
+
+    rich_table = rich.table.Table("id", "name", title="Text Indexes")
+    for index in indexes:
+        rich_table.add_row(index.id, index.name)
+    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.text_indexes.sync_index(index_id, SyncIndexRequest(resources=resources))
+    CONSOLE.print(f"Triggered sync job {response.worker_id} for index {index_id}.")
+
+
+@app.command(name="serve", help="Spin up a worker to serve an index.")
+def serve(
+    project: ProjectArg,
+    index: Annotated[str | None, Option(help="Index name.")] = None,
+    region: Annotated[GcpRegion, Option(help="GCP region for the worker.")] = GcpRegion.US_EAST4,
+    cpu: Annotated[CPU, Option(help="CPU resources for the worker.")] = CPU.ONE,
+    memory: Annotated[Memory, Option(help="Memory resources for the worker in MB.")] = Memory.MB_512,
+):
+    """Spin up a worker."""
+    index_id = get_index_id(project, index)
+    request = CreateWorkerRequest(cpu=cpu, memory=memory, region=region)
+    response = state.spiral.api.text_indexes.create_worker(index_id, request)
+    CONSOLE.print(f"Created worker {response.worker_id} for {index_id}.")
+
+
+@app.command(name="workers", help="List search workers.")
+def workers(
+    project: ProjectArg,
+    index: Annotated[str | None, Option(help="Index name.")] = None,
+):
+    """List text search workers."""
+    index_id = get_index_id(project, index)
+    worker_ids = state.spiral.api.text_indexes.list_workers(index_id)
+
+    rich_table = rich.table.Table("Worker ID", "URL", title=f"Text Search Workers for {index_id}")
+    for worker_id in worker_ids:
+        try:
+            worker = state.spiral.api.text_indexes.get_worker(worker_id)
+            rich_table.add_row(
+                worker_id,
+                worker.url,
+            )
+        except Exception:
+            rich_table.add_row(
+                worker_id,
+                "Unavailable",
+            )
+    CONSOLE.print(rich_table)
+
+
+@app.command(name="shutdown", help="Shutdown a search worker.")
+def shutdown(worker_id: str):
+    """Shutdown a worker."""
+    state.spiral.api.text_indexes.shutdown_worker(worker_id)
+    CONSOLE.print(f"Requested worker {worker_id} to shutdown.")

spiral/cli/types.py

@@ -0,0 +1,50 @@
+from typing import Annotated
+
+import questionary
+import typer
+from questionary import Choice
+from typer import Argument
+
+from spiral.api.types import OrgId, ProjectId
+from spiral.cli import ERR_CONSOLE, state
+
+
+def ask_project(title="Select a project"):
+    projects = list(state.settings.api.project.list())
+
+    if not projects:
+        ERR_CONSOLE.print("No projects found")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        title,
+        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=ask_project)]
+
+
+def _org_default():
+    memberships = list(state.settings.api.organization.list_memberships())
+
+    if not memberships:
+        ERR_CONSOLE.print("No organizations found")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        "Select an organization",
+        choices=[
+            Choice(
+                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[OrgId, Argument(help="Organization ID", show_default=False, default_factory=_org_default)]

spiral/cli/workloads.py

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