feat: Add a 'create secrets' command in the CLI

Author: Robin-Van-de-Merghel

diracx-cli/pyproject.toml

@@ -39,6 +39,7 @@ dirac = "diracx.cli:app"
 [project.entry-points."diracx.cli"]
 jobs = "diracx.cli.jobs:app"
 config = "diracx.cli.config:app"
+pilots = "diracx.cli.pilots:app"
 
 [project.entry-points."diracx.cli.hidden"]
 internal = "diracx.cli.internal:app"

diracx-cli/src/diracx/cli/pilots.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+__all__ = ("app",)
+
+import asyncio
+import json
+from typing import Annotated, Optional
+
+import typer
+
+from diracx.client.aio import AsyncDiracClient
+
+from .utils import AsyncTyper
+
+app = AsyncTyper()
+
+
+async def installation_metadata():
+    async with AsyncDiracClient() as api:
+        return await api.well_known.get_installation_metadata()
+
+
+def vo_callback(vo: str | None) -> str:
+    metadata = asyncio.run(installation_metadata())
+    vos = list(metadata.virtual_organizations)
+    if not vo:
+        raise typer.BadParameter(
+            f"VO must be specified, available options are: {' '.join(vos)}"
+        )
+    if vo not in vos:
+        raise typer.BadParameter(
+            f"Unknown VO {vo}, available options are: {' '.join(vos)}"
+        )
+    return vo
+
+
+@app.async_command()
+async def generate_pilot_secrets(
+    vo: Annotated[
+        str,
+        typer.Argument(callback=vo_callback, help="Virtual Organization name"),
+    ],
+    n: Annotated[
+        int,
+        typer.Argument(help="Number of secrets to generate."),
+    ],
+    expiration_minutes: Optional[int] = typer.Option(
+        60,
+        help="Expiration in minutes of the secrets.",
+    ),
+    max_use: Optional[int] = typer.Option(
+        60,
+        help="Number of uses max for a secret.",
+    ),
+):
+    async with AsyncDiracClient() as api:
+        secrets = await api.pilots.create_pilot_secrets(
+            n=n,
+            expiration_minutes=expiration_minutes,
+            pilot_secret_use_count_max=max_use,
+            vo=vo,
+        )
+        # Convert each model to dict
+        secrets_dict = [secret.as_dict() for secret in secrets]
+
+        print(json.dumps(secrets_dict, indent=2))

diracx-client/src/diracx/client/patches/pilots/aio.py

@@ -16,16 +16,18 @@
 from azure.core.tracing.decorator_async import distributed_trace_async
 
 from ..._generated.aio.operations._operations import PilotsOperations as _PilotsOperations
-from ..._generated.models._models import PilotCredentialsInfo
+from ..._generated.models._models import PilotCredentialsInfo, PilotSecretsInfo
 from .common import (
     make_search_body,
     make_summary_body,
     make_add_pilot_stamps_body,
     make_update_pilot_fields_body,
+    make_create_pilot_secrets_body,
     SearchKwargs,
     SummaryKwargs,
     AddPilotStampsKwargs,
-    UpdatePilotFieldsKwargs
+    UpdatePilotFieldsKwargs,
+    CreatePilotSecretsKwargs,
 )
 
 # We're intentionally ignoring overrides here because we want to change the interface.
@@ -52,3 +54,8 @@ async def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> list
     async def update_pilot_fields(self, **kwargs: Unpack[UpdatePilotFieldsKwargs]) -> None:
         """TODO"""
         return await super().update_pilot_fields(**make_update_pilot_fields_body(**kwargs))
+
+    @distributed_trace_async
+    async def create_pilot_secrets(self, **kwargs: Unpack[CreatePilotSecretsKwargs]) -> list[PilotSecretsInfo]:
+        """TODO"""
+        return await super().create_pilot_secrets(**make_create_pilot_secrets_body(**kwargs))

diracx-client/src/diracx/client/patches/pilots/common.py

@@ -11,6 +11,8 @@
     "make_add_pilot_stamps_body",
     "UpdatePilotFieldsKwargs",
     "make_update_pilot_fields_body"
+    "CreatePilotSecretsKwargs",
+    "make_create_pilot_secrets_body"
 ]
 
 import json
@@ -146,3 +148,31 @@ def make_update_pilot_fields_body(**kwargs: Unpack[UpdatePilotFieldsKwargs]) ->
     result: UnderlyingUpdatePilotFields = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
     result.update(cast(ResponseExtra, kwargs))
     return result
+
+# ------------------ CreatePilotSecrets ------------------
+
+class CreatePilotSecretsBody(TypedDict, total=False):
+    n: int
+    expiration_minutes: int | None
+    pilot_secret_use_count_max: int | None
+    vo: str
+
+class CreatePilotSecretsKwargs(CreatePilotSecretsBody, ResponseExtra): ...
+
+class UnderlyingCreatePilotSecrets(ResponseExtra, total=False):
+    # FIXME: The autorest-generated has a bug that it expected IO[bytes] despite
+    # the code being generated to support IO[bytes] | bytes.
+    body: IO[bytes]
+
+def make_create_pilot_secrets_body(**kwargs: Unpack[CreatePilotSecretsKwargs]) -> UnderlyingCreatePilotSecrets:
+    body: CreatePilotSecretsBody = {}
+    for key in CreatePilotSecretsBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["n", "expiration_minutes", "pilot_secret_use_count_max", "vo"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingCreatePilotSecrets = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(ResponseExtra, kwargs))
+    return result

diracx-client/src/diracx/client/patches/pilots/sync.py

@@ -16,16 +16,18 @@
 from azure.core.tracing.decorator import distributed_trace
 
 from ..._generated.operations._operations import PilotsOperations as _PilotsOperations
-from ..._generated.models._models import PilotCredentialsInfo
+from ..._generated.models._models import PilotCredentialsInfo, PilotSecretsInfo
 from .common import (
     make_search_body,
     make_summary_body,
     make_add_pilot_stamps_body,
     make_update_pilot_fields_body,
+    make_create_pilot_secrets_body,
     SearchKwargs,
     SummaryKwargs,
     AddPilotStampsKwargs,
-    UpdatePilotFieldsKwargs
+    UpdatePilotFieldsKwargs,
+    CreatePilotSecretsKwargs,
 )
 
 # We're intentionally ignoring overrides here because we want to change the interface.
@@ -52,3 +54,8 @@ def add_pilot_stamps(self, **kwargs: Unpack[AddPilotStampsKwargs]) -> list[Pilot
     def update_pilot_fields(self, **kwargs: Unpack[UpdatePilotFieldsKwargs]) -> None:
         """TODO"""
         return super().update_pilot_fields(**make_update_pilot_fields_body(**kwargs))
+
+    @distributed_trace
+    def create_pilot_secrets(self, **kwargs: Unpack[CreatePilotSecretsKwargs]) -> list[PilotSecretsInfo]:
+        """TODO"""
+        return super().create_pilot_secrets(**make_create_pilot_secrets_body(**kwargs))