Skip to content

Commit

Permalink
feat(diracx): move bl from __init__ files
Browse files Browse the repository at this point in the history
  • Loading branch information
aldbr authored and chrisburr committed Dec 20, 2024
1 parent 15a1926 commit d0402ec
Show file tree
Hide file tree
Showing 62 changed files with 2,518 additions and 2,311 deletions.
2 changes: 2 additions & 0 deletions .github/workflows/make_release.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
#!/usr/bin/env python
from __future__ import annotations

import argparse
from pathlib import Path

Expand Down
2 changes: 1 addition & 1 deletion diracx-cli/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ dependencies = [
"gitpython",
"pydantic>=2.10",
"rich",
"typer",
"typer>=0.12.4",
"pyyaml",
]
dynamic = ["version"]
Expand Down
137 changes: 3 additions & 134 deletions diracx-cli/src/diracx/cli/__init__.py
Original file line number Diff line number Diff line change
@@ -1,144 +1,13 @@
import asyncio
import json
import os
from datetime import datetime, timedelta, timezone
from typing import Annotated, Optional
from __future__ import annotations

import typer

from diracx.client.aio import DiracClient
from diracx.client.models import DeviceFlowErrorResponse
from diracx.core.extensions import select_from_extension
from diracx.core.preferences import get_diracx_preferences
from diracx.core.utils import read_credentials, write_credentials

from .utils import AsyncTyper

app = AsyncTyper()


async def installation_metadata():
async with DiracClient() as api:
return await api.well_known.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 login(
vo: Annotated[
Optional[str],
typer.Argument(callback=vo_callback, help="Virtual Organization name"),
] = None,
group: Optional[str] = typer.Option(
None,
help="Group name within the VO. If not provided, the default group for the VO will be used.",
),
property: Optional[list[str]] = typer.Option(
None,
help=(
"List of properties to add to the default properties of the group. "
"If not provided, default properties of the group will be used."
),
),
):
"""Login to the DIRAC system using the device flow.
- If only VO is provided: Uses the default group and its properties for the VO.
- If VO and group are provided: Uses the specified group and its properties for the VO.

- If VO and properties are provided: Uses the default group and combines its properties with the
provided properties.
from .auth import app

- If VO, group, and properties are provided: Uses the specified group and combines its properties with the
provided properties.
"""
scopes = [f"vo:{vo}"]
if group:
scopes.append(f"group:{group}")
if property:
scopes += [f"property:{p}" for p in property]

print(f"Logging in with scopes: {scopes}")
async with DiracClient() as api:
data = await api.auth.initiate_device_flow(
client_id=api.client_id,
scope=" ".join(scopes),
)
print("Now go to:", data.verification_uri_complete)
expires = datetime.now(tz=timezone.utc) + timedelta(
seconds=data.expires_in - 30
)
while expires > datetime.now(tz=timezone.utc):
print(".", end="", flush=True)
response = await api.auth.token(device_code=data.device_code, client_id=api.client_id) # type: ignore
if isinstance(response, DeviceFlowErrorResponse):
if response.error == "authorization_pending":
# TODO: Setting more than 5 seconds results in an error
# Related to keep-alive disconnects from uvicon (--timeout-keep-alive)
await asyncio.sleep(2)
continue
raise RuntimeError(f"Device flow failed with {response}")
break
else:
raise RuntimeError("Device authorization flow expired")

# Save credentials
write_credentials(response)
credentials_path = get_diracx_preferences().credentials_path
print(f"Saved credentials to {credentials_path}")
print("\nLogin successful!")


@app.async_command()
async def whoami():
async with DiracClient() as api:
user_info = await api.auth.userinfo()
# TODO: Add a RICH output format
print(json.dumps(user_info.as_dict(), indent=2))


@app.async_command()
async def logout():
async with DiracClient() as api:
credentials_path = get_diracx_preferences().credentials_path
if credentials_path.exists():
credentials = read_credentials(credentials_path)

# Revoke refresh token
try:
await api.auth.revoke_refresh_token(credentials.refresh_token)
except Exception as e:
print(f"Error revoking the refresh token {e!r}")
pass

# Remove credentials
credentials_path.unlink(missing_ok=True)
print(f"Removed credentials from {credentials_path}")
print("\nLogout successful!")


@app.callback()
def callback(output_format: Optional[str] = None):
if output_format is not None:
os.environ["DIRACX_OUTPUT_FORMAT"] = output_format
__all__ = ("app",)


# Load all the sub commands

cli_names = set(
[entry_point.name for entry_point in select_from_extension(group="diracx.cli")]
)
Expand Down
2 changes: 2 additions & 0 deletions diracx-cli/src/diracx/cli/__main__.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from __future__ import annotations

Check warning on line 1 in diracx-cli/src/diracx/cli/__main__.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/__main__.py#L1

Added line #L1 was not covered by tests

from . import app

if __name__ == "__main__":
Expand Down
142 changes: 142 additions & 0 deletions diracx-cli/src/diracx/cli/auth.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,142 @@
# Can't using PEP-604 with typer: https://github.com/tiangolo/typer/issues/348
# from __future__ import annotations
from __future__ import annotations

__all__ = ("app",)

import asyncio
import json
import os
from datetime import datetime, timedelta, timezone
from typing import Annotated, Optional

import typer

from diracx.client.aio import DiracClient
from diracx.client.models import DeviceFlowErrorResponse
from diracx.core.preferences import get_diracx_preferences
from diracx.core.utils import read_credentials, write_credentials

from .utils import AsyncTyper

app = AsyncTyper()


async def installation_metadata():
async with DiracClient() as api:
return await api.well_known.installation_metadata()

Check warning on line 27 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L26-L27

Added lines #L26 - L27 were not covered by tests


def vo_callback(vo: str | None) -> str:
metadata = asyncio.run(installation_metadata())
vos = list(metadata.virtual_organizations)
if not vo:
raise typer.BadParameter(

Check warning on line 34 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L31-L34

Added lines #L31 - L34 were not covered by tests
f"VO must be specified, available options are: {' '.join(vos)}"
)
if vo not in vos:
raise typer.BadParameter(

Check warning on line 38 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L37-L38

Added lines #L37 - L38 were not covered by tests
f"Unknown VO {vo}, available options are: {' '.join(vos)}"
)
return vo

Check warning on line 41 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L41

Added line #L41 was not covered by tests


@app.async_command()
async def login(
vo: Annotated[
Optional[str],
typer.Argument(callback=vo_callback, help="Virtual Organization name"),
] = None,
group: Optional[str] = typer.Option(
None,
help="Group name within the VO. If not provided, the default group for the VO will be used.",
),
property: Optional[list[str]] = typer.Option(
None,
help=(
"List of properties to add to the default properties of the group. "
"If not provided, default properties of the group will be used."
),
),
):
"""Login to the DIRAC system using the device flow.
- If only VO is provided: Uses the default group and its properties for the VO.
- If VO and group are provided: Uses the specified group and its properties for the VO.
- If VO and properties are provided: Uses the default group and combines its properties with the
provided properties.
- If VO, group, and properties are provided: Uses the specified group and combines its properties with the
provided properties.
"""
scopes = [f"vo:{vo}"]
if group:
scopes.append(f"group:{group}")
if property:
scopes += [f"property:{p}" for p in property]

Check warning on line 78 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L74-L78

Added lines #L74 - L78 were not covered by tests

print(f"Logging in with scopes: {scopes}")
async with DiracClient() as api:
data = await api.auth.initiate_device_flow(

Check warning on line 82 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L80-L82

Added lines #L80 - L82 were not covered by tests
client_id=api.client_id,
scope=" ".join(scopes),
)
print("Now go to:", data.verification_uri_complete)
expires = datetime.now(tz=timezone.utc) + timedelta(

Check warning on line 87 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L86-L87

Added lines #L86 - L87 were not covered by tests
seconds=data.expires_in - 30
)
while expires > datetime.now(tz=timezone.utc):
print(".", end="", flush=True)
response = await api.auth.token(device_code=data.device_code, client_id=api.client_id) # type: ignore
if isinstance(response, DeviceFlowErrorResponse):
if response.error == "authorization_pending":

Check warning on line 94 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L90-L94

Added lines #L90 - L94 were not covered by tests
# TODO: Setting more than 5 seconds results in an error
# Related to keep-alive disconnects from uvicon (--timeout-keep-alive)
await asyncio.sleep(2)
continue
raise RuntimeError(f"Device flow failed with {response}")
break

Check warning on line 100 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L97-L100

Added lines #L97 - L100 were not covered by tests
else:
raise RuntimeError("Device authorization flow expired")

Check warning on line 102 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L102

Added line #L102 was not covered by tests

# Save credentials
write_credentials(response)
credentials_path = get_diracx_preferences().credentials_path
print(f"Saved credentials to {credentials_path}")
print("\nLogin successful!")

Check warning on line 108 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L105-L108

Added lines #L105 - L108 were not covered by tests


@app.async_command()
async def whoami():
async with DiracClient() as api:
user_info = await api.auth.userinfo()

Check warning on line 114 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L113-L114

Added lines #L113 - L114 were not covered by tests
# TODO: Add a RICH output format
print(json.dumps(user_info.as_dict(), indent=2))

Check warning on line 116 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L116

Added line #L116 was not covered by tests


@app.async_command()
async def logout():
async with DiracClient() as api:
credentials_path = get_diracx_preferences().credentials_path
if credentials_path.exists():
credentials = read_credentials(credentials_path)

Check warning on line 124 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L121-L124

Added lines #L121 - L124 were not covered by tests

# Revoke refresh token
try:
await api.auth.revoke_refresh_token(credentials.refresh_token)
except Exception as e:
print(f"Error revoking the refresh token {e!r}")
pass

Check warning on line 131 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L127-L131

Added lines #L127 - L131 were not covered by tests

# Remove credentials
credentials_path.unlink(missing_ok=True)
print(f"Removed credentials from {credentials_path}")
print("\nLogout successful!")

Check warning on line 136 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L134-L136

Added lines #L134 - L136 were not covered by tests


@app.callback()
def callback(output_format: Optional[str] = None):
if output_format is not None:
os.environ["DIRACX_OUTPUT_FORMAT"] = output_format

Check warning on line 142 in diracx-cli/src/diracx/cli/auth.py

View check run for this annotation

Codecov / codecov/patch

diracx-cli/src/diracx/cli/auth.py#L142

Added line #L142 was not covered by tests
1 change: 1 addition & 0 deletions diracx-cli/src/diracx/cli/config.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
# Can't using PEP-604 with typer: https://github.com/tiangolo/typer/issues/348
# from __future__ import annotations
from __future__ import annotations

__all__ = ("dump",)

Expand Down
Loading

0 comments on commit d0402ec

Please sign in to comment.