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

spiral/cli/fs.py

@@ -2,46 +2,111 @@ from typing import Annotated
 
 import questionary
 import rich
+from pydantic import SecretStr
 from typer import Option
 
-from spiral.api.filesystems import BuiltinFileSystem, GetFileSystem, UpdateFileSystem
+from spiral.api.filesystems import (
+    AWSSecretAccessKey,
+    BuiltinFileSystem,
+    GCPServiceAccount,
+    UpdateGCSFileSystem,
+    UpdateS3FileSystem,
+    UpstreamFileSystem,
+)
 from spiral.cli import AsyncTyper, state
-from spiral.cli.types import ProjectArg
+from spiral.cli.types import ProjectArg, ask_project
 
-app = AsyncTyper()
+app = AsyncTyper(short_help="File Systems.")
 
 
 @app.command(help="Show the file system configured for project.")
 def show(project: ProjectArg):
-    res = state.settings.api.file_system.get_file_system(GetFileSystem.Request(project_id=project))
-    match res.file_system:
+    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(res.file_system)
+            rich.print(file_system)
 
 
-def _provider_default():
+def ask_provider():
     res = state.settings.api.file_system.list_providers()
-    return questionary.select("Select a file system provider", choices=res.providers).ask()
+    return questionary.select("Select a file system provider", choices=res).ask()
 
 
-ProviderOpt = Annotated[
+BuiltinProviderOpt = Annotated[
     str,
-    Option(help="Built-in provider to use for the file system.", show_default=False, default_factory=_provider_default),
+    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 file system.")
-def update(project: ProjectArg, provider: ProviderOpt):
-    res = state.settings.api.file_system.update_file_system(
-        UpdateFileSystem.Request(project_id=project, file_system=BuiltinFileSystem(provider=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():
-    res = state.settings.api.file_system.list_providers()
-    for provider in res.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)

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

@@ -1,13 +1,22 @@
+import jwt
 from rich import print
 
-from spiral.cli import OptionalStr, state
+from spiral.cli import state
 
 
-def command(org_id: OptionalStr = None, force: bool = False, refresh: bool = False):
-    tokens = state.settings.spiraldb.device_auth().authenticate(force=force, refresh=refresh, organization_id=org_id)
-    print(tokens)
+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.spiraldb.device_auth().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

@@ -1,3 +1,4 @@
+import json
 from collections.abc import Iterable
 from typing import TypeVar
 
@@ -17,10 +18,17 @@ def table_of_models(cls: type[M], messages: Iterable[M], fields: list[str] = Non
     cols = fields or cls.model_fields.keys()
     table = Table(*cols, title=title or f"{cls.__name__}s")
     for msg in messages:
-        table.add_row(*[getattr(msg, col, "") for col in cols])
+        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.
 

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

@@ -1,3 +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)