pyspiral 0.1.0__cp310-abi3-macosx_11_0_arm64.whl

spiral/authn/github_.py

@@ -0,0 +1,33 @@
+import os
+
+import httpx
+
+from spiral.api import Authn
+
+
+class GitHubActionsProvider(Authn):
+    AUDIENCE = "https://iss.spiraldb.com"
+
+    def __init__(self):
+        self._gh_token = None
+
+    def token(self) -> str | None:
+        if self._gh_token is not None:
+            return self._gh_token
+
+        if os.environ.get("GITHUB_ACTIONS") == "true":
+            # Next, we check to see if we're running in GitHub actions and if so, grab an ID token.
+            if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" in os.environ and "ACTIONS_ID_TOKEN_REQUEST_URL" in os.environ:
+                if not hasattr(self, "__gh_token"):
+                    resp = httpx.get(
+                        f"{os.environ['ACTIONS_ID_TOKEN_REQUEST_URL']}&audience={self.AUDIENCE}",
+                        headers={"Authorization": f'Bearer {os.environ["ACTIONS_ID_TOKEN_REQUEST_TOKEN"]}'},
+                    )
+                    if not resp.is_success:
+                        raise ValueError(f"Failed to get GitHub Actions ID token: {resp.text}", resp)
+                    self._gh_token = resp.json()["value"]
+            else:
+                raise ValueError("Please set 'id-token: write' permission for this GitHub Actions workflow.")
+
+        # For now, we don't exchange the token for a Spiral one.
+        return self._gh_token

spiral/authn/modal_.py

@@ -0,0 +1,18 @@
+import os
+
+from spiral.api import Authn
+
+
+class ModalProvider(Authn):
+    def __init__(self):
+        self._modal_token = None
+
+    def token(self) -> str | None:
+        if self._modal_token is not None:
+            return self._modal_token
+
+        if os.environ.get("MODAL_IDENTITY_TOKEN") is not None:
+            self._modal_token = os.environ["MODAL_IDENTITY_TOKEN"]
+
+        # For now, we don't exchange the token for a Spiral one.
+        return self._modal_token

spiral/catalog.py

@@ -0,0 +1,78 @@
+from typing import TYPE_CHECKING
+
+import jwt
+
+from spiral.api import SpiralAPI
+from spiral.api.projects import CreateProject
+from spiral.settings import Settings, settings
+
+if TYPE_CHECKING:
+    from spiral.project import Project
+    from spiral.table import Table
+
+_default = object()
+
+
+class Spiral:
+    def __init__(self, config: Settings | None = None):
+        self._config = config or settings()
+        self._api = self._config.api
+
+        self._org = None
+
+    @property
+    def config(self) -> Settings:
+        return self._config
+
+    @property
+    def api(self) -> SpiralAPI:
+        return self._api
+
+    @property
+    def organization(self) -> str:
+        if self._org is None:
+            token_payload = jwt.decode(self._config.authn.token(), options={"verify_signature": False})
+            if "org_id" not in token_payload:
+                raise ValueError("Please create an organization.")
+            self._org = token_payload["org_id"]
+        return self._org
+
+    def list_projects(self) -> list["Project"]:
+        """List project IDs."""
+        from .project import Project
+
+        return [Project(self, id=p.id, name=p.name) for p in self.api.project.list()]
+
+    def list_project_ids(self) -> list[str]:
+        """List project IDs."""
+        return [p.id for p in self.list_projects()]
+
+    def create_project(
+        self,
+        id_prefix: str | None = None,
+        *,
+        org: str | None = None,
+        name: str | None = None,
+    ) -> "Project":
+        """Create a project in the current, or given, organization."""
+        from .project import Project
+
+        org = org or self.organization
+        res = self.api.project.create(CreateProject.Request(organization_id=org, id_prefix=id_prefix, name=name))
+        return Project(self, res.project.id, name=res.project.name)
+
+    def project(self, project_id: str) -> "Project":
+        """Open an existing project."""
+        from .project import Project
+
+        # We avoid an API call since we'd just be fetching a human-readable name. Seems a waste in most cases.
+        return Project(self, id=project_id, name=project_id)
+
+    def table(self, identifier: str) -> "Table":
+        """Open a table with a "project.dataset.table" identifier."""
+        parts = identifier.split(".")
+        if len(parts) != 3:
+            raise ValueError(f"Invalid table identifier: {identifier}")
+        project_id, dataset, table = parts
+
+        return self.project(project_id).table(f"{dataset}.{table}")

spiral/cli/__init__.py

@@ -0,0 +1,82 @@
+import asyncio
+import functools
+import inspect
+from typing import IO, Optional
+
+import rich
+import typer
+from click import ClickException
+from grpclib import GRPCError
+from httpx import HTTPStatusError
+
+# We need to use Optional[str] since Typer doesn't support str | None.
+OptionalStr = Optional[str]  # noqa: UP007
+
+
+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)
+        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,21 @@
+from rich import print
+
+from spiral.api.admin import SyncMemberships, SyncOrgs
+from spiral.cli import AsyncTyper, state
+
+app = AsyncTyper()
+
+
+@app.command()
+def sync(orgs: bool = False, memberships: bool = False):
+    run_all = True
+    if any([orgs, memberships]):
+        run_all = False
+
+    if run_all or orgs:
+        for org_id in state.settings.api._admin.sync_orgs(SyncOrgs.Request()):
+            print(org_id)
+
+    if run_all or memberships:
+        for membership in state.settings.api._admin.sync_memberships(SyncMemberships.Request()):
+            print(membership)

spiral/cli/app.py

@@ -0,0 +1,48 @@
+import logging
+import os
+from logging.handlers import RotatingFileHandler
+
+from spiral.cli import AsyncTyper, admin, console, fs, login, org, project, state, table, token, workload
+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(org.app, name="org")
+app.add_typer(project.app, name="project")
+app.add_typer(table.app, name="table")
+app.add_typer(workload.app, name="workload")
+app.add_typer(token.app, name="token")
+app.command("console")(console.command)
+app.command("login")(login.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.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.api 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,47 @@
+from typing import Annotated
+
+import questionary
+import rich
+from typer import Option
+
+from spiral.api.filesystems import BuiltinFileSystem, GetFileSystem, UpdateFileSystem
+from spiral.cli import AsyncTyper, state
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@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:
+        case BuiltinFileSystem(provider=provider):
+            rich.print(f"provider: {provider}")
+        case _:
+            rich.print(res.file_system)
+
+
+def _provider_default():
+    res = state.settings.api.file_system.list_providers()
+    questionary.select("Select a file system provider", choices=res.providers).ask()
+
+
+ProviderOpt = Annotated[
+    str,
+    Option(help="Built-in provider to use for the file system.", show_default=False, default_factory=_provider_default),
+]
+
+
+@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))
+    )
+    rich.print(res.file_system.builtin)
+
+
+@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:
+        rich.print(provider)

spiral/cli/login.py

@@ -0,0 +1,13 @@
+from rich import print
+
+from spiral.cli import OptionalStr, 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 logout():
+    state.settings.spiraldb.device_auth().logout()
+    print("Logged out.")

spiral/cli/org.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 CreateOrganization, InviteUser, OrganizationRole, PortalLink
+from spiral.cli import AsyncTyper, OptionalStr, state
+from spiral.cli.types import OrganizationArg
+
+app = AsyncTyper()
+
+
+@app.command(help="Switch the active organization.")
+def switch(org_id: OrganizationArg):
+    state.settings.spiraldb.device_auth().authenticate(refresh=True, organization_id=org_id)
+    rich.print(f"Switched to organization: {org_id}")
+
+
+@app.command(help="Create a new organization.")
+def create(
+    name: Annotated[OptionalStr, Option(help="The human-readable name of the organization.")] = None,
+):
+    res = state.settings.api.organization.create_organization(CreateOrganization.Request(name=name))
+
+    # Authenticate to the new organization
+    state.settings.spiraldb.device_auth().authenticate(refresh=True, organization_id=res.organization.id)
+
+    rich.print(f"{res.organization.name} [dim]{res.organization.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_user_memberships():
+        table.add_row("👉" if m.organization.id == org_id else "", m.organization.id, m.organization.name, m.role)
+
+    rich.print(table)
+
+
+@app.command(help="Invite a user to the organization.")
+def invite(email: str, role: OrganizationRole = "member", expires_in_days: int = 7):
+    state.settings.api.organization.invite_user(
+        InviteUser.Request(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(PortalLink.Intent.SSO)
+
+
+@app.command(help="Configure directory services for your organization.")
+def directory():
+    _do_action(PortalLink.Intent.DIRECTORY)
+
+
+@app.command(help="Configure audit logs for your organization.")
+def audit_logs():
+    _do_action(PortalLink.Intent.AUDIT_LOGS)
+
+
+@app.command(help="Configure log streams for your organization.")
+def log_streams():
+    _do_action(PortalLink.Intent.LOG_STREAMS)
+
+
+@app.command(help="Configure domains for your organization.")
+def domains():
+    _do_action(PortalLink.Intent.DOMAIN_VERIFICATION)
+
+
+def _do_action(intent: PortalLink.Intent):
+    res = state.settings.api.organization.portal_link(PortalLink.Request(intent=intent))
+    rich.print(f"Opening the configuration portal:\n{res.url}")
+    webbrowser.open(res.url)
+
+
+def current_org_id():
+    org_id = jwt.decode(state.settings.authn.token(), options={"verify_signature": False}).get("org_id")
+    if not org_id:
+        rich.print("[red]You are not logged in to an organization.[/red]")
+        raise typer.Exit(1)
+    return org_id

spiral/cli/printer.py

@@ -0,0 +1,45 @@
+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(*[getattr(msg, col, "") for col in cols])
+    return table
+
+
+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/project.py

@@ -0,0 +1,107 @@
+from typing import Annotated
+
+import rich
+import typer
+from typer import Option
+
+from spiral.api.organizations import OrganizationRole
+from spiral.api.projects import CreateProject, Grant, GrantRole, ListGrants, Project
+from spiral.cli import AsyncTyper, OptionalStr, printer, state
+from spiral.cli.org import current_org_id
+from spiral.cli.types import ProjectArg
+
+app = AsyncTyper()
+
+
+@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[
+        OptionalStr, Option(help="An optional ID prefix to which a random number will be appended.")
+    ] = None,
+    org_id: Annotated[OptionalStr, Option(help="Organization ID in which to create the project.")] = None,
+    name: Annotated[OptionalStr, Option(help="Friendly name for the project.")] = None,
+):
+    res = state.settings.api.project.create(
+        CreateProject.Request(organization_id=org_id or current_org_id(), 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[
+        OptionalStr, Option(help="Pass an organization ID to grant a role to an organization user(s).")
+    ] = None,
+    user_id: Annotated[
+        OptionalStr, 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[
+        OptionalStr,
+        Option(help="Pass an organization role when using --org-id to grant a role to all users with that role."),
+    ] = None,
+    workload_id: Annotated[OptionalStr, Option(help="Pass a workload ID to grant a role to a workload.")] = None,
+    github: Annotated[
+        OptionalStr, Option(help="Pass an `<org>/<repo>` string to grant a role to a job running in GitHub Actions.")
+    ] = None,
+    modal: Annotated[
+        OptionalStr,
+        Option(help="Pass a `<workspace_id>/<env_name>` string to grant a role to a job running in Modal environment."),
+    ] = None,
+    conditions: list[str] = Option(
+        default_factory=list,
+        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}) != 1:
+        raise typer.BadParameter("Only one of --org-id, --github or --modal may be specified.")
+
+    if github:
+        org, repo = github.split("/", 1)
+        conditions = {GrantRole.GitHubClaim(k): v for k, v in dict(c.split("=", 1) for c in conditions).items()}
+        principal = GrantRole.GitHubPrincipal(org=org, repo=repo, conditions=conditions)
+    elif modal:
+        workspace_id, environment_name = modal.split("/", 1)
+        conditions = {GrantRole.ModalClaim(k): v for k, v in dict(c.split("=", 1) for c in conditions).items()}
+        principal = GrantRole.ModalPrincipal(
+            workspace_id=workspace_id, environment_name=environment_name, conditions=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 = GrantRole.OrgUserPrincipal(org_id=org_id, user_id=user_id)
+        elif org_role is None:
+            principal = GrantRole.OrgRolePrincipal(org_id=org_id, role=OrganizationRole(org_role))
+        else:
+            raise NotImplementedError("Only user or role principal is supported at this time.")
+    elif workload_id:
+        principal = GrantRole.WorkloadPrincipal(workload_id=workload_id)
+    else:
+        raise NotImplementedError("Only organization, GitHub or Modal principal is supported at this time.")
+
+    state.settings.api.project.grant_role(
+        GrantRole.Request(
+            project_id=project,
+            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(ListGrants.Request(project_id=project)))
+    rich.print(printer.table_of_models(Grant, project_grants, title="Project Grants"))