feat: DIRAC can associate a pilot with a job

Author: Robin-Van-de-Merghel

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

@@ -55,6 +55,7 @@
     build_jobs_summary_request,
     build_jobs_unassign_bulk_jobs_sandboxes_request,
     build_jobs_unassign_job_sandboxes_request,
+    build_pilots_associate_pilot_with_jobs_request,
     build_pilots_associate_pilots_with_secrets_request,
     build_pilots_create_pilot_secrets_request,
     build_pilots_patch_pilot_data_request,
@@ -3202,3 +3203,110 @@ async def update_pilot_fields(
 
         if cls:
             return cls(pipeline_response, None, {})  # type: ignore
+
+    @overload
+    async def associate_pilot_with_jobs(
+        self,
+        body: _models.BodyPilotsAssociatePilotWithJobs,
+        *,
+        content_type: str = "application/json",
+        **kwargs: Any,
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :param body: Required.
+        :type body: ~_generated.models.BodyPilotsAssociatePilotWithJobs
+        :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 associate_pilot_with_jobs(
+        self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :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: None
+        :rtype: None
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @distributed_trace_async
+    async def associate_pilot_with_jobs(
+        self,
+        body: Union[_models.BodyPilotsAssociatePilotWithJobs, IO[bytes]],
+        **kwargs: Any,
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :param body: Is either a BodyPilotsAssociatePilotWithJobs type or a IO[bytes] type. Required.
+        :type body: ~_generated.models.BodyPilotsAssociatePilotWithJobs or IO[bytes]
+        :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, "BodyPilotsAssociatePilotWithJobs")
+
+        _request = build_pilots_associate_pilot_with_jobs_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 = (
+            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

diracx-client/src/diracx/client/_generated/models/__init__.py

@@ -16,6 +16,7 @@
     BodyAuthGetOidcTokenGrantType,
     BodyAuthPilotLogin,
     BodyAuthRefreshPilotTokens,
+    BodyPilotsAssociatePilotWithJobs,
     BodyPilotsAssociatePilotsWithSecrets,
     BodyPilotsCreatePilotSecrets,
     BodyPilotsRegisterNewPilotsToDb,
@@ -74,6 +75,7 @@
     "BodyAuthGetOidcTokenGrantType",
     "BodyAuthPilotLogin",
     "BodyAuthRefreshPilotTokens",
+    "BodyPilotsAssociatePilotWithJobs",
     "BodyPilotsAssociatePilotsWithSecrets",
     "BodyPilotsCreatePilotSecrets",
     "BodyPilotsRegisterNewPilotsToDb",

diracx-client/src/diracx/client/_generated/models/_models.py

@@ -190,6 +190,41 @@ def __init__(
         self.pilot_secrets = pilot_secrets
 
 
+class BodyPilotsAssociatePilotWithJobs(_serialization.Model):
+    """Body_pilots_associate_pilot_with_jobs.
+
+    All required parameters must be populated in order to send to server.
+
+    :ivar pilot_stamp: The stamp of the pilot. Required.
+    :vartype pilot_stamp: str
+    :ivar pilot_jobs_ids: The jobs we want to add to the pilot. Required.
+    :vartype pilot_jobs_ids: list[int]
+    """
+
+    _validation = {
+        "pilot_stamp": {"required": True},
+        "pilot_jobs_ids": {"required": True},
+    }
+
+    _attribute_map = {
+        "pilot_stamp": {"key": "pilot_stamp", "type": "str"},
+        "pilot_jobs_ids": {"key": "pilot_jobs_ids", "type": "[int]"},
+    }
+
+    def __init__(
+        self, *, pilot_stamp: str, pilot_jobs_ids: List[int], **kwargs: Any
+    ) -> None:
+        """
+        :keyword pilot_stamp: The stamp of the pilot. Required.
+        :paramtype pilot_stamp: str
+        :keyword pilot_jobs_ids: The jobs we want to add to the pilot. Required.
+        :paramtype pilot_jobs_ids: list[int]
+        """
+        super().__init__(**kwargs)
+        self.pilot_stamp = pilot_stamp
+        self.pilot_jobs_ids = pilot_jobs_ids
+
+
 class BodyPilotsCreatePilotSecrets(_serialization.Model):
     """Body_pilots_create_pilot_secrets.
 

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

@@ -824,6 +824,26 @@ def build_pilots_update_pilot_fields_request(**kwargs: Any) -> HttpRequest:
     return HttpRequest(method="PATCH", url=_url, headers=_headers, **kwargs)
 
 
+def build_pilots_associate_pilot_with_jobs_request(
+    **kwargs: Any,
+) -> HttpRequest:  # pylint: disable=name-too-long
+    _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
+
+    content_type: Optional[str] = kwargs.pop(
+        "content_type", _headers.pop("Content-Type", None)
+    )
+    # Construct URL
+    _url = "/api/pilots/fields/jobs"
+
+    # Construct headers
+    if content_type is not None:
+        _headers["Content-Type"] = _SERIALIZER.header(
+            "content_type", content_type, "str"
+        )
+
+    return HttpRequest(method="PATCH", url=_url, headers=_headers, **kwargs)
+
+
 class WellKnownOperations:
     """
     .. warning::
@@ -3385,7 +3405,7 @@ class PilotsOperations:
 
     models = _models
 
-    def __init__(self, *args, **kwargs):
+    def __init__(self, *args, **kwargs) -> None:
         input_args = list(args)
         self._client: PipelineClient = (
             input_args.pop(0) if input_args else kwargs.pop("client")
@@ -3950,3 +3970,110 @@ def update_pilot_fields(  # pylint: disable=inconsistent-return-statements
 
         if cls:
             return cls(pipeline_response, None, {})  # type: ignore
+
+    @overload
+    def associate_pilot_with_jobs(
+        self,
+        body: _models.BodyPilotsAssociatePilotWithJobs,
+        *,
+        content_type: str = "application/json",
+        **kwargs: Any,
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :param body: Required.
+        :type body: ~_generated.models.BodyPilotsAssociatePilotWithJobs
+        :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
+    def associate_pilot_with_jobs(
+        self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :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: None
+        :rtype: None
+        :raises ~azure.core.exceptions.HttpResponseError:
+        """
+
+    @distributed_trace
+    def associate_pilot_with_jobs(  # pylint: disable=inconsistent-return-statements
+        self,
+        body: Union[_models.BodyPilotsAssociatePilotWithJobs, IO[bytes]],
+        **kwargs: Any,
+    ) -> None:
+        """Associate Pilot With Jobs.
+
+        Associate Pilot With Jobs.
+
+        :param body: Is either a BodyPilotsAssociatePilotWithJobs type or a IO[bytes] type. Required.
+        :type body: ~_generated.models.BodyPilotsAssociatePilotWithJobs or IO[bytes]
+        :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, "BodyPilotsAssociatePilotWithJobs")
+
+        _request = build_pilots_associate_pilot_with_jobs_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 [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

diracx-core/src/diracx/core/exceptions.py

@@ -165,3 +165,13 @@ class SecretHasExpiredError(GenericError):
 class SecretAlreadyExistsError(GenericError):
     head = "Secret"
     tail = "already exists"
+
+
+class PilotJobsNotFoundError(GenericError):
+    head = "Pilots or Jobs"
+    tail = "not found"
+
+
+class PilotAlreadyAssociatedWithJobError(GenericError):
+    head = "Pilot is already associated with a job"
+    tail = ""