feat: integration authz headers in Dirac client

src/diracx/cli/__init__.py

@@ -1,23 +1,17 @@
 from __future__ import annotations
 
-import asyncio
-import json
 import os
-from datetime import datetime, timedelta, timezone
 from typing import Optional
 
 from typer import Option
 
 from diracx.client.aio import Dirac
-from diracx.client.models import DeviceFlowErrorResponse
 
 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 +28,16 @@ async def login(
         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:
-        data = await api.auth.initiate_device_flow(
-            client_id="myDIRACClientID",
-            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"
-            )
-            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}")
+    async with Dirac() as api:
+        await api.login(scopes)
+    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 Dirac() as api:
+        await api.logout()
+    print("\nLogout successful!")
 
 
 @app.callback()

src/diracx/cli/jobs.py

@@ -13,9 +13,8 @@
 
 from diracx.client.aio import Dirac
 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 @@ async def search(
     condition: Annotated[list[SearchSpec], Option(parser=parse_condition)] = [],
     all: bool = False,
 ):
-    async with Dirac(endpoint=get_diracx_preferences().url) as api:
+    async with Dirac() 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 @@ def display_rich(data, unit: str) -> None:
 
 @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 Dirac() 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,8 +1,7 @@
 from __future__ import annotations
 
-__all__ = ("AsyncTyper", "CREDENTIALS_PATH", "get_auth_headers")
+__all__ = ("AsyncTyper", "CREDENTIALS_PATH")
 
-import json
 from asyncio import run
 from functools import wraps
 from pathlib import Path
@@ -23,11 +22,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.
 # --------------------------------------------------------------------------
 

src/diracx/client/_vendor.py

@@ -1,5 +1,5 @@
 # --------------------------------------------------------------------------
-# 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/aio/__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/aio/_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/aio/_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.
 # --------------------------------------------------------------------------