feat: Support diracx pilots for jobs

Author: Robin-Van-de-Merghel

diracx-client/src/diracx/client/_generated/aio/operations/_operations.py

@@ -58,7 +58,10 @@
     build_pilots_clear_pilots_request,
     build_pilots_create_pilot_secrets_request,
     build_pilots_delete_pilots_request,
+    build_pilots_pilot_add_heartbeat_request,
     build_pilots_pilot_login_request,
+    build_pilots_pilot_patch_metadata_request,
+    build_pilots_pilot_set_job_statuses_request,
     build_pilots_refresh_pilot_tokens_request,
     build_pilots_search_request,
     build_pilots_update_pilot_fields_request,
@@ -3176,3 +3179,383 @@ async def refresh_pilot_tokens(
             return cls(pipeline_response, deserialized, {})  # type: ignore
 
         return deserialized  # type: ignore
+
+    @overload
+    async def pilot_set_job_statuses(
+        self,
+        body: Dict[str, Dict[str, _models.JobStatusUpdate]],
+        *,
+        force: bool = False,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> _models.SetJobStatusReturn:
+        """Pilot Set Job Statuses.
+
+        Pilot Set Job Statuses.
+
+        :param body: Required.
+        :type body: dict[str, dict[str, ~_generated.models.JobStatusUpdate]]
+        :keyword force: Default value is False.
+        :paramtype force: bool
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :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
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @overload
+    async def pilot_set_job_statuses(
+        self,
+        body: IO[bytes],
+        *,
+        force: bool = False,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> _models.SetJobStatusReturn:
+        """Pilot Set Job Statuses.
+
+        Pilot Set Job Statuses.
+
+        :param body: Required.
+        :type body: IO[bytes]
+        :keyword force: Default value is False.
+        :paramtype force: bool
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :keyword content_type: Body Parameter content-type. Content type parameter for binary body.
+         Default value is "application/json".
+        :paramtype content_type: str
+        :return: SetJobStatusReturn
+        :rtype: ~_generated.models.SetJobStatusReturn
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @distributed_trace_async
+    async def pilot_set_job_statuses(
+        self,
+        body: Union[Dict[str, Dict[str, _models.JobStatusUpdate]], IO[bytes]],
+        *,
+        force: bool = False,
+        authorization: Optional[str] = None,
+        **kwargs: Any
+    ) -> _models.SetJobStatusReturn:
+        """Pilot Set Job Statuses.
+
+        Pilot 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]
+        :keyword force: Default value is False.
+        :paramtype force: bool
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :return: SetJobStatusReturn
+        :rtype: ~_generated.models.SetJobStatusReturn
+        :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[_models.SetJobStatusReturn] = 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, "{{JobStatusUpdate}}")
+
+        _request = build_pilots_pilot_set_job_statuses_request(
+            force=force,
+            authorization=authorization,
+            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 = await 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("SetJobStatusReturn", pipeline_response.http_response)
+
+        if cls:
+            return cls(pipeline_response, deserialized, {})  # type: ignore
+
+        return deserialized  # type: ignore
+
+    @overload
+    async def pilot_add_heartbeat(
+        self,
+        body: Dict[str, _models.HeartbeatData],
+        *,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> List[_models.JobCommand]:
+        """Pilot 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 authorization: Default value is None.
+        :paramtype authorization: str
+        :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
+    async def pilot_add_heartbeat(
+        self,
+        body: IO[bytes],
+        *,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> List[_models.JobCommand]:
+        """Pilot 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 authorization: Default value is None.
+        :paramtype authorization: str
+        :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_async
+    async def pilot_add_heartbeat(
+        self,
+        body: Union[Dict[str, _models.HeartbeatData], IO[bytes]],
+        *,
+        authorization: Optional[str] = None,
+        **kwargs: Any
+    ) -> List[_models.JobCommand]:
+        """Pilot 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]
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :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_pilots_pilot_add_heartbeat_request(
+            authorization=authorization,
+            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 = await 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
+
+    @overload
+    async def pilot_patch_metadata(
+        self,
+        body: Dict[str, Dict[str, Any]],
+        *,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> None:
+        """Pilot Patch Metadata.
+
+        Pilot Patch Metadata.
+
+        :param body: Required.
+        :type body: dict[str, dict[str, any]]
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
+         Default value is "application/json".
+        :paramtype content_type: str
+        :return: None
+        :rtype: None
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @overload
+    async def pilot_patch_metadata(
+        self,
+        body: IO[bytes],
+        *,
+        authorization: Optional[str] = None,
+        content_type: str = "application/json",
+        **kwargs: Any
+    ) -> None:
+        """Pilot Patch Metadata.
+
+        Pilot Patch Metadata.
+
+        :param body: Required.
+        :type body: IO[bytes]
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :keyword content_type: Body Parameter content-type. Content type parameter for binary body.
+         Default value is "application/json".
+        :paramtype content_type: str
+        :return: None
+        :rtype: None
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @distributed_trace_async
+    async def pilot_patch_metadata(
+        self, body: Union[Dict[str, Dict[str, Any]], IO[bytes]], *, authorization: Optional[str] = None, **kwargs: Any
+    ) -> None:
+        """Pilot Patch Metadata.
+
+        Pilot Patch Metadata.
+
+        :param body: Is either a {str: {str: Any}} type or a IO[bytes] type. Required.
+        :type body: dict[str, dict[str, any]] or IO[bytes]
+        :keyword authorization: Default value is None.
+        :paramtype authorization: str
+        :return: None
+        :rtype: None
+        :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[None] = 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, "{{object}}")
+
+        _request = build_pilots_pilot_patch_metadata_request(
+            authorization=authorization,
+            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 = await self._client._pipeline.run(  # pylint: disable=protected-access
+            _request, stream=_stream, **kwargs
+        )
+
+        response = pipeline_response.http_response
+
+        if response.status_code not in [204]:
+            map_error(status_code=response.status_code, response=response, error_map=error_map)
+            raise HttpResponseError(response=response)
+
+        if cls:
+            return cls(pipeline_response, None, {})  # type: ignore