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

diracx/client/{generated → _generated}/operations/_operations.py

@@ -1,7 +1,7 @@
 # pylint: disable=too-many-lines
 # coding=utf-8
 # --------------------------------------------------------------------------
-# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.10.4, generator: @autorest/python@6.32.1)
+# Code generated by Microsoft (R) AutoRest Code Generator (autorest: 3.10.4, generator: @autorest/python@6.32.2)
 # Changes may cause incorrect behavior and will be lost if the code is regenerated.
 # --------------------------------------------------------------------------
 from io import IOBase
@@ -508,6 +508,27 @@ def build_jobs_set_job_statuses_request(
     )
 
 
+def build_jobs_add_heartbeat_request(**kwargs: Any) -> HttpRequest:
+    _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
+
+    content_type: Optional[str] = kwargs.pop(
+        "content_type", _headers.pop("Content-Type", None)
+    )
+    accept = _headers.pop("Accept", "application/json")
+
+    # Construct URL
+    _url = "/api/jobs/heartbeat"
+
+    # Construct headers
+    if content_type is not None:
+        _headers["Content-Type"] = _SERIALIZER.header(
+            "content_type", content_type, "str"
+        )
+    _headers["Accept"] = _SERIALIZER.header("accept", accept, "str")
+
+    return HttpRequest(method="PATCH", url=_url, headers=_headers, **kwargs)
+
+
 def build_jobs_reschedule_jobs_request(
     *, job_ids: List[int], reset_jobs: bool = False, **kwargs: Any
 ) -> HttpRequest:
@@ -630,7 +651,7 @@ class WellKnownOperations:
         **DO NOT** instantiate this class directly.
 
         Instead, you should access the following operations through
-        :class:`~generated.Dirac`'s
+        :class:`~_generated.Dirac`'s
         :attr:`well_known` attribute.
     """
 
@@ -658,7 +679,7 @@ class WellKnownOperations:
         OpenID Connect discovery endpoint.
 
         :return: OpenIDConfiguration
-        :rtype: ~generated.models.OpenIDConfiguration
+        :rtype: ~_generated.models.OpenIDConfiguration
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -711,7 +732,7 @@ class WellKnownOperations:
         Get metadata about the dirac installation.
 
         :return: Metadata
-        :rtype: ~generated.models.Metadata
+        :rtype: ~_generated.models.Metadata
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -762,7 +783,7 @@ class AuthOperations:  # pylint: disable=abstract-class-instantiated
         **DO NOT** instantiate this class directly.
 
         Instead, you should access the following operations through
-        :class:`~generated.Dirac`'s
+        :class:`~_generated.Dirac`'s
         :attr:`auth` attribute.
     """
 
@@ -825,7 +846,7 @@ class AuthOperations:  # pylint: disable=abstract-class-instantiated
         :keyword scope: Required.
         :paramtype scope: str
         :return: InitiateDeviceFlowResponse
-        :rtype: ~generated.models.InitiateDeviceFlowResponse
+        :rtype: ~_generated.models.InitiateDeviceFlowResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -1160,7 +1181,7 @@ class AuthOperations:  # pylint: disable=abstract-class-instantiated
         Get information about the user's identity.
 
         :return: UserInfoResponse
-        :rtype: ~generated.models.UserInfoResponse
+        :rtype: ~_generated.models.UserInfoResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -1386,7 +1407,7 @@ class ConfigOperations:
         **DO NOT** instantiate this class directly.
 
         Instead, you should access the following operations through
-        :class:`~generated.Dirac`'s
+        :class:`~_generated.Dirac`'s
         :attr:`config` attribute.
     """
 
@@ -1493,7 +1514,7 @@ class JobsOperations:
         **DO NOT** instantiate this class directly.
 
         Instead, you should access the following operations through
-        :class:`~generated.Dirac`'s
+        :class:`~_generated.Dirac`'s
         :attr:`jobs` attribute.
     """
 
@@ -1533,12 +1554,12 @@ class JobsOperations:
         should be used to upload the sandbox to the storage backend.
 
         :param body: Required.
-        :type body: ~generated.models.SandboxInfo
+        :type body: ~_generated.models.SandboxInfo
         :keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
          Default value is "application/json".
         :paramtype content_type: str
         :return: SandboxUploadResponse
-        :rtype: ~generated.models.SandboxUploadResponse
+        :rtype: ~_generated.models.SandboxUploadResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -1562,7 +1583,7 @@ class JobsOperations:
          Default value is "application/json".
         :paramtype content_type: str
         :return: SandboxUploadResponse
-        :rtype: ~generated.models.SandboxUploadResponse
+        :rtype: ~_generated.models.SandboxUploadResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -1581,9 +1602,9 @@ class JobsOperations:
         should be used to upload the sandbox to the storage backend.
 
         :param body: Is either a SandboxInfo type or a IO[bytes] type. Required.
-        :type body: ~generated.models.SandboxInfo or IO[bytes]
+        :type body: ~_generated.models.SandboxInfo or IO[bytes]
         :return: SandboxUploadResponse
-        :rtype: ~generated.models.SandboxUploadResponse
+        :rtype: ~_generated.models.SandboxUploadResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -1660,7 +1681,7 @@ class JobsOperations:
         :keyword pfn: Required.
         :paramtype pfn: str
         :return: SandboxDownloadResponse
-        :rtype: ~generated.models.SandboxDownloadResponse
+        :rtype: ~_generated.models.SandboxDownloadResponse
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -1882,7 +1903,7 @@ class JobsOperations:
         :param job_id: Required.
         :type job_id: int
         :param sandbox_type: Known values are: "input" and "output". Required.
-        :type sandbox_type: str or ~generated.models.SandboxType
+        :type sandbox_type: str or ~_generated.models.SandboxType
         :return: list of any
         :rtype: list[any]
         :raises ~azure.core.exceptions.HttpResponseError:
@@ -2066,14 +2087,14 @@ class JobsOperations:
         Set Job Statuses.
 
         :param body: Required.
-        :type body: dict[str, dict[str, ~generated.models.JobStatusUpdate]]
+        :type body: dict[str, dict[str, ~_generated.models.JobStatusUpdate]]
         :keyword force: Default value is False.
         :paramtype force: bool
         :keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
          Default value is "application/json".
         :paramtype content_type: str
         :return: SetJobStatusReturn
-        :rtype: ~generated.models.SetJobStatusReturn
+        :rtype: ~_generated.models.SetJobStatusReturn
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -2098,7 +2119,7 @@ class JobsOperations:
          Default value is "application/json".
         :paramtype content_type: str
         :return: SetJobStatusReturn
-        :rtype: ~generated.models.SetJobStatusReturn
+        :rtype: ~_generated.models.SetJobStatusReturn
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -2115,11 +2136,11 @@ class JobsOperations:
         Set Job Statuses.
 
         :param body: Is either a {str: {str: JobStatusUpdate}} type or a IO[bytes] type. Required.
-        :type body: dict[str, dict[str, ~generated.models.JobStatusUpdate]] or IO[bytes]
+        :type body: dict[str, dict[str, ~_generated.models.JobStatusUpdate]] or IO[bytes]
         :keyword force: Default value is False.
         :paramtype force: bool
         :return: SetJobStatusReturn
-        :rtype: ~generated.models.SetJobStatusReturn
+        :rtype: ~_generated.models.SetJobStatusReturn
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {
@@ -2180,6 +2201,135 @@ class JobsOperations:
 
         return deserialized  # type: ignore
 
+    @overload
+    def add_heartbeat(
+        self,
+        body: Dict[str, _models.HeartbeatData],
+        *,
+        content_type: str = "application/json",
+        **kwargs: Any,
+    ) -> List[_models.JobCommand]:
+        """Add Heartbeat.
+
+        Register a heartbeat from the job.
+
+        This endpoint is used by the JobAgent to send heartbeats to the WMS and to
+        receive job commands from the WMS. It also results in stalled jobs being
+        restored to the RUNNING status.
+
+        The ``data`` parameter and return value are mappings keyed by job ID.
+
+        :param body: Required.
+        :type body: dict[str, ~_generated.models.HeartbeatData]
+        :keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
+         Default value is "application/json".
+        :paramtype content_type: str
+        :return: list of JobCommand
+        :rtype: list[~_generated.models.JobCommand]
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @overload
+    def add_heartbeat(
+        self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
+    ) -> List[_models.JobCommand]:
+        """Add Heartbeat.
+
+        Register a heartbeat from the job.
+
+        This endpoint is used by the JobAgent to send heartbeats to the WMS and to
+        receive job commands from the WMS. It also results in stalled jobs being
+        restored to the RUNNING status.
+
+        The ``data`` parameter and return value are mappings keyed by job ID.
+
+        :param body: Required.
+        :type body: IO[bytes]
+        :keyword content_type: Body Parameter content-type. Content type parameter for binary body.
+         Default value is "application/json".
+        :paramtype content_type: str
+        :return: list of JobCommand
+        :rtype: list[~_generated.models.JobCommand]
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @distributed_trace
+    def add_heartbeat(
+        self, body: Union[Dict[str, _models.HeartbeatData], IO[bytes]], **kwargs: Any
+    ) -> List[_models.JobCommand]:
+        """Add Heartbeat.
+
+        Register a heartbeat from the job.
+
+        This endpoint is used by the JobAgent to send heartbeats to the WMS and to
+        receive job commands from the WMS. It also results in stalled jobs being
+        restored to the RUNNING status.
+
+        The ``data`` parameter and return value are mappings keyed by job ID.
+
+        :param body: Is either a {str: HeartbeatData} type or a IO[bytes] type. Required.
+        :type body: dict[str, ~_generated.models.HeartbeatData] or IO[bytes]
+        :return: list of JobCommand
+        :rtype: list[~_generated.models.JobCommand]
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+        error_map: MutableMapping = {
+            401: ClientAuthenticationError,
+            404: ResourceNotFoundError,
+            409: ResourceExistsError,
+            304: ResourceNotModifiedError,
+        }
+        error_map.update(kwargs.pop("error_map", {}) or {})
+
+        _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
+        _params = kwargs.pop("params", {}) or {}
+
+        content_type: Optional[str] = kwargs.pop(
+            "content_type", _headers.pop("Content-Type", None)
+        )
+        cls: ClsType[List[_models.JobCommand]] = kwargs.pop("cls", None)
+
+        content_type = content_type or "application/json"
+        _json = None
+        _content = None
+        if isinstance(body, (IOBase, bytes)):
+            _content = body
+        else:
+            _json = self._serialize.body(body, "{HeartbeatData}")
+
+        _request = build_jobs_add_heartbeat_request(
+            content_type=content_type,
+            json=_json,
+            content=_content,
+            headers=_headers,
+            params=_params,
+        )
+        _request.url = self._client.format_url(_request.url)
+
+        _stream = False
+        pipeline_response: PipelineResponse = (
+            self._client._pipeline.run(  # pylint: disable=protected-access
+                _request, stream=_stream, **kwargs
+            )
+        )
+
+        response = pipeline_response.http_response
+
+        if response.status_code not in [200]:
+            map_error(
+                status_code=response.status_code, response=response, error_map=error_map
+            )
+            raise HttpResponseError(response=response)
+
+        deserialized = self._deserialize(
+            "[JobCommand]", pipeline_response.http_response
+        )
+
+        if cls:
+            return cls(pipeline_response, deserialized, {})  # type: ignore
+
+        return deserialized  # type: ignore
+
     @distributed_trace
     def reschedule_jobs(
         self, *, job_ids: List[int], reset_jobs: bool = False, **kwargs: Any
@@ -2361,7 +2511,7 @@ class JobsOperations:
         **TODO: Add more docs**.
 
         :param body: Default value is None.
-        :type body: ~generated.models.JobSearchParams
+        :type body: ~_generated.models.JobSearchParams
         :keyword page: Default value is 1.
         :paramtype page: int
         :keyword per_page: Default value is 100.
@@ -2420,7 +2570,7 @@ class JobsOperations:
         **TODO: Add more docs**.
 
         :param body: Is either a JobSearchParams type or a IO[bytes] type. Default value is None.
-        :type body: ~generated.models.JobSearchParams or IO[bytes]
+        :type body: ~_generated.models.JobSearchParams or IO[bytes]
         :keyword page: Default value is 1.
         :paramtype page: int
         :keyword per_page: Default value is 100.
@@ -2508,7 +2658,7 @@ class JobsOperations:
         Show information suitable for plotting.
 
         :param body: Required.
-        :type body: ~generated.models.JobSummaryParams
+        :type body: ~_generated.models.JobSummaryParams
         :keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
          Default value is "application/json".
         :paramtype content_type: str
@@ -2544,7 +2694,7 @@ class JobsOperations:
         Show information suitable for plotting.
 
         :param body: Is either a JobSummaryParams type or a IO[bytes] type. Required.
-        :type body: ~generated.models.JobSummaryParams or IO[bytes]
+        :type body: ~_generated.models.JobSummaryParams or IO[bytes]
         :return: any
         :rtype: any
         :raises ~azure.core.exceptions.HttpResponseError:
@@ -2618,7 +2768,7 @@ class JobsOperations:
          Default value is "application/json".
         :paramtype content_type: str
         :return: list of InsertedJob
-        :rtype: list[~generated.models.InsertedJob]
+        :rtype: list[~_generated.models.InsertedJob]
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -2636,7 +2786,7 @@ class JobsOperations:
          Default value is "application/json".
         :paramtype content_type: str
         :return: list of InsertedJob
-        :rtype: list[~generated.models.InsertedJob]
+        :rtype: list[~_generated.models.InsertedJob]
         :raises ~azure.core.exceptions.HttpResponseError:
         """
 
@@ -2651,7 +2801,7 @@ class JobsOperations:
         :param body: Is either a [str] type or a IO[bytes] type. Required.
         :type body: list[str] or IO[bytes]
         :return: list of InsertedJob
-        :rtype: list[~generated.models.InsertedJob]
+        :rtype: list[~_generated.models.InsertedJob]
         :raises ~azure.core.exceptions.HttpResponseError:
         """
         error_map: MutableMapping = {

diracx/client/_generated/operations/_patch.py

@@ -0,0 +1,26 @@
+# ------------------------------------
+# 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
+"""
+from __future__ import annotations
+
+__all__ = [
+    "AuthOperations",
+    "JobsOperations",
+]  # Add all objects you want publicly available to users at this package level
+
+from ...patches.auth.sync import AuthOperations
+from ...patches.jobs.sync import JobsOperations
+
+
+def patch_sdk():
+    """Do not remove from this file.
+
+    `patch_sdk` is a last resort escape hatch that allows you to do customizations
+    you can't accomplish using the techniques described in
+    https://aka.ms/azsdk/python/dpcodegen/python/customize
+    """

diracx/client/aio.py

@@ -1,3 +1,13 @@
-from .patches.aio import DiracClient
+from __future__ import absolute_import
 
-__all__ = ("DiracClient",)
+__all__ = [
+    "AsyncDiracClient",
+    "async_operations",
+]
+
+from ._generated.aio import Dirac
+from ._generated.aio import operations as async_operations
+
+
+class AsyncDiracClient(Dirac):
+    pass

diracx/client/models.py

@@ -1,7 +1,4 @@
-from .generated.models import *  # pylint: disable=unused-wildcard-import
+from __future__ import absolute_import
 
-
-# TODO: replace with postprocess
-from .generated.models import DeviceFlowErrorResponse
-
-__all__ = ("DeviceFlowErrorResponse",)
+from ._generated.models import *
+from ._generated.models import __all__

diracx/client/patches/auth/aio.py

@@ -0,0 +1,45 @@
+"""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__ = [
+    "AuthOperations",
+]
+
+from azure.core.pipeline import PipelineResponse
+from azure.core.tracing.decorator_async import distributed_trace_async
+from diracx.core.models import TokenResponse
+
+from ..._generated.aio.operations._operations import (
+    _models,
+    AuthOperations as _AuthOperations,
+)
+from .common import prepare_request, handle_response
+
+
+class AuthOperations(_AuthOperations):
+    @distributed_trace_async
+    async def get_oidc_token(
+        self, device_code: str, client_id: str, **kwargs
+    ) -> TokenResponse | _models.DeviceFlowErrorResponse:
+        request = prepare_request(
+            device_code=device_code,
+            client_id=client_id,
+            format_url=self._client.format_url,
+        )
+
+        pipeline_response: PipelineResponse = (
+            await self._client._pipeline.run(  # pylint: disable=protected-access
+                request, stream=False, **kwargs
+            )
+        )
+
+        response = handle_response(pipeline_response, self._deserialize)
+        if isinstance(response, _models.DeviceFlowErrorResponse):
+            return response
+        return TokenResponse.model_validate(response.as_dict())

diracx/client/patches/auth/common.py

@@ -0,0 +1,56 @@
+"""Utilities which are common to the sync and async auth operator patches."""
+
+from __future__ import annotations
+
+__all__ = [
+    "prepare_request",
+    "handle_response",
+]
+
+from typing import Any
+
+from azure.core.exceptions import map_error, HttpResponseError
+from azure.core.pipeline import PipelineResponse
+from azure.core.rest import HttpRequest
+from azure.core.utils import case_insensitive_dict
+
+from ..._generated.models import TokenResponse, DeviceFlowErrorResponse
+from ..._generated.operations._operations import _SERIALIZER
+
+
+def build_token_request(**kwargs: Any) -> HttpRequest:
+    _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
+
+    accept = _headers.pop("Accept", "application/json")
+
+    _url = "/api/auth/token"
+
+    _headers["Accept"] = _SERIALIZER.header("accept", accept, "str")
+
+    return HttpRequest(method="POST", url=_url, headers=_headers, **kwargs)
+
+
+def prepare_request(device_code, client_id, format_url) -> HttpRequest:
+    request = build_token_request(
+        data={
+            "grant_type": "urn:ietf:params:oauth:grant-type:device_code",
+            "device_code": device_code,
+            "client_id": client_id,
+        },
+    )
+    request.url = format_url(request.url)
+    return request
+
+
+def handle_response(
+    pipeline_response: PipelineResponse, deserialize
+) -> TokenResponse | DeviceFlowErrorResponse:
+    response = pipeline_response.http_response
+
+    if response.status_code == 200:
+        return deserialize("TokenResponse", pipeline_response)
+    elif response.status_code == 400:
+        return deserialize("DeviceFlowErrorResponse", pipeline_response)
+    else:
+        map_error(status_code=response.status_code, response=response, error_map={})
+        raise HttpResponseError(response=response)

diracx/client/patches/auth/sync.py

@@ -0,0 +1,41 @@
+"""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__ = [
+    "AuthOperations",
+]
+
+from azure.core.pipeline import PipelineResponse
+from azure.core.tracing.decorator import distributed_trace
+
+from ..._generated.operations._operations import (
+    _models,
+    AuthOperations as _AuthOperations,
+)
+from .common import prepare_request, handle_response
+
+
+class AuthOperations(_AuthOperations):
+    @distributed_trace
+    async def get_oidc_token(
+        self, device_code: str, client_id: str, **kwargs
+    ) -> _models.TokenResponse | _models.DeviceFlowErrorResponse:
+        request = prepare_request(
+            device_code=device_code,
+            client_id=client_id,
+            format_url=self._client.format_url,
+        )
+
+        pipeline_response: PipelineResponse = (
+            self._client._pipeline.run(  # pylint: disable=protected-access
+                request, stream=False, **kwargs
+            )
+        )
+
+        return handle_response(pipeline_response, self._deserialize)