pyspiral 0.4.0__pp310-pypy310_pp73-macosx_10_12_x86_64.whl

spiral/cli/iceberg/tables.py

@@ -0,0 +1,60 @@
+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 Tables.")
+
+
+@app.command(help="List tables.")
+def ls(
+    project: ProjectArg,
+    namespace: Annotated[str | None, Argument(help="Show only tables in the given namespace.")] = None,
+):
+    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:
+        print(
+            f"The namespace, {repr(project)}.{repr(namespace)}, does not exist or you lack the "
+            f"`iceberg:view` permission to list tables in it.",
+            file=sys.stderr,
+        )
+        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))
+    rich.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.")],
+):
+    catalog = state.spiral.iceberg.catalog()
+
+    try:
+        tbl = catalog.load_table((project, namespace, table))
+    except pyiceberg.exceptions.NoSuchTableError:
+        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)
+    rich.print(rich_table)

spiral/cli/indexes/__init__.py

@@ -0,0 +1,19 @@
+import rich
+
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper(short_help="Indexes.")
+
+
+@app.command(help="List indexes.")
+def ls(
+    project: ProjectArg,
+):
+    """List indexes."""
+    indexes = state.spiral.project(project).indexes.list_indexes()
+
+    rich_table = rich.table.Table("id", "name", title="Indexes")
+    for index in indexes:
+        rich_table.add_row(index.id, index.name)
+    rich.print(rich_table)

spiral/cli/login.py

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

spiral/cli/orgs.py

@@ -0,0 +1,90 @@
+import webbrowser
+from typing import Annotated
+
+import jwt
+import rich
+import typer
+from rich.table import Table
+from typer import Option
+
+from spiral.api.organizations import CreateOrgRequest, InviteUserRequest, OrgRole, PortalLinkIntent, PortalLinkRequest
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import OrganizationArg
+
+app = AsyncTyper(short_help="Org admin.")
+
+
+@app.command(help="Switch the active organization.")
+def switch(org_id: OrganizationArg):
+    state.settings.device_code_auth.authenticate(org_id=org_id)
+    rich.print(f"Switched to organization: {org_id}")
+
+
+@app.command(help="Create a new organization.")
+def create(
+    name: Annotated[str | None, Option(help="The human-readable name of the organization.")] = None,
+):
+    res = state.settings.api.organization.create(CreateOrgRequest(name=name))
+
+    # Authenticate to the new organization
+    state.settings.device_code_auth.authenticate(org_id=res.org.id)
+
+    rich.print(f"{res.org.name} [dim]{res.org.id}[/dim]")
+
+
+@app.command(help="List organizations.")
+def ls():
+    org_id = current_org_id()
+
+    table = Table("", "id", "name", "role", title="Organizations")
+    for m in state.settings.api.organization.list_memberships():
+        table.add_row("👉" if m.org.id == org_id else "", m.org.id, m.org.name, m.role)
+
+    rich.print(table)
+
+
+@app.command(help="Invite a user to the organization.")
+def invite(email: str, role: OrgRole = OrgRole.MEMBER, expires_in_days: int = 7):
+    state.settings.api.organization.invite_user(
+        InviteUserRequest(email=email, role=role, expires_in_days=expires_in_days)
+    )
+    rich.print(f"Invited {email} as a {role.value}.")
+
+
+@app.command(help="Configure single sign-on for your organization.")
+def sso():
+    _do_action(PortalLinkIntent.SSO)
+
+
+@app.command(help="Configure directory services for your organization.")
+def directory():
+    _do_action(PortalLinkIntent.DIRECTORY_SYNC)
+
+
+@app.command(help="Configure audit logs for your organization.")
+def audit_logs():
+    _do_action(PortalLinkIntent.AUDIT_LOGS)
+
+
+@app.command(help="Configure log streams for your organization.")
+def log_streams():
+    _do_action(PortalLinkIntent.LOG_STREAMS)
+
+
+@app.command(help="Configure domains for your organization.")
+def domains():
+    _do_action(PortalLinkIntent.DOMAIN_VERIFICATION)
+
+
+def _do_action(intent: PortalLinkIntent):
+    res = state.settings.api.organization.portal_link(PortalLinkRequest(intent=intent))
+    rich.print(f"Opening the configuration portal:\n{res.url}")
+    webbrowser.open(res.url)
+
+
+def current_org_id():
+    if token := state.settings.authn.token():
+        if org_id := jwt.decode(token.expose_secret(), options={"verify_signature": False}).get("org_id"):
+            return org_id
+    rich.print("[red]You are not logged in to an organization.[/red]")
+    raise typer.Exit(1)

spiral/cli/printer.py

@@ -0,0 +1,53 @@
+import json
+from collections.abc import Iterable
+from typing import TypeVar
+
+import betterproto
+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=betterproto.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()._betterproto.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(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,136 @@
+from typing import Annotated
+
+import rich
+import typer
+from typer import Option
+
+from spiral.api.organizations import OrgRole
+from spiral.api.projects import (
+    CreateProjectRequest,
+    GCPPrincipalConditions,
+    GitHubConditions,
+    GitHubPrincipalConditions,
+    Grant,
+    GrantRoleRequest,
+    ModalConditions,
+    ModalPrincipalConditions,
+    OrgRolePrincipalConditions,
+    OrgUserPrincipalConditions,
+    Project,
+    WorkloadPrincipalConditions,
+)
+from spiral.cli import 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())
+    rich.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))
+    rich.print(f"Created project {res.project.id}")
+
+
+@app.command(help="Grant a role on a project.")
+def grant(
+    project: ProjectArg,
+    role: Annotated[str, Option(help="Role to grant.")],
+    org_id: Annotated[
+        str | None, Option(help="Pass an organization ID to grant a role to an organization user(s).")
+    ] = None,
+    user_id: 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[
+        str | None,
+        Option(help="Pass an organization 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: Annotated[
+        str | None,
+        Option(help="Pass a `<service_account_email>/<unique_id>` to grant a role to a GCP service account."),
+    ] = None,
+    conditions: list[str] | None = Option(
+        default=None,
+        help="`<key>=<value>` token conditions to apply to the grant when using --github or --modal.",
+    ),
+):
+    # Check mutual exclusion
+    if sum(int(bool(opt)) for opt in {org_id, workload_id, github, modal, gcp}) != 1:
+        raise typer.BadParameter("Only one of --org-id, --github or --modal 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 {user_id, org_role}) != 1:
+            raise typer.BadParameter("Only one of --user-id or --org-role may be specified.")
+
+        if user_id is not None:
+            principal = OrgUserPrincipalConditions(org_id=org_id, user_id=user_id)
+        elif org_role is not None:
+            principal = OrgRolePrincipalConditions(org_id=org_id, role=OrgRole(org_role))
+        else:
+            raise NotImplementedError("Only user or role principal is supported at this time.")
+
+    elif workload_id:
+        principal = WorkloadPrincipalConditions(workload_id=workload_id)
+
+    elif gcp:
+        service_account, unique_id = gcp.split("/", 1)
+        principal = GCPPrincipalConditions(service_account=service_account, unique_id=unique_id)
+
+    else:
+        raise ValueError("Invalid grant principal")
+
+    state.settings.api.project.grant_role(
+        project,
+        GrantRoleRequest(
+            role_id=role,
+            principal=principal,
+        ),
+    )
+
+    rich.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))
+    rich.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/__init__.py

@@ -0,0 +1,121 @@
+from typing import Annotated
+
+import questionary
+import rich
+import typer
+from questionary import Choice
+from typer import Argument, Option
+
+from spiral import Spiral
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import ProjectArg
+from spiral.tables import Table
+
+app = AsyncTyper(short_help="Spiral Tables.")
+
+
+def ask_table(project_id, title="Select a table"):
+    tables = list(state.spiral.project(project_id).tables.list_tables())
+
+    if not tables:
+        rich.print("[red]No tables found[/red]")
+        raise typer.Exit(1)
+
+    return questionary.select(
+        title,
+        choices=[
+            Choice(title=f"{table.dataset}.{table.table}", value=f"{table.project_id}.{table.dataset}.{table.table}")
+            for table in tables
+        ],
+    ).ask()
+
+
+@app.command(help="List tables.")
+def ls(
+    project: ProjectArg,
+):
+    tables = Spiral().project(project).tables.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)
+    rich.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,
+):
+    _, table = _get_table(project, table, dataset)
+    rich.print(table.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,
+):
+    _, table = _get_table(project, table, dataset)
+    rich.print(table.schema)
+
+
+@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,
+):
+    identifier, table = _get_table(project, table, dataset)
+    table.maintenance().flush_wal()
+    print(f"Flushed WAL for table {identifier} in project {project}.")
+
+
+@app.command(help="Display scan.")
+def debug(
+    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.")] = ".",
+):
+    _, table = _get_table(project, table, dataset)
+    if column_group != ".":
+        projection = table[column_group]
+    else:
+        projection = table
+    scan = table.scan(projection)
+
+    scan._debug()
+
+
+@app.command(help="Display manifests.")
+def manifests(
+    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.")] = ".",
+):
+    _, table = _get_table(project, table, dataset)
+    if column_group != ".":
+        projection = table[column_group]
+    else:
+        projection = table
+    scan = projection.scan()
+
+    scan._dump_manifests()
+
+
+def _get_table(
+    project: ProjectArg,
+    table: Annotated[str | None, Option(help="Table name.")] = None,
+    dataset: Annotated[str | None, Option(help="Dataset name.")] = None,
+) -> (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).tables.table(identifier)

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

@@ -0,0 +1,51 @@
+from typing import Annotated
+
+import questionary
+import rich
+import typer
+from questionary import Choice
+from typer import Argument
+
+from spiral.api.types import OrgId, ProjectId
+from spiral.cli import state
+
+
+def ask_project(title="Select a project"):
+    projects = list(state.settings.api.project.list())
+
+    if not projects:
+        rich.print("[red]No projects found[/red]")
+        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:
+        rich.print("[red]No organizations found[/red]")
+        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,59 @@
+from typing import Annotated
+
+import pyperclip
+import questionary
+import rich
+from questionary import Choice
+from typer import Argument, Option
+
+from spiral.api.workloads import CreateWorkloadRequest, IssueWorkloadCredentialsResponse, Workload
+from spiral.cli import 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))
+    rich.print(f"Created workload {res.workload.id}")
+
+
+@app.command(help="List workloads.")
+def ls(
+    project: ProjectArg,
+):
+    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 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)
+            rich.print("[green]Secret copied to clipboard![/green]")
+            break
+        elif choice == 2:
+            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.client_id}")