pyspiral 0.1.0__cp310-abi3-macosx_11_0_arm64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (81) hide show
  1. pyspiral-0.1.0.dist-info/METADATA +48 -0
  2. pyspiral-0.1.0.dist-info/RECORD +81 -0
  3. pyspiral-0.1.0.dist-info/WHEEL +4 -0
  4. pyspiral-0.1.0.dist-info/entry_points.txt +2 -0
  5. spiral/__init__.py +11 -0
  6. spiral/_lib.abi3.so +0 -0
  7. spiral/adbc.py +386 -0
  8. spiral/api/__init__.py +221 -0
  9. spiral/api/admin.py +29 -0
  10. spiral/api/filesystems.py +125 -0
  11. spiral/api/organizations.py +90 -0
  12. spiral/api/projects.py +160 -0
  13. spiral/api/tables.py +94 -0
  14. spiral/api/tokens.py +56 -0
  15. spiral/api/workloads.py +45 -0
  16. spiral/arrow.py +209 -0
  17. spiral/authn/__init__.py +0 -0
  18. spiral/authn/authn.py +89 -0
  19. spiral/authn/device.py +206 -0
  20. spiral/authn/github_.py +33 -0
  21. spiral/authn/modal_.py +18 -0
  22. spiral/catalog.py +78 -0
  23. spiral/cli/__init__.py +82 -0
  24. spiral/cli/__main__.py +4 -0
  25. spiral/cli/admin.py +21 -0
  26. spiral/cli/app.py +48 -0
  27. spiral/cli/console.py +95 -0
  28. spiral/cli/fs.py +47 -0
  29. spiral/cli/login.py +13 -0
  30. spiral/cli/org.py +90 -0
  31. spiral/cli/printer.py +45 -0
  32. spiral/cli/project.py +107 -0
  33. spiral/cli/state.py +3 -0
  34. spiral/cli/table.py +20 -0
  35. spiral/cli/token.py +27 -0
  36. spiral/cli/types.py +53 -0
  37. spiral/cli/workload.py +59 -0
  38. spiral/config.py +26 -0
  39. spiral/core/__init__.py +0 -0
  40. spiral/core/core/__init__.pyi +53 -0
  41. spiral/core/manifests/__init__.pyi +53 -0
  42. spiral/core/metastore/__init__.pyi +91 -0
  43. spiral/core/spec/__init__.pyi +257 -0
  44. spiral/dataset.py +239 -0
  45. spiral/debug.py +251 -0
  46. spiral/expressions/__init__.py +222 -0
  47. spiral/expressions/base.py +149 -0
  48. spiral/expressions/http.py +86 -0
  49. spiral/expressions/io.py +100 -0
  50. spiral/expressions/list_.py +68 -0
  51. spiral/expressions/refs.py +44 -0
  52. spiral/expressions/str_.py +39 -0
  53. spiral/expressions/struct.py +57 -0
  54. spiral/expressions/tiff.py +223 -0
  55. spiral/expressions/udf.py +46 -0
  56. spiral/grpc_.py +32 -0
  57. spiral/project.py +137 -0
  58. spiral/proto/_/__init__.py +0 -0
  59. spiral/proto/_/arrow/__init__.py +0 -0
  60. spiral/proto/_/arrow/flight/__init__.py +0 -0
  61. spiral/proto/_/arrow/flight/protocol/__init__.py +0 -0
  62. spiral/proto/_/arrow/flight/protocol/sql/__init__.py +1990 -0
  63. spiral/proto/_/scandal/__init__.py +223 -0
  64. spiral/proto/_/spfs/__init__.py +36 -0
  65. spiral/proto/_/spiral/__init__.py +0 -0
  66. spiral/proto/_/spiral/table/__init__.py +225 -0
  67. spiral/proto/_/spiraldb/__init__.py +0 -0
  68. spiral/proto/_/spiraldb/metastore/__init__.py +499 -0
  69. spiral/proto/__init__.py +0 -0
  70. spiral/proto/scandal/__init__.py +45 -0
  71. spiral/proto/spiral/__init__.py +0 -0
  72. spiral/proto/spiral/table/__init__.py +96 -0
  73. spiral/proto/substrait/__init__.py +3399 -0
  74. spiral/proto/substrait/extensions/__init__.py +115 -0
  75. spiral/proto/util.py +41 -0
  76. spiral/py.typed +0 -0
  77. spiral/scan_.py +168 -0
  78. spiral/settings.py +157 -0
  79. spiral/substrait_.py +275 -0
  80. spiral/table.py +157 -0
  81. spiral/types_.py +6 -0
@@ -0,0 +1,33 @@
1
+ import os
2
+
3
+ import httpx
4
+
5
+ from spiral.api import Authn
6
+
7
+
8
+ class GitHubActionsProvider(Authn):
9
+ AUDIENCE = "https://iss.spiraldb.com"
10
+
11
+ def __init__(self):
12
+ self._gh_token = None
13
+
14
+ def token(self) -> str | None:
15
+ if self._gh_token is not None:
16
+ return self._gh_token
17
+
18
+ if os.environ.get("GITHUB_ACTIONS") == "true":
19
+ # Next, we check to see if we're running in GitHub actions and if so, grab an ID token.
20
+ if "ACTIONS_ID_TOKEN_REQUEST_TOKEN" in os.environ and "ACTIONS_ID_TOKEN_REQUEST_URL" in os.environ:
21
+ if not hasattr(self, "__gh_token"):
22
+ resp = httpx.get(
23
+ f"{os.environ['ACTIONS_ID_TOKEN_REQUEST_URL']}&audience={self.AUDIENCE}",
24
+ headers={"Authorization": f'Bearer {os.environ["ACTIONS_ID_TOKEN_REQUEST_TOKEN"]}'},
25
+ )
26
+ if not resp.is_success:
27
+ raise ValueError(f"Failed to get GitHub Actions ID token: {resp.text}", resp)
28
+ self._gh_token = resp.json()["value"]
29
+ else:
30
+ raise ValueError("Please set 'id-token: write' permission for this GitHub Actions workflow.")
31
+
32
+ # For now, we don't exchange the token for a Spiral one.
33
+ return self._gh_token
spiral/authn/modal_.py ADDED
@@ -0,0 +1,18 @@
1
+ import os
2
+
3
+ from spiral.api import Authn
4
+
5
+
6
+ class ModalProvider(Authn):
7
+ def __init__(self):
8
+ self._modal_token = None
9
+
10
+ def token(self) -> str | None:
11
+ if self._modal_token is not None:
12
+ return self._modal_token
13
+
14
+ if os.environ.get("MODAL_IDENTITY_TOKEN") is not None:
15
+ self._modal_token = os.environ["MODAL_IDENTITY_TOKEN"]
16
+
17
+ # For now, we don't exchange the token for a Spiral one.
18
+ return self._modal_token
spiral/catalog.py ADDED
@@ -0,0 +1,78 @@
1
+ from typing import TYPE_CHECKING
2
+
3
+ import jwt
4
+
5
+ from spiral.api import SpiralAPI
6
+ from spiral.api.projects import CreateProject
7
+ from spiral.settings import Settings, settings
8
+
9
+ if TYPE_CHECKING:
10
+ from spiral.project import Project
11
+ from spiral.table import Table
12
+
13
+ _default = object()
14
+
15
+
16
+ class Spiral:
17
+ def __init__(self, config: Settings | None = None):
18
+ self._config = config or settings()
19
+ self._api = self._config.api
20
+
21
+ self._org = None
22
+
23
+ @property
24
+ def config(self) -> Settings:
25
+ return self._config
26
+
27
+ @property
28
+ def api(self) -> SpiralAPI:
29
+ return self._api
30
+
31
+ @property
32
+ def organization(self) -> str:
33
+ if self._org is None:
34
+ token_payload = jwt.decode(self._config.authn.token(), options={"verify_signature": False})
35
+ if "org_id" not in token_payload:
36
+ raise ValueError("Please create an organization.")
37
+ self._org = token_payload["org_id"]
38
+ return self._org
39
+
40
+ def list_projects(self) -> list["Project"]:
41
+ """List project IDs."""
42
+ from .project import Project
43
+
44
+ return [Project(self, id=p.id, name=p.name) for p in self.api.project.list()]
45
+
46
+ def list_project_ids(self) -> list[str]:
47
+ """List project IDs."""
48
+ return [p.id for p in self.list_projects()]
49
+
50
+ def create_project(
51
+ self,
52
+ id_prefix: str | None = None,
53
+ *,
54
+ org: str | None = None,
55
+ name: str | None = None,
56
+ ) -> "Project":
57
+ """Create a project in the current, or given, organization."""
58
+ from .project import Project
59
+
60
+ org = org or self.organization
61
+ res = self.api.project.create(CreateProject.Request(organization_id=org, id_prefix=id_prefix, name=name))
62
+ return Project(self, res.project.id, name=res.project.name)
63
+
64
+ def project(self, project_id: str) -> "Project":
65
+ """Open an existing project."""
66
+ from .project import Project
67
+
68
+ # We avoid an API call since we'd just be fetching a human-readable name. Seems a waste in most cases.
69
+ return Project(self, id=project_id, name=project_id)
70
+
71
+ def table(self, identifier: str) -> "Table":
72
+ """Open a table with a "project.dataset.table" identifier."""
73
+ parts = identifier.split(".")
74
+ if len(parts) != 3:
75
+ raise ValueError(f"Invalid table identifier: {identifier}")
76
+ project_id, dataset, table = parts
77
+
78
+ return self.project(project_id).table(f"{dataset}.{table}")
spiral/cli/__init__.py ADDED
@@ -0,0 +1,82 @@
1
+ import asyncio
2
+ import functools
3
+ import inspect
4
+ from typing import IO, Optional
5
+
6
+ import rich
7
+ import typer
8
+ from click import ClickException
9
+ from grpclib import GRPCError
10
+ from httpx import HTTPStatusError
11
+
12
+ # We need to use Optional[str] since Typer doesn't support str | None.
13
+ OptionalStr = Optional[str] # noqa: UP007
14
+
15
+
16
+ class AsyncTyper(typer.Typer):
17
+ """Wrapper to allow async functions to be used as commands.
18
+
19
+ We also pre-bake some configuration.
20
+
21
+ Per https://github.com/tiangolo/typer/issues/88#issuecomment-1732469681
22
+ """
23
+
24
+ def __init__(self, **kwargs):
25
+ super().__init__(
26
+ no_args_is_help=True,
27
+ pretty_exceptions_enable=False,
28
+ **kwargs,
29
+ )
30
+
31
+ def callback(self, *args, **kwargs):
32
+ decorator = super().callback(*args, **kwargs)
33
+ for wrapper in (_wrap_exceptions, _maybe_run_async):
34
+ decorator = functools.partial(wrapper, decorator)
35
+ return decorator
36
+
37
+ def command(self, *args, **kwargs):
38
+ decorator = super().command(*args, **kwargs)
39
+ for wrapper in (_wrap_exceptions, _maybe_run_async):
40
+ decorator = functools.partial(wrapper, decorator)
41
+ return decorator
42
+
43
+
44
+ class _ClickGRPCException(ClickException):
45
+ def __init__(self, err: GRPCError):
46
+ super().__init__(err.message)
47
+ self.err = err
48
+ self.exit_code = 1
49
+
50
+ def format_message(self) -> str:
51
+ if self.err.details:
52
+ return f"{self.message}: {self.err.details}"
53
+ return self.message
54
+
55
+ def show(self, file: IO[str] | None = None) -> None:
56
+ rich.print(f"Error: {self.format_message()}", file=file)
57
+
58
+
59
+ def _maybe_run_async(decorator, f):
60
+ if inspect.iscoroutinefunction(f):
61
+
62
+ @functools.wraps(f)
63
+ def runner(*args, **kwargs):
64
+ return asyncio.run(f(*args, **kwargs))
65
+
66
+ decorator(runner)
67
+ else:
68
+ decorator(f)
69
+ return f
70
+
71
+
72
+ def _wrap_exceptions(decorator, f):
73
+ @functools.wraps(f)
74
+ def runner(*args, **kwargs):
75
+ try:
76
+ return f(*args, **kwargs)
77
+ except HTTPStatusError as e:
78
+ raise ClickException(str(e))
79
+ except GRPCError as e:
80
+ raise _ClickGRPCException(e)
81
+
82
+ return decorator(runner)
spiral/cli/__main__.py ADDED
@@ -0,0 +1,4 @@
1
+ from spiral.cli.app import main
2
+
3
+ if __name__ == "__main__":
4
+ main()
spiral/cli/admin.py ADDED
@@ -0,0 +1,21 @@
1
+ from rich import print
2
+
3
+ from spiral.api.admin import SyncMemberships, SyncOrgs
4
+ from spiral.cli import AsyncTyper, state
5
+
6
+ app = AsyncTyper()
7
+
8
+
9
+ @app.command()
10
+ def sync(orgs: bool = False, memberships: bool = False):
11
+ run_all = True
12
+ if any([orgs, memberships]):
13
+ run_all = False
14
+
15
+ if run_all or orgs:
16
+ for org_id in state.settings.api._admin.sync_orgs(SyncOrgs.Request()):
17
+ print(org_id)
18
+
19
+ if run_all or memberships:
20
+ for membership in state.settings.api._admin.sync_memberships(SyncMemberships.Request()):
21
+ print(membership)
spiral/cli/app.py ADDED
@@ -0,0 +1,48 @@
1
+ import logging
2
+ import os
3
+ from logging.handlers import RotatingFileHandler
4
+
5
+ from spiral.cli import AsyncTyper, admin, console, fs, login, org, project, state, table, token, workload
6
+ from spiral.settings import LOG_DIR, Settings
7
+
8
+ app = AsyncTyper(name="spiral")
9
+
10
+
11
+ @app.callback()
12
+ def _callback(verbose: bool = False):
13
+ if verbose:
14
+ logging.getLogger().setLevel(level=logging.INFO)
15
+
16
+ # Load the settings (we reload in the callback to support testing under different env vars)
17
+ state.settings = Settings()
18
+
19
+
20
+ app.add_typer(fs.app, name="fs")
21
+ app.add_typer(org.app, name="org")
22
+ app.add_typer(project.app, name="project")
23
+ app.add_typer(table.app, name="table")
24
+ app.add_typer(workload.app, name="workload")
25
+ app.add_typer(token.app, name="token")
26
+ app.command("console")(console.command)
27
+ app.command("login")(login.command)
28
+
29
+ # Register unless we're building docs. Because Typer docs command does not skip hidden commands...
30
+ if not bool(os.environ.get("SPIRAL_DOCS", False)):
31
+ app.add_typer(admin.app, name="admin", hidden=True)
32
+ app.command("logout", hidden=True)(login.logout)
33
+
34
+
35
+ def main():
36
+ # Setup rotating CLI logging.
37
+ # NOTE(ngates): we should do the same for the Spiral client? Maybe move this logic elsewhere?
38
+ LOG_DIR.mkdir(parents=True, exist_ok=True)
39
+ logging.basicConfig(
40
+ level=logging.DEBUG,
41
+ handlers=[RotatingFileHandler(LOG_DIR / "cli.log", maxBytes=2**20, backupCount=10)],
42
+ )
43
+
44
+ app()
45
+
46
+
47
+ if __name__ == "__main__":
48
+ main()
spiral/cli/console.py ADDED
@@ -0,0 +1,95 @@
1
+ import os
2
+ import subprocess
3
+
4
+ from spiral import Spiral
5
+ from spiral.adbc import ADBCFlightServer, SpiralADBCServer
6
+ from spiral.api import wait_for_port
7
+
8
+
9
+ def command():
10
+ """Launch a SQL console to query Spiral tables."""
11
+
12
+ # To avoid taking a dependency on Harlequin, we install it on-demand using
13
+ # either uvx or pipx.
14
+ harlequin_args = _uvx()
15
+ if harlequin_args is None:
16
+ harlequin_args = _pipx()
17
+ if harlequin_args is None:
18
+ raise ValueError("Please install pipx to continue\n\tSee https://github.com/pypa/pipx")
19
+
20
+ # Set up a pipe to send the server port to the child process.
21
+ r, w = os.pipe()
22
+
23
+ pid = os.fork()
24
+ if pid == 0: # In the child
25
+ os.close(w)
26
+ port = int.from_bytes(os.read(r, 4), "big")
27
+
28
+ # Wait for the server to be up.
29
+ wait_for_port(port)
30
+
31
+ os.execv(
32
+ harlequin_args[0],
33
+ harlequin_args
34
+ + [
35
+ "-a",
36
+ "adbc",
37
+ "--driver-type",
38
+ "flightsql",
39
+ f"grpc://localhost:{port}",
40
+ ],
41
+ )
42
+ else:
43
+ os.close(r)
44
+
45
+ # I can't get the Flight server to stop writing to stdout. So we need to spawn a new process I think and
46
+ # then hope we can kill it?
47
+ fd = os.open("/dev/null", os.O_WRONLY)
48
+ os.dup2(fd, 1)
49
+ os.dup2(fd, 2)
50
+
51
+ # In the parent, we launch the Flight SQL server and send the port to the child
52
+ server = ADBCFlightServer(SpiralADBCServer(Spiral()))
53
+ os.write(w, server.port.to_bytes(4, "big"))
54
+
55
+ # Then wait for the console app to exit
56
+ os.waitpid(pid, 0)
57
+
58
+
59
+ def _pipx() -> list[str] | None:
60
+ """Run harlequin via pipx."""
61
+ res = subprocess.run(["which", "pipx"], stdout=subprocess.PIPE)
62
+ if res.returncode != 0:
63
+ return None
64
+ # raise ValueError("Please install pipx to continue\n\tSee https://github.com/pypa/pipx")
65
+ pipx = res.stdout.strip()
66
+
67
+ return [
68
+ pipx,
69
+ "run",
70
+ "--pip-args",
71
+ "adbc_driver_flightsql",
72
+ "--pip-args",
73
+ # for now, we pin rich
74
+ "rich<=13.9.1",
75
+ "harlequin[adbc]",
76
+ ]
77
+
78
+
79
+ def _uvx() -> list[str] | None:
80
+ """Run harlequin via uvx."""
81
+ res = subprocess.run(["which", "uvx"], stdout=subprocess.PIPE)
82
+ if res.returncode != 0:
83
+ return None
84
+ uvx = res.stdout.strip()
85
+
86
+ return [
87
+ uvx,
88
+ "--with",
89
+ "adbc_driver_flightsql",
90
+ "--with",
91
+ "rich<=13.9.1",
92
+ "--from",
93
+ "harlequin[adbc]",
94
+ "harlequin",
95
+ ]
spiral/cli/fs.py ADDED
@@ -0,0 +1,47 @@
1
+ from typing import Annotated
2
+
3
+ import questionary
4
+ import rich
5
+ from typer import Option
6
+
7
+ from spiral.api.filesystems import BuiltinFileSystem, GetFileSystem, UpdateFileSystem
8
+ from spiral.cli import AsyncTyper, state
9
+ from spiral.cli.types import ProjectArg
10
+
11
+ app = AsyncTyper()
12
+
13
+
14
+ @app.command(help="Show the file system configured for project.")
15
+ def show(project: ProjectArg):
16
+ res = state.settings.api.file_system.get_file_system(GetFileSystem.Request(project_id=project))
17
+ match res.file_system:
18
+ case BuiltinFileSystem(provider=provider):
19
+ rich.print(f"provider: {provider}")
20
+ case _:
21
+ rich.print(res.file_system)
22
+
23
+
24
+ def _provider_default():
25
+ res = state.settings.api.file_system.list_providers()
26
+ questionary.select("Select a file system provider", choices=res.providers).ask()
27
+
28
+
29
+ ProviderOpt = Annotated[
30
+ str,
31
+ Option(help="Built-in provider to use for the file system.", show_default=False, default_factory=_provider_default),
32
+ ]
33
+
34
+
35
+ @app.command(help="Update a project's file system.")
36
+ def update(project: ProjectArg, provider: ProviderOpt):
37
+ res = state.settings.api.file_system.update_file_system(
38
+ UpdateFileSystem.Request(project_id=project, file_system=BuiltinFileSystem(provider=provider))
39
+ )
40
+ rich.print(res.file_system.builtin)
41
+
42
+
43
+ @app.command(help="Lists the available built-in file system providers.")
44
+ def list_providers():
45
+ res = state.settings.api.file_system.list_providers()
46
+ for provider in res.providers:
47
+ rich.print(provider)
spiral/cli/login.py ADDED
@@ -0,0 +1,13 @@
1
+ from rich import print
2
+
3
+ from spiral.cli import OptionalStr, state
4
+
5
+
6
+ def command(org_id: OptionalStr = None, force: bool = False, refresh: bool = False):
7
+ tokens = state.settings.spiraldb.device_auth().authenticate(force=force, refresh=refresh, organization_id=org_id)
8
+ print(tokens)
9
+
10
+
11
+ def logout():
12
+ state.settings.spiraldb.device_auth().logout()
13
+ print("Logged out.")
spiral/cli/org.py ADDED
@@ -0,0 +1,90 @@
1
+ import webbrowser
2
+ from typing import Annotated
3
+
4
+ import jwt
5
+ import rich
6
+ import typer
7
+ from rich.table import Table
8
+ from typer import Option
9
+
10
+ from spiral.api.organizations import CreateOrganization, InviteUser, OrganizationRole, PortalLink
11
+ from spiral.cli import AsyncTyper, OptionalStr, state
12
+ from spiral.cli.types import OrganizationArg
13
+
14
+ app = AsyncTyper()
15
+
16
+
17
+ @app.command(help="Switch the active organization.")
18
+ def switch(org_id: OrganizationArg):
19
+ state.settings.spiraldb.device_auth().authenticate(refresh=True, organization_id=org_id)
20
+ rich.print(f"Switched to organization: {org_id}")
21
+
22
+
23
+ @app.command(help="Create a new organization.")
24
+ def create(
25
+ name: Annotated[OptionalStr, Option(help="The human-readable name of the organization.")] = None,
26
+ ):
27
+ res = state.settings.api.organization.create_organization(CreateOrganization.Request(name=name))
28
+
29
+ # Authenticate to the new organization
30
+ state.settings.spiraldb.device_auth().authenticate(refresh=True, organization_id=res.organization.id)
31
+
32
+ rich.print(f"{res.organization.name} [dim]{res.organization.id}[/dim]")
33
+
34
+
35
+ @app.command(help="List organizations.")
36
+ def ls():
37
+ org_id = current_org_id()
38
+
39
+ table = Table("", "id", "name", "role", title="Organizations")
40
+ for m in state.settings.api.organization.list_user_memberships():
41
+ table.add_row("👉" if m.organization.id == org_id else "", m.organization.id, m.organization.name, m.role)
42
+
43
+ rich.print(table)
44
+
45
+
46
+ @app.command(help="Invite a user to the organization.")
47
+ def invite(email: str, role: OrganizationRole = "member", expires_in_days: int = 7):
48
+ state.settings.api.organization.invite_user(
49
+ InviteUser.Request(email=email, role=role, expires_in_days=expires_in_days)
50
+ )
51
+ rich.print(f"Invited {email} as a {role.value}.")
52
+
53
+
54
+ @app.command(help="Configure single sign-on for your organization.")
55
+ def sso():
56
+ _do_action(PortalLink.Intent.SSO)
57
+
58
+
59
+ @app.command(help="Configure directory services for your organization.")
60
+ def directory():
61
+ _do_action(PortalLink.Intent.DIRECTORY)
62
+
63
+
64
+ @app.command(help="Configure audit logs for your organization.")
65
+ def audit_logs():
66
+ _do_action(PortalLink.Intent.AUDIT_LOGS)
67
+
68
+
69
+ @app.command(help="Configure log streams for your organization.")
70
+ def log_streams():
71
+ _do_action(PortalLink.Intent.LOG_STREAMS)
72
+
73
+
74
+ @app.command(help="Configure domains for your organization.")
75
+ def domains():
76
+ _do_action(PortalLink.Intent.DOMAIN_VERIFICATION)
77
+
78
+
79
+ def _do_action(intent: PortalLink.Intent):
80
+ res = state.settings.api.organization.portal_link(PortalLink.Request(intent=intent))
81
+ rich.print(f"Opening the configuration portal:\n{res.url}")
82
+ webbrowser.open(res.url)
83
+
84
+
85
+ def current_org_id():
86
+ org_id = jwt.decode(state.settings.authn.token(), options={"verify_signature": False}).get("org_id")
87
+ if not org_id:
88
+ rich.print("[red]You are not logged in to an organization.[/red]")
89
+ raise typer.Exit(1)
90
+ return org_id
spiral/cli/printer.py ADDED
@@ -0,0 +1,45 @@
1
+ from collections.abc import Iterable
2
+ from typing import TypeVar
3
+
4
+ import betterproto
5
+ from pydantic import BaseModel
6
+ from rich.console import ConsoleRenderable, Group
7
+ from rich.padding import Padding
8
+ from rich.pretty import Pretty
9
+ from rich.table import Table
10
+
11
+ T = TypeVar("T", bound=betterproto.Message)
12
+ M = TypeVar("M", bound=BaseModel)
13
+
14
+
15
+ def table_of_models(cls: type[M], messages: Iterable[M], fields: list[str] = None, title: str = None) -> Table:
16
+ """Centralized logic for printing tables of Pydantic models."""
17
+ cols = fields or cls.model_fields.keys()
18
+ table = Table(*cols, title=title or f"{cls.__name__}s")
19
+ for msg in messages:
20
+ table.add_row(*[getattr(msg, col, "") for col in cols])
21
+ return table
22
+
23
+
24
+ def table_of_protos(cls: type[T], messages: Iterable[T], fields: list[str] = None, title: str = None) -> Table:
25
+ """Centralized logic for printing tables of proto messages.
26
+
27
+ TODO(ngates): add a CLI switch to emit JSON results instead of tables.
28
+ """
29
+ cols = fields or cls()._betterproto.sorted_field_names
30
+ table = Table(*cols, title=title or f"{cls.__name__}s")
31
+ for msg in messages:
32
+ table.add_row(*[getattr(msg, col, "") for col in cols])
33
+ return table
34
+
35
+
36
+ def proto(message: T, title: str = None, fields: list[str] = None) -> ConsoleRenderable:
37
+ """Centralized logic for printing a single proto message."""
38
+ value = Pretty({k: v for k, v in message.to_dict().items() if not fields or k in fields})
39
+ if title:
40
+ return Group(
41
+ f"[bold]{title}[/bold]",
42
+ Padding.indent(value, level=2),
43
+ )
44
+ else:
45
+ return value
spiral/cli/project.py ADDED
@@ -0,0 +1,107 @@
1
+ from typing import Annotated
2
+
3
+ import rich
4
+ import typer
5
+ from typer import Option
6
+
7
+ from spiral.api.organizations import OrganizationRole
8
+ from spiral.api.projects import CreateProject, Grant, GrantRole, ListGrants, Project
9
+ from spiral.cli import AsyncTyper, OptionalStr, printer, state
10
+ from spiral.cli.org import current_org_id
11
+ from spiral.cli.types import ProjectArg
12
+
13
+ app = AsyncTyper()
14
+
15
+
16
+ @app.command(help="List projects.")
17
+ def ls():
18
+ projects = list(state.settings.api.project.list())
19
+ rich.print(printer.table_of_models(Project, projects))
20
+
21
+
22
+ @app.command(help="Create a new project.")
23
+ def create(
24
+ id_prefix: Annotated[
25
+ OptionalStr, Option(help="An optional ID prefix to which a random number will be appended.")
26
+ ] = None,
27
+ org_id: Annotated[OptionalStr, Option(help="Organization ID in which to create the project.")] = None,
28
+ name: Annotated[OptionalStr, Option(help="Friendly name for the project.")] = None,
29
+ ):
30
+ res = state.settings.api.project.create(
31
+ CreateProject.Request(organization_id=org_id or current_org_id(), id_prefix=id_prefix, name=name)
32
+ )
33
+ rich.print(f"Created project {res.project.id}")
34
+
35
+
36
+ @app.command(help="Grant a role on a project.")
37
+ def grant(
38
+ project: ProjectArg,
39
+ role: Annotated[str, Option(help="Role to grant.")],
40
+ org_id: Annotated[
41
+ OptionalStr, Option(help="Pass an organization ID to grant a role to an organization user(s).")
42
+ ] = None,
43
+ user_id: Annotated[
44
+ OptionalStr, Option(help="Pass a user ID when using --org-id to grant a role to grant a role to a user.")
45
+ ] = None,
46
+ org_role: Annotated[
47
+ OptionalStr,
48
+ Option(help="Pass an organization role when using --org-id to grant a role to all users with that role."),
49
+ ] = None,
50
+ workload_id: Annotated[OptionalStr, Option(help="Pass a workload ID to grant a role to a workload.")] = None,
51
+ github: Annotated[
52
+ OptionalStr, Option(help="Pass an `<org>/<repo>` string to grant a role to a job running in GitHub Actions.")
53
+ ] = None,
54
+ modal: Annotated[
55
+ OptionalStr,
56
+ Option(help="Pass a `<workspace_id>/<env_name>` string to grant a role to a job running in Modal environment."),
57
+ ] = None,
58
+ conditions: list[str] = Option(
59
+ default_factory=list,
60
+ help="`<key>=<value>` token conditions to apply to the grant when using --github or --modal.",
61
+ ),
62
+ ):
63
+ # Check mutual exclusion
64
+ if sum(int(bool(opt)) for opt in {org_id, workload_id, github, modal}) != 1:
65
+ raise typer.BadParameter("Only one of --org-id, --github or --modal may be specified.")
66
+
67
+ if github:
68
+ org, repo = github.split("/", 1)
69
+ conditions = {GrantRole.GitHubClaim(k): v for k, v in dict(c.split("=", 1) for c in conditions).items()}
70
+ principal = GrantRole.GitHubPrincipal(org=org, repo=repo, conditions=conditions)
71
+ elif modal:
72
+ workspace_id, environment_name = modal.split("/", 1)
73
+ conditions = {GrantRole.ModalClaim(k): v for k, v in dict(c.split("=", 1) for c in conditions).items()}
74
+ principal = GrantRole.ModalPrincipal(
75
+ workspace_id=workspace_id, environment_name=environment_name, conditions=conditions
76
+ )
77
+ elif org_id:
78
+ # Check mutual exclusion
79
+ if sum(int(bool(opt)) for opt in {user_id, org_role}) != 1:
80
+ raise typer.BadParameter("Only one of --user-id or --org-role may be specified.")
81
+
82
+ if user_id is not None:
83
+ principal = GrantRole.OrgUserPrincipal(org_id=org_id, user_id=user_id)
84
+ elif org_role is None:
85
+ principal = GrantRole.OrgRolePrincipal(org_id=org_id, role=OrganizationRole(org_role))
86
+ else:
87
+ raise NotImplementedError("Only user or role principal is supported at this time.")
88
+ elif workload_id:
89
+ principal = GrantRole.WorkloadPrincipal(workload_id=workload_id)
90
+ else:
91
+ raise NotImplementedError("Only organization, GitHub or Modal principal is supported at this time.")
92
+
93
+ state.settings.api.project.grant_role(
94
+ GrantRole.Request(
95
+ project_id=project,
96
+ role_id=role,
97
+ principal=principal,
98
+ )
99
+ )
100
+
101
+ rich.print(f"Granted role {role} on project {project}")
102
+
103
+
104
+ @app.command(help="List project grants.")
105
+ def grants(project: ProjectArg):
106
+ project_grants = list(state.settings.api.project.list_grants(ListGrants.Request(project_id=project)))
107
+ rich.print(printer.table_of_models(Grant, project_grants, title="Project Grants"))