diracx-client 0.0.1a33__py3-none-any.whl → 0.0.1a35__py3-none-any.whl

diracx/client/patches/{aio/utils.py → client/aio.py}

@@ -1,38 +1,27 @@
-# ------------------------------------
-# Copyright (c) Microsoft Corporation.
-# Licensed under the MIT License.
-# ------------------------------------
-"""Customize generated code here.
-
-Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize
-"""
+"""Patches for the autorest-generated client to enable authentication."""
+
 from __future__ import annotations
 
-import abc
 from importlib.metadata import PackageNotFoundError, distribution
-from types import TracebackType
 from pathlib import Path
-
-from typing import Any, List, Optional, cast
+from types import TracebackType
+from typing import Any, cast
 
 from azure.core.credentials import AccessToken
 from azure.core.credentials_async import AsyncTokenCredential
 from azure.core.pipeline import PipelineRequest
 from azure.core.pipeline.policies import AsyncBearerTokenCredentialPolicy
-
 from diracx.core.preferences import get_diracx_preferences, DiracxPreferences
 
-from ..utils import (
-    get_openid_configuration,
-    get_token,
-)
+from ..utils import get_openid_configuration, get_token
+from ..._generated.aio._client import Dirac as _Dirac
 
-__all__: List[str] = [
-    "DiracClient",
-]  # Add all objects you want publicly available to users at this package level
+__all__ = [
+    "Dirac",
+]
 
 
-class DiracTokenCredential(AsyncTokenCredential):
+class AsyncDiracTokenCredential(AsyncTokenCredential):
     """Tailor get_token() for our context"""
 
     def __init__(
@@ -51,8 +40,8 @@ class DiracTokenCredential(AsyncTokenCredential):
     async def get_token(
         self,
         *scopes: str,
-        claims: Optional[str] = None,
-        tenant_id: Optional[str] = None,
+        claims: str | None = None,
+        tenant_id: str | None = None,
         **kwargs: Any,
     ) -> AccessToken:
         return get_token(
@@ -81,18 +70,17 @@ class DiracTokenCredential(AsyncTokenCredential):
         pass
 
 
-class DiracBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
+class AsyncDiracBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
     """Custom AsyncBearerTokenCredentialPolicy tailored for our use case.
 
     * It does not ensure the connection is done through https.
     * It does not ensure that an access token is available.
     """
 
-    # Make mypy happy
-    _token: Optional[AccessToken] = None
+    _token: AccessToken | None = None
 
     def __init__(
-        self, credential: DiracTokenCredential, *scopes: str, **kwargs: Any
+        self, credential: AsyncDiracTokenCredential, *scopes: str, **kwargs: Any
     ) -> None:
         super().__init__(credential, *scopes, **kwargs)
 
@@ -104,9 +92,10 @@ class DiracBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
         :type request: ~azure.core.pipeline.PipelineRequest
         :raises: :class:`~azure.core.exceptions.ServiceRequestError`
         """
-        # Make mypy happy
         if not isinstance(self._credential, AsyncTokenCredential):
-            return
+            raise NotImplementedError(
+                "AsyncDiracBearerTokenCredentialPolicy only supports AsyncTokenCredential"
+            )
 
         self._token = await self._credential.get_token("", token=self._token)
         if not self._token.token:
@@ -119,10 +108,7 @@ class DiracBearerTokenCredentialPolicy(AsyncBearerTokenCredentialPolicy):
         )
 
 
-class DiracClientMixin(metaclass=abc.ABCMeta):
-    """This class inherits from the generated Dirac client and adds support for tokens,
-    so that the caller does not need to configure it by itself.
-    """
+class Dirac(_Dirac):
 
     def __init__(
         self,
@@ -156,14 +142,10 @@ class DiracClientMixin(metaclass=abc.ABCMeta):
             "DiracX-Client-Version"
         ] = self.client_version
 
-        # Initialize Dirac with a Dirac-specific token credential policy
-        # We need to ignore types here because mypy complains that we give
-        # too many arguments to "object" constructor as this is a mixin
-
-        super().__init__(  # type: ignore
+        super().__init__(
             endpoint=self._endpoint,
-            authentication_policy=DiracBearerTokenCredentialPolicy(
-                DiracTokenCredential(
+            authentication_policy=AsyncDiracBearerTokenCredentialPolicy(
+                AsyncDiracTokenCredential(
                     location=diracx_preferences.credentials_path,
                     token_endpoint=openid_configuration["token_endpoint"],
                     client_id=self._client_id,

diracx/client/patches/client/common.py

@@ -0,0 +1,196 @@
+"""Utilities which are common to the sync and async client patches."""
+
+from __future__ import annotations
+
+__all__ = [
+    "DiracAuthMixin",
+]
+
+import fcntl
+import json
+import os
+from datetime import datetime, timedelta, timezone
+from dataclasses import dataclass
+from enum import Enum
+from pathlib import Path
+from typing import TextIO
+from urllib import parse
+
+import httpx
+import jwt
+from azure.core.credentials import AccessToken
+from diracx.core.utils import EXPIRES_GRACE_SECONDS, serialize_credentials
+from diracx.core.models import TokenResponse
+
+
+class TokenStatus(Enum):
+    VALID = "valid"
+    REFRESH = "refresh"
+    INVALID = "invalid"
+
+
+@dataclass
+class TokenResult:
+    status: TokenStatus
+    access_token: AccessToken | None = None
+    refresh_token: str | None = None
+
+
+def get_openid_configuration(
+    endpoint: str, *, verify: bool | str = True
+) -> dict[str, str]:
+    """Get the openid configuration from the .well-known endpoint"""
+    response = httpx.get(
+        url=parse.urljoin(endpoint, ".well-known/openid-configuration"),
+        verify=verify,
+    )
+    if not response.is_success:
+        raise RuntimeError("Cannot fetch any information from the .well-known endpoint")
+    return response.json()
+
+
+def get_token(
+    location: Path,
+    token: AccessToken | None,
+    token_endpoint: str,
+    client_id: str,
+    verify: bool | str,
+) -> AccessToken:
+    """Get the access token if available and still valid."""
+    # Immediately return the token if it is available and still valid
+    if token and is_token_valid(token):
+        return token
+
+    if not location.exists():
+        # If we are here, it means the credentials path does not exist
+        # we suppose access token is not needed to perform the request
+        # we return an empty token to align with the expected return type
+        return AccessToken(token="", expires_on=0)
+
+    with open(location, "r+") as f:
+        # Acquire exclusive lock
+        fcntl.flock(f, fcntl.LOCK_EX)
+        try:
+            response = extract_token_from_credentials(f, token)
+            if response.status == TokenStatus.VALID and response.access_token:
+                # Lock is released in the finally block
+                return response.access_token
+
+            if response.status == TokenStatus.REFRESH and response.refresh_token:
+                # If we are here, it means the token needs to be refreshed
+                token_response = refresh_token(
+                    token_endpoint,
+                    client_id,
+                    response.refresh_token,
+                    verify=verify,
+                )
+
+                # Write the new credentials to the file
+                f.seek(0)
+                f.truncate()
+                f.write(serialize_credentials(token_response))
+                f.flush()
+                os.fsync(f.fileno())
+
+                # Get an AccessToken instance
+                return AccessToken(
+                    token=token_response.access_token,
+                    expires_on=int(
+                        (
+                            datetime.now(tz=timezone.utc)
+                            + timedelta(
+                                seconds=token_response.expires_in
+                                - EXPIRES_GRACE_SECONDS
+                            )
+                        ).timestamp()
+                    ),
+                )
+            # If we are here, it means the token is not available or not valid anymore
+            return AccessToken(token="", expires_on=0)
+        finally:
+            # Release the lock
+            fcntl.flock(f, fcntl.LOCK_UN)
+
+
+def refresh_token(
+    token_endpoint: str,
+    client_id: str,
+    refresh_token: str,
+    *,
+    verify: bool | str = True,
+) -> TokenResponse:
+    """Refresh the access token using the refresh_token flow."""
+    response = httpx.post(
+        url=token_endpoint,
+        data={
+            "client_id": client_id,
+            "grant_type": "refresh_token",
+            "refresh_token": refresh_token,
+        },
+        verify=verify,
+    )
+
+    if response.status_code != 200:
+        raise RuntimeError(
+            f"An issue occured while refreshing your access token: {response.json()['detail']}"
+        )
+
+    res = response.json()
+    return TokenResponse(
+        access_token=res["access_token"],
+        expires_in=res["expires_in"],
+        token_type=res.get("token_type"),
+        refresh_token=res.get("refresh_token"),
+    )
+
+
+def is_token_valid(token: AccessToken) -> bool:
+    """Condition to get a new token"""
+    return (
+        datetime.fromtimestamp(token.expires_on, tz=timezone.utc)
+        - datetime.now(tz=timezone.utc)
+    ).total_seconds() > 300
+
+
+def extract_token_from_credentials(
+    token_file_descriptor: TextIO, token: AccessToken | None
+) -> TokenResult:
+    """Get token if available and still valid."""
+    # If we are here, it means the token is not available or not valid anymore
+    # We try to get it from the file
+    try:
+        credentials = json.load(token_file_descriptor)
+    except json.JSONDecodeError:
+        return TokenResult(TokenStatus.INVALID)
+
+    try:
+        token = AccessToken(
+            token=credentials["access_token"],
+            expires_on=credentials["expires_on"],
+        )
+        refresh_token = credentials["refresh_token"]
+    except KeyError:
+        return TokenResult(TokenStatus.INVALID)
+
+    # We check the validity of the tokens
+    if is_token_valid(token):
+        return TokenResult(TokenStatus.VALID, access_token=token)
+
+    if is_refresh_token_valid(refresh_token):
+        return TokenResult(TokenStatus.REFRESH, refresh_token=refresh_token)
+
+    # If we are here, it means the refresh token is not valid anymore
+    return TokenResult(TokenStatus.INVALID)
+
+
+def is_refresh_token_valid(refresh_token: str | None) -> bool:
+    """Check if the refresh token is still valid."""
+    if not refresh_token:
+        return False
+    # Decode the refresh token
+    refresh_payload = jwt.decode(refresh_token, options={"verify_signature": False})
+    if not refresh_payload or "exp" not in refresh_payload:
+        return False
+
+    # Check the expiration time
+    return refresh_payload["exp"] > datetime.now(tz=timezone.utc).timestamp()

diracx/client/patches/client/sync.py

@@ -0,0 +1,141 @@
+"""Patches for the autorest-generated client to enable authentication."""
+
+from __future__ import annotations
+
+__all__ = [
+    "Dirac",
+]
+
+from importlib.metadata import PackageNotFoundError, distribution
+from pathlib import Path
+from typing import Any, Optional
+
+from azure.core.credentials import AccessToken, TokenCredential
+from azure.core.pipeline import PipelineRequest
+from azure.core.pipeline.policies import BearerTokenCredentialPolicy
+from diracx.core.preferences import DiracxPreferences, get_diracx_preferences
+
+from .common import get_openid_configuration, get_token
+from ..._generated._client import Dirac as _Dirac
+
+
+class SyncDiracTokenCredential(TokenCredential):
+    """Tailor get_token() for our context"""
+
+    def __init__(
+        self,
+        location: Path,
+        token_endpoint: str,
+        client_id: str,
+        *,
+        verify: bool | str = True,
+    ) -> None:
+        self.location = location
+        self.verify = verify
+        self.token_endpoint = token_endpoint
+        self.client_id = client_id
+
+    def get_token(
+        self,
+        *scopes: str,
+        claims: Optional[str] = None,
+        tenant_id: Optional[str] = None,
+        **kwargs: Any,
+    ) -> AccessToken:
+        return get_token(
+            self.location,
+            kwargs.get("token"),
+            self.token_endpoint,
+            self.client_id,
+            self.verify,
+        )
+
+
+class SyncDiracBearerTokenCredentialPolicy(BearerTokenCredentialPolicy):
+    """Custom BearerTokenCredentialPolicy tailored for our use case.
+
+    * It does not ensure the connection is done through https.
+    * It does not ensure that an access token is available.
+    """
+
+    # Make mypy happy
+    _token: Optional[AccessToken] = None
+
+    def __init__(
+        self, credential: SyncDiracTokenCredential, *scopes: str, **kwargs: Any
+    ) -> None:
+        super().__init__(credential, *scopes, **kwargs)
+
+    def on_request(self, request: PipelineRequest) -> None:
+        """Authorization Bearer is optional here.
+        :param request: The pipeline request object to be modified.
+        :type request: ~azure.core.pipeline.PipelineRequest
+        :raises: :class:`~azure.core.exceptions.ServiceRequestError`
+        """
+        if not isinstance(self._credential, TokenCredential):
+            raise NotImplementedError(
+                "SyncDiracBearerTokenCredentialPolicy only supports TokenCredential"
+            )
+
+        self._token = self._credential.get_token("", token=self._token)
+        if not self._token.token:
+            # If we are here, it means the token is not available
+            # we suppose it is not needed to perform the request
+            return
+
+        self._update_headers(request.http_request.headers, self._token.token)
+
+
+class Dirac(_Dirac):
+    """This class inherits from the generated Dirac client and adds support for tokens,
+    so that the caller does not need to configure it by itself.
+    """
+
+    def __init__(
+        self,
+        endpoint: str | None = None,
+        client_id: str | None = None,
+        diracx_preferences: DiracxPreferences | None = None,
+        verify: bool | str = True,
+        **kwargs: Any,
+    ) -> None:
+        diracx_preferences = diracx_preferences or get_diracx_preferences()
+        self._endpoint = str(endpoint or diracx_preferences.url)
+        if verify is True and diracx_preferences.ca_path:
+            verify = str(diracx_preferences.ca_path)
+        kwargs["connection_verify"] = verify
+        self._client_id = client_id or "myDIRACClientID"
+
+        # Get .well-known configuration
+        openid_configuration = get_openid_configuration(self._endpoint, verify=verify)
+
+        try:
+            self.client_version = distribution("diracx").version
+        except PackageNotFoundError:
+            try:
+                self.client_version = distribution("diracx-client").version
+            except PackageNotFoundError:
+                print("Error while getting client version")
+                self.client_version = "Unknown"
+
+        # Setting default headers
+        kwargs.setdefault("base_headers", {})[
+            "DiracX-Client-Version"
+        ] = self.client_version
+
+        super().__init__(
+            endpoint=self._endpoint,
+            authentication_policy=SyncDiracBearerTokenCredentialPolicy(
+                SyncDiracTokenCredential(
+                    location=diracx_preferences.credentials_path,
+                    token_endpoint=openid_configuration["token_endpoint"],
+                    client_id=self._client_id,
+                    verify=verify,
+                ),
+            ),
+            **kwargs,
+        )
+
+    @property
+    def client_id(self):
+        return self._client_id

diracx/client/patches/jobs/aio.py

@@ -0,0 +1,34 @@
+"""Patches for the autorest-generated jobs client.
+
+This file can be used to customize the generated code for the jobs client.
+When adding new classes to this file, make sure to also add them to the
+__all__ list in the corresponding file in the patches directory.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "JobsOperations",
+]
+
+from typing import Any, Unpack
+
+from azure.core.tracing.decorator_async import distributed_trace_async
+
+from ..._generated.aio.operations._operations import JobsOperations as _JobsOperations
+from .common import make_search_body, make_summary_body, SearchKwargs, SummaryKwargs
+
+# We're intentionally ignoring overrides here because we want to change the interface.
+# mypy: disable-error-code=override
+
+
+class JobsOperations(_JobsOperations):
+    @distributed_trace_async
+    async def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return await super().search(**make_search_body(**kwargs))
+
+    @distributed_trace_async
+    async def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return await super().summary(**make_summary_body(**kwargs))

diracx/client/patches/jobs/common.py

@@ -0,0 +1,85 @@
+"""Utilities which are common to the sync and async jobs operator patches."""
+
+from __future__ import annotations
+
+__all__ = [
+    "make_search_body",
+    "SearchKwargs",
+    "make_summary_body",
+    "SummaryKwargs",
+]
+
+import json
+from io import BytesIO
+from typing import Any, IO, TypedDict, Unpack, cast, Literal
+
+from diracx.core.models import SearchSpec
+
+
+class ResponseExtra(TypedDict, total=False):
+    content_type: str
+    headers: dict[str, str]
+    params: dict[str, str]
+    cls: Any
+
+
+class SearchBody(TypedDict, total=False):
+    parameters: list[str] | None
+    search: list[SearchSpec] | None
+    sort: list[str] | None
+
+
+class SearchExtra(ResponseExtra, total=False):
+    page: int
+    per_page: int
+
+
+class SearchKwargs(SearchBody, SearchExtra): ...
+
+
+class UnderlyingSearchArgs(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_search_body(**kwargs: Unpack[SearchKwargs]) -> UnderlyingSearchArgs:
+    body: SearchBody = {}
+    for key in SearchBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["parameters", "search", "sort"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingSearchArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(SearchExtra, kwargs))
+    return result
+
+
+class SummaryBody(TypedDict, total=False):
+    grouping: list[str]
+    search: list[str]
+
+
+class SummaryKwargs(SummaryBody, ResponseExtra): ...
+
+
+class UnderlyingSummaryArgs(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_summary_body(**kwargs: Unpack[SummaryKwargs]) -> UnderlyingSummaryArgs:
+    body: SummaryBody = {}
+    for key in SummaryBody.__optional_keys__:
+        if key not in kwargs:
+            continue
+        key = cast(Literal["grouping", "search"], key)
+        value = kwargs.pop(key)
+        if value is not None:
+            body[key] = value
+    result: UnderlyingSummaryArgs = {"body": BytesIO(json.dumps(body).encode("utf-8"))}
+    result.update(cast(ResponseExtra, kwargs))
+    return result

diracx/client/patches/jobs/sync.py

@@ -0,0 +1,34 @@
+"""Patches for the autorest-generated jobs client.
+
+This file can be used to customize the generated code for the jobs client.
+When adding new classes to this file, make sure to also add them to the
+__all__ list in the corresponding file in the patches directory.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "JobsOperations",
+]
+
+from typing import Any, Unpack
+
+from azure.core.tracing.decorator import distributed_trace
+
+from ..._generated.operations._operations import JobsOperations as _JobsOperations
+from .common import make_search_body, make_summary_body, SearchKwargs, SummaryKwargs
+
+# We're intentionally ignoring overrides here because we want to change the interface.
+# mypy: disable-error-code=override
+
+
+class JobsOperations(_JobsOperations):
+    @distributed_trace
+    def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return super().search(**make_search_body(**kwargs))
+
+    @distributed_trace
+    def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]:
+        """TODO"""
+        return super().summary(**make_summary_body(**kwargs))

diracx/client/py.typed

@@ -1 +0,0 @@
-# Marker file for PEP 561.

diracx/client/sync.py

@@ -0,0 +1,13 @@
+from __future__ import absolute_import
+
+__all__ = [
+    "SyncDiracClient",
+    "sync_operations",
+]
+
+from ._generated import Dirac
+from ._generated import operations as sync_operations
+
+
+class SyncDiracClient(Dirac):
+    pass

{diracx_client-0.0.1a33.dist-info → diracx_client-0.0.1a35.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: diracx-client
-Version: 0.0.1a33
+Version: 0.0.1a35
 Summary: TODO
 License: GPL-3.0-only
 Classifier: Intended Audience :: Science/Research
@@ -9,11 +9,10 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: System :: Distributed Computing
 Requires-Python: >=3.11
-Description-Content-Type: text/markdown
 Requires-Dist: azure-core
 Requires-Dist: diracx-core
-Requires-Dist: isodate
 Requires-Dist: httpx
+Requires-Dist: isodate
 Provides-Extra: testing
-Requires-Dist: diracx-testing; extra == "testing"
+Requires-Dist: diracx-testing; extra == 'testing'
 Provides-Extra: types