feat: integration authz headers in Dirac client

.github/workflows/integration.yml

@@ -31,8 +31,6 @@ jobs:
         run: |
           pip install typer pyyaml gitpython packaging
           git clone https://github.com/DIRACGrid/DIRAC.git -b "${{ matrix.dirac-branch }}" /tmp/DIRACRepo
-          # HACK: Workaround the issues which will be fixed in https://github.com/DIRACGrid/diracx/pull/87
-          git --git-dir /tmp/DIRACRepo/.git checkout 208bb7983c3530fc0aa125df6011fc32c3f482fd
           # We need to cd in the directory for the integration_tests.py to work
       - name: Prepare environment
         run: cd /tmp/DIRACRepo && ./integration_tests.py prepare-environment "TEST_DIRACX=Yes"  --extra-module "diracx=${GITHUB_WORKSPACE}"

src/diracx/cli/__init__.py

@@ -3,21 +3,21 @@
 import asyncio
 import json
 import os
-from datetime import datetime, timedelta, timezone
+from datetime import datetime, timedelta
 from typing import Optional
 
 from typer import Option
 
-from diracx.client.aio import Dirac
+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 write_credentials
 
 from . import internal, jobs
-from .utils import CREDENTIALS_PATH, AsyncTyper
+from .utils import AsyncTyper
 
 app = AsyncTyper()
 
-EXPIRES_GRACE_SECONDS = 15
-
 
 @app.async_command()
 async def login(
@@ -34,50 +34,52 @@
         scopes += [f"property:{p}" for p in property]
 
     print(f"Logging in with scopes: {scopes}")
-    # TODO set endpoint URL from preferences
-    async with Dirac(endpoint="http://localhost:8000") as api:
+    async with DiracClient() as api:
         data = await api.auth.initiate_device_flow(
-            client_id="myDIRACClientID",
+            client_id=api.client_id,
             audience="Dirac server",
             scope=" ".join(scopes),
         )
         print("Now go to:", data.verification_uri_complete)
         expires = datetime.now() + timedelta(seconds=data.expires_in - 30)
         while expires > datetime.now():
             print(".", end="", flush=True)
-            response = await api.auth.token(  # type: ignore
-                vo, device_code=data.device_code, client_id="myDIRACClientID"
-            )
+            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}")
-            print("\nLogin successful!")
             break
         else:
             raise RuntimeError("Device authorization flow expired")
 
-    CREDENTIALS_PATH.parent.mkdir(parents=True, exist_ok=True)
-    expires = datetime.now(tz=timezone.utc) + timedelta(
-        seconds=response.expires_in - EXPIRES_GRACE_SECONDS
-    )
-    credential_data = {
-        "access_token": response.access_token,
-        "refresh_token": response.refresh_token,
-        "expires": expires.isoformat(),
-    }
-    CREDENTIALS_PATH.write_text(json.dumps(credential_data))
-    print(f"Saved credentials to {CREDENTIALS_PATH}")
+        # 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 logout():
-    CREDENTIALS_PATH.unlink(missing_ok=True)
-    # TODO: This should also revoke the refresh token
-    print(f"Removed credentials from {CREDENTIALS_PATH}")
+    async with DiracClient() as api:
+        credentials_path = get_diracx_preferences().credentials_path
+        if credentials_path:
+            credentials = json.loads(credentials_path.read_text())
+
+            # Revoke refresh token
+            try:
+                await api.auth.revoke_refresh_token(credentials["refresh_token"])
+            except Exception:
+                pass
+
+            # Remove credentials
+            credentials_path.unlink(missing_ok=True)
+            print(f"Removed credentials from {credentials_path}")
+    print("\nLogout successful!")
 
 
 @app.callback()

src/diracx/cli/jobs.py

@@ -11,11 +11,10 @@
 from rich.table import Table
 from typer import FileText, Option
 
-from diracx.client.aio import Dirac
+from diracx.client.aio import DiracClient
 from diracx.core.models import ScalarSearchOperator, SearchSpec, VectorSearchOperator
-from diracx.core.preferences import get_diracx_preferences
 
-from .utils import AsyncTyper, get_auth_headers
+from .utils import AsyncTyper
 
 app = AsyncTyper()
 
@@ -54,11 +53,10 @@
     condition: Annotated[list[SearchSpec], Option(parser=parse_condition)] = [],
     all: bool = False,
 ):
-    async with Dirac(endpoint=get_diracx_preferences().url) as api:
+    async with DiracClient() as api:
         jobs = await api.jobs.search(
             parameters=None if all else parameter,
             search=condition if condition else None,
-            headers=get_auth_headers(),
         )
     display(jobs, "jobs")
 
@@ -104,10 +102,9 @@
 
 @app.async_command()
 async def submit(jdl: list[FileText]):
-    async with Dirac(endpoint="http://localhost:8000") as api:
-        jobs = await api.jobs.submit_bulk_jobs(
-            [x.read() for x in jdl], headers=get_auth_headers()
-        )
+    async with DiracClient() as api:
+        # api.valid(enforce_https=False)
+        jobs = await api.jobs.submit_bulk_jobs([x.read() for x in jdl])
     print(
         f"Inserted {len(jobs)} jobs with ids: {','.join(map(str, (job.job_id for job in jobs)))}"
     )

src/diracx/cli/utils.py

@@ -1,16 +1,12 @@
 from __future__ import annotations
 
-__all__ = ("AsyncTyper", "CREDENTIALS_PATH", "get_auth_headers")
+__all__ = ("AsyncTyper",)
 
-import json
 from asyncio import run
 from functools import wraps
-from pathlib import Path
 
 import typer
 
-CREDENTIALS_PATH = Path.home() / ".cache" / "diracx" / "credentials.json"
-
 
 class AsyncTyper(typer.Typer):
     def async_command(self, *args, **kwargs):
@@ -23,11 +19,3 @@ def sync_func(*_args, **_kwargs):
             return async_func
 
         return decorator
-
-
-def get_auth_headers():
-    # TODO: Use autorest's actual mechanism for this
-    if not CREDENTIALS_PATH.exists():
-        raise NotImplementedError("Login first")
-    credentials = json.loads(CREDENTIALS_PATH.read_text())
-    return {"Authorization": f"Bearer {credentials['access_token']}"}

src/diracx/client/__init__.py

@@ -1,6 +1,6 @@
 # coding=utf-8
 # --------------------------------------------------------------------------
-# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.5, generator: @autorest/[email protected])
+# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.7, generator: @autorest/[email protected])
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 

src/diracx/client/_client.py

@@ -1,6 +1,6 @@
 # coding=utf-8
 # --------------------------------------------------------------------------
-# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.5, generator: @autorest/[email protected])
+# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.7, generator: @autorest/[email protected])
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 

src/diracx/client/_configuration.py

@@ -1,6 +1,6 @@
 # coding=utf-8
 # --------------------------------------------------------------------------
-# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.5, generator: @autorest/[email protected])
+# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.9.7, generator: @autorest/[email protected])
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------