diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 13309f1a5..1c8788fd3 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -69,7 +69,9 @@ repos: rev: v2.4.3 hooks: - id: codespell - args: ["-w"] + # Configuration lives in [tool.codespell] in pyproject.toml + additional_dependencies: + - tomli exclude: pixi\.lock$ - repo: https://github.com/adamchainz/blacken-docs diff --git a/diracx-client/src/diracx/client/_generated/_client.py b/diracx-client/src/diracx/client/_generated/_client.py index c5e641f24..e145960ae 100644 --- a/diracx-client/src/diracx/client/_generated/_client.py +++ b/diracx-client/src/diracx/client/_generated/_client.py @@ -15,7 +15,14 @@ from . import models as _models from ._configuration import DiracConfiguration from ._utils.serialization import Deserializer, Serializer -from .operations import AuthOperations, ConfigOperations, JobsOperations, RssOperations, WellKnownOperations +from .operations import ( + AuthOperations, + ConfigOperations, + JobsOperations, + PilotsOperations, + RssOperations, + WellKnownOperations, +) class Dirac: # pylint: disable=client-accepts-api-version-keyword @@ -29,6 +36,8 @@ class Dirac: # pylint: disable=client-accepts-api-version-keyword :vartype config: _generated.operations.ConfigOperations :ivar jobs: JobsOperations operations :vartype jobs: _generated.operations.JobsOperations + :ivar pilots: PilotsOperations operations + :vartype pilots: _generated.operations.PilotsOperations :ivar rss: RssOperations operations :vartype rss: _generated.operations.RssOperations :keyword endpoint: Service URL. Required. Default value is "". @@ -67,6 +76,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self.auth = AuthOperations(self._client, self._config, self._serialize, self._deserialize) self.config = ConfigOperations(self._client, self._config, self._serialize, self._deserialize) self.jobs = JobsOperations(self._client, self._config, self._serialize, self._deserialize) + self.pilots = PilotsOperations(self._client, self._config, self._serialize, self._deserialize) self.rss = RssOperations(self._client, self._config, self._serialize, self._deserialize) def send_request(self, request: HttpRequest, *, stream: bool = False, **kwargs: Any) -> HttpResponse: diff --git a/diracx-client/src/diracx/client/_generated/aio/_client.py b/diracx-client/src/diracx/client/_generated/aio/_client.py index f2d2b46bf..b3187208e 100644 --- a/diracx-client/src/diracx/client/_generated/aio/_client.py +++ b/diracx-client/src/diracx/client/_generated/aio/_client.py @@ -15,7 +15,14 @@ from .. import models as _models from .._utils.serialization import Deserializer, Serializer from ._configuration import DiracConfiguration -from .operations import AuthOperations, ConfigOperations, JobsOperations, RssOperations, WellKnownOperations +from .operations import ( + AuthOperations, + ConfigOperations, + JobsOperations, + PilotsOperations, + RssOperations, + WellKnownOperations, +) class Dirac: # pylint: disable=client-accepts-api-version-keyword @@ -29,6 +36,8 @@ class Dirac: # pylint: disable=client-accepts-api-version-keyword :vartype config: _generated.aio.operations.ConfigOperations :ivar jobs: JobsOperations operations :vartype jobs: _generated.aio.operations.JobsOperations + :ivar pilots: PilotsOperations operations + :vartype pilots: _generated.aio.operations.PilotsOperations :ivar rss: RssOperations operations :vartype rss: _generated.aio.operations.RssOperations :keyword endpoint: Service URL. Required. Default value is "". @@ -67,6 +76,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self.auth = AuthOperations(self._client, self._config, self._serialize, self._deserialize) self.config = ConfigOperations(self._client, self._config, self._serialize, self._deserialize) self.jobs = JobsOperations(self._client, self._config, self._serialize, self._deserialize) + self.pilots = PilotsOperations(self._client, self._config, self._serialize, self._deserialize) self.rss = RssOperations(self._client, self._config, self._serialize, self._deserialize) def send_request( diff --git a/diracx-client/src/diracx/client/_generated/aio/operations/__init__.py b/diracx-client/src/diracx/client/_generated/aio/operations/__init__.py index 77674abec..48074896e 100644 --- a/diracx-client/src/diracx/client/_generated/aio/operations/__init__.py +++ b/diracx-client/src/diracx/client/_generated/aio/operations/__init__.py @@ -14,6 +14,7 @@ from ._operations import AuthOperations # type: ignore from ._operations import ConfigOperations # type: ignore from ._operations import JobsOperations # type: ignore +from ._operations import PilotsOperations # type: ignore from ._operations import RssOperations # type: ignore from ._patch import __all__ as _patch_all @@ -25,6 +26,7 @@ "AuthOperations", "ConfigOperations", "JobsOperations", + "PilotsOperations", "RssOperations", ] __all__.extend([p for p in _patch_all if p not in __all__]) # pyright: ignore diff --git a/diracx-client/src/diracx/client/_generated/aio/operations/_operations.py b/diracx-client/src/diracx/client/_generated/aio/operations/_operations.py index 4efe2dd6f..09957588e 100644 --- a/diracx-client/src/diracx/client/_generated/aio/operations/_operations.py +++ b/diracx-client/src/diracx/client/_generated/aio/operations/_operations.py @@ -51,6 +51,10 @@ build_jobs_summary_request, build_jobs_unassign_bulk_jobs_sandboxes_request, build_jobs_unassign_job_sandboxes_request, + build_pilots_register_pilot_request, + build_pilots_search_request, + build_pilots_summary_request, + build_pilots_update_pilot_metadata_request, build_rss_get_compute_status_request, build_rss_get_fts_status_request, build_rss_get_site_status_request, @@ -1957,6 +1961,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: ~_generated.models.SearchParams :keyword page: Default value is 1. @@ -1997,6 +2007,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: IO[bytes] :keyword page: Default value is 1. @@ -2036,6 +2052,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. :type body: ~_generated.models.SearchParams or IO[bytes] :keyword page: Default value is 1. @@ -2313,7 +2335,7 @@ async def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any response = pipeline_response.http_response - if response.status_code not in [200]: + if response.status_code not in [201]: map_error(status_code=response.status_code, response=response, error_map=error_map) raise HttpResponseError(response=response) @@ -2325,6 +2347,500 @@ async def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any return deserialized # type: ignore +class PilotsOperations: + """ + .. warning:: + **DO NOT** instantiate this class directly. + + Instead, you should access the following operations through + :class:`~_generated.aio.Dirac`'s + :attr:`pilots` attribute. + """ + + models = _models + + def __init__(self, *args, **kwargs) -> None: + input_args = list(args) + self._client: AsyncPipelineClient = input_args.pop(0) if input_args else kwargs.pop("client") + self._config: DiracConfiguration = input_args.pop(0) if input_args else kwargs.pop("config") + self._serialize: Serializer = input_args.pop(0) if input_args else kwargs.pop("serializer") + self._deserialize: Deserializer = input_args.pop(0) if input_args else kwargs.pop("deserializer") + + @overload + async def register_pilot( + self, body: _models.PilotRegistrationParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Required. + :type body: ~_generated.models.PilotRegistrationParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def register_pilot(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def register_pilot(self, body: Union[_models.PilotRegistrationParams, IO[bytes]], **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Is either a PilotRegistrationParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.PilotRegistrationParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "PilotRegistrationParams") + + _request = build_pilots_register_pilot_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 [201]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + @overload + async def update_pilot_metadata( + self, body: dict[str, _models.PilotMetadata], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Required. + :type body: dict[str, ~_generated.models.PilotMetadata] + :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 update_pilot_metadata( + self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :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 update_pilot_metadata( + self, body: Union[dict[str, _models.PilotMetadata], IO[bytes]], **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Is either a {str: PilotMetadata} type or a IO[bytes] type. Required. + :type body: dict[str, ~_generated.models.PilotMetadata] 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, "{PilotMetadata}") + + _request = build_pilots_update_pilot_metadata_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 + + @overload + async def search( + self, + body: Optional[_models.SearchParams] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: ~_generated.models.SearchParams + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def search( + self, + body: Optional[IO[bytes]] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def search( + self, + body: Optional[Union[_models.SearchParams, IO[bytes]]] = None, + *, + page: int = 1, + per_page: int = 100, + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. + :type body: ~_generated.models.SearchParams or IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :return: list of dict mapping str to any + :rtype: list[dict[str, any]] + :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)) + content_type = content_type if body else None + cls: ClsType[list[dict[str, Any]]] = kwargs.pop("cls", None) + + content_type = content_type or "application/json" if body else None + _json = None + _content = None + if isinstance(body, (IOBase, bytes)): + _content = body + else: + if body is not None: + _json = self._serialize.body(body, "SearchParams") + else: + _json = None + + _request = build_pilots_search_request( + page=page, + per_page=per_page, + 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, 206]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + response_headers = {} + if response.status_code == 206: + response_headers["Content-Range"] = self._deserialize("str", response.headers.get("Content-Range")) + + deserialized = self._deserialize("[{object}]", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + async def summary( + self, body: _models.SummaryParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Required. + :type body: ~_generated.models.SummaryParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def summary(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def summary(self, body: Union[_models.SummaryParams, IO[bytes]], **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Is either a SummaryParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.SummaryParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "SummaryParams") + + _request = build_pilots_summary_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 [200]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + class RssOperations: """ .. warning:: diff --git a/diracx-client/src/diracx/client/_generated/aio/operations/_patch.py b/diracx-client/src/diracx/client/_generated/aio/operations/_patch.py index a408e57d2..c950258bc 100644 --- a/diracx-client/src/diracx/client/_generated/aio/operations/_patch.py +++ b/diracx-client/src/diracx/client/_generated/aio/operations/_patch.py @@ -6,15 +6,18 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ + from __future__ import annotations __all__ = [ "AuthOperations", "JobsOperations", + "PilotsOperations", ] # Add all objects you want publicly available to users at this package level from ....patches.auth.aio import AuthOperations from ....patches.jobs.aio import JobsOperations +from ....patches.pilots.aio import PilotsOperations def patch_sdk(): diff --git a/diracx-client/src/diracx/client/_generated/models/__init__.py b/diracx-client/src/diracx/client/_generated/models/__init__.py index b4c06cc69..7b7fb1792 100644 --- a/diracx-client/src/diracx/client/_generated/models/__init__.py +++ b/diracx-client/src/diracx/client/_generated/models/__init__.py @@ -32,6 +32,8 @@ JobStatusUpdate, Metadata, OpenIDConfiguration, + PilotMetadata, + PilotRegistrationParams, SandboxDownloadResponse, SandboxInfo, SandboxUploadResponse, @@ -61,6 +63,7 @@ from ._enums import ( # type: ignore ChecksumAlgorithm, JobStatus, + PilotStatus, SandboxFormat, SandboxType, ScalarSearchOperator, @@ -92,6 +95,8 @@ "JobStatusUpdate", "Metadata", "OpenIDConfiguration", + "PilotMetadata", + "PilotRegistrationParams", "SandboxDownloadResponse", "SandboxInfo", "SandboxUploadResponse", @@ -118,6 +123,7 @@ "VectorSearchSpec", "ChecksumAlgorithm", "JobStatus", + "PilotStatus", "SandboxFormat", "SandboxType", "ScalarSearchOperator", diff --git a/diracx-client/src/diracx/client/_generated/models/_enums.py b/diracx-client/src/diracx/client/_generated/models/_enums.py index b83473639..849d3252a 100644 --- a/diracx-client/src/diracx/client/_generated/models/_enums.py +++ b/diracx-client/src/diracx/client/_generated/models/_enums.py @@ -34,6 +34,19 @@ class JobStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta): RESCHEDULED = "Rescheduled" +class PilotStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """PilotStatus.""" + + SUBMITTED = "Submitted" + WAITING = "Waiting" + RUNNING = "Running" + DONE = "Done" + FAILED = "Failed" + DELETED = "Deleted" + ABORTED = "Aborted" + UNKNOWN = "Unknown" + + class SandboxFormat(str, Enum, metaclass=CaseInsensitiveEnumMeta): """SandboxFormat.""" diff --git a/diracx-client/src/diracx/client/_generated/models/_models.py b/diracx-client/src/diracx/client/_generated/models/_models.py index 730f15d6d..0245f0415 100644 --- a/diracx-client/src/diracx/client/_generated/models/_models.py +++ b/diracx-client/src/diracx/client/_generated/models/_models.py @@ -1053,6 +1053,185 @@ def __init__( self.code_challenge_methods_supported = code_challenge_methods_supported +class PilotMetadata(_serialization.Model): + """Mutable metadata attached to a pilot. + + The pilot is identified by its stamp, passed alongside this model + (e.g. as the mapping key on ``PATCH /api/pilots/metadata``\\ ). Every + field is optional; when absent it is left untouched by an update. + + :ivar status_reason: Human-readable reason for the current status. + :vartype status_reason: str + :ivar status: Current pilot status. Known values are: "Submitted", "Waiting", "Running", + "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :vartype status: str or ~_generated.models.PilotStatus + :ivar bench_mark: Pilot benchmark value. + :vartype bench_mark: float + :ivar destination_site: Destination site. + :vartype destination_site: str + :ivar queue: Batch queue name. + :vartype queue: str + :ivar grid_site: Grid site. + :vartype grid_site: str + :ivar grid_type: Grid type. + :vartype grid_type: str + :ivar accounting_sent: Whether accounting has been sent for this pilot. + :vartype accounting_sent: bool + :ivar current_job_id: ID of the job currently running on this pilot. + :vartype current_job_id: int + """ + + _validation = { + "status_reason": {"max_length": 255}, + "destination_site": {"max_length": 128}, + "queue": {"max_length": 128}, + "grid_site": {"max_length": 128}, + "grid_type": {"max_length": 32}, + } + + _attribute_map = { + "status_reason": {"key": "StatusReason", "type": "str"}, + "status": {"key": "Status", "type": "str"}, + "bench_mark": {"key": "BenchMark", "type": "float"}, + "destination_site": {"key": "DestinationSite", "type": "str"}, + "queue": {"key": "Queue", "type": "str"}, + "grid_site": {"key": "GridSite", "type": "str"}, + "grid_type": {"key": "GridType", "type": "str"}, + "accounting_sent": {"key": "AccountingSent", "type": "bool"}, + "current_job_id": {"key": "CurrentJobID", "type": "int"}, + } + + def __init__( + self, + *, + status_reason: Optional[str] = None, + status: Optional[Union[str, "_models.PilotStatus"]] = None, + bench_mark: Optional[float] = None, + destination_site: Optional[str] = None, + queue: Optional[str] = None, + grid_site: Optional[str] = None, + grid_type: Optional[str] = None, + accounting_sent: Optional[bool] = None, + current_job_id: Optional[int] = None, + **kwargs: Any + ) -> None: + """ + :keyword status_reason: Human-readable reason for the current status. + :paramtype status_reason: str + :keyword status: Current pilot status. Known values are: "Submitted", "Waiting", "Running", + "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :paramtype status: str or ~_generated.models.PilotStatus + :keyword bench_mark: Pilot benchmark value. + :paramtype bench_mark: float + :keyword destination_site: Destination site. + :paramtype destination_site: str + :keyword queue: Batch queue name. + :paramtype queue: str + :keyword grid_site: Grid site. + :paramtype grid_site: str + :keyword grid_type: Grid type. + :paramtype grid_type: str + :keyword accounting_sent: Whether accounting has been sent for this pilot. + :paramtype accounting_sent: bool + :keyword current_job_id: ID of the job currently running on this pilot. + :paramtype current_job_id: int + """ + super().__init__(**kwargs) + self.status_reason = status_reason + self.status = status + self.bench_mark = bench_mark + self.destination_site = destination_site + self.queue = queue + self.grid_site = grid_site + self.grid_type = grid_type + self.accounting_sent = accounting_sent + self.current_job_id = current_job_id + + +class PilotRegistrationParams(_serialization.Model): + """Body of ``POST /api/pilots/`` to register a single pilot. + + The ``max_length`` constraints mirror the column sizes of the legacy + ``PilotAgents`` table so that oversized values are rejected with a 422 + instead of a backend-dependent database error. + + All required parameters must be populated in order to send to server. + + :ivar pilot_stamp: Stamp of the pilot to create. Required. + :vartype pilot_stamp: str + :ivar vo: Pilot virtual organization. Required. + :vartype vo: str + :ivar grid_type: Grid type of the pilot. + :vartype grid_type: str + :ivar grid_site: Pilot grid site. + :vartype grid_site: str + :ivar destination_site: Pilot destination site. + :vartype destination_site: str + :ivar pilot_reference: CE job reference of the pilot; defaults to the stamp. + :vartype pilot_reference: str + :ivar pilot_status: Initial status of the pilot. Known values are: "Submitted", "Waiting", + "Running", "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :vartype pilot_status: str or ~_generated.models.PilotStatus + """ + + _validation = { + "pilot_stamp": {"required": True, "max_length": 32, "min_length": 1}, + "vo": {"required": True, "max_length": 128, "min_length": 1}, + "grid_type": {"max_length": 32}, + "grid_site": {"max_length": 128}, + "destination_site": {"max_length": 128}, + "pilot_reference": {"max_length": 255}, + } + + _attribute_map = { + "pilot_stamp": {"key": "pilot_stamp", "type": "str"}, + "vo": {"key": "vo", "type": "str"}, + "grid_type": {"key": "grid_type", "type": "str"}, + "grid_site": {"key": "grid_site", "type": "str"}, + "destination_site": {"key": "destination_site", "type": "str"}, + "pilot_reference": {"key": "pilot_reference", "type": "str"}, + "pilot_status": {"key": "pilot_status", "type": "str"}, + } + + def __init__( + self, + *, + pilot_stamp: str, + vo: str, + grid_type: str = "DIRAC", + grid_site: str = "Unknown", + destination_site: str = "NotAssigned", + pilot_reference: Optional[str] = None, + pilot_status: Optional[Union[str, "_models.PilotStatus"]] = None, + **kwargs: Any + ) -> None: + """ + :keyword pilot_stamp: Stamp of the pilot to create. Required. + :paramtype pilot_stamp: str + :keyword vo: Pilot virtual organization. Required. + :paramtype vo: str + :keyword grid_type: Grid type of the pilot. + :paramtype grid_type: str + :keyword grid_site: Pilot grid site. + :paramtype grid_site: str + :keyword destination_site: Pilot destination site. + :paramtype destination_site: str + :keyword pilot_reference: CE job reference of the pilot; defaults to the stamp. + :paramtype pilot_reference: str + :keyword pilot_status: Initial status of the pilot. Known values are: "Submitted", "Waiting", + "Running", "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :paramtype pilot_status: str or ~_generated.models.PilotStatus + """ + super().__init__(**kwargs) + self.pilot_stamp = pilot_stamp + self.vo = vo + self.grid_type = grid_type + self.grid_site = grid_site + self.destination_site = destination_site + self.pilot_reference = pilot_reference + self.pilot_status = pilot_status + + class SandboxDownloadResponse(_serialization.Model): """SandboxDownloadResponse. diff --git a/diracx-client/src/diracx/client/_generated/operations/__init__.py b/diracx-client/src/diracx/client/_generated/operations/__init__.py index 77674abec..48074896e 100644 --- a/diracx-client/src/diracx/client/_generated/operations/__init__.py +++ b/diracx-client/src/diracx/client/_generated/operations/__init__.py @@ -14,6 +14,7 @@ from ._operations import AuthOperations # type: ignore from ._operations import ConfigOperations # type: ignore from ._operations import JobsOperations # type: ignore +from ._operations import PilotsOperations # type: ignore from ._operations import RssOperations # type: ignore from ._patch import __all__ as _patch_all @@ -25,6 +26,7 @@ "AuthOperations", "ConfigOperations", "JobsOperations", + "PilotsOperations", "RssOperations", ] __all__.extend([p for p in _patch_all if p not in __all__]) # pyright: ignore diff --git a/diracx-client/src/diracx/client/_generated/operations/_operations.py b/diracx-client/src/diracx/client/_generated/operations/_operations.py index 69089682a..077560861 100644 --- a/diracx-client/src/diracx/client/_generated/operations/_operations.py +++ b/diracx-client/src/diracx/client/_generated/operations/_operations.py @@ -565,6 +565,78 @@ def build_jobs_submit_jdl_jobs_request(**kwargs: Any) -> HttpRequest: return HttpRequest(method="POST", url=_url, headers=_headers, **kwargs) +def build_pilots_register_pilot_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/pilots/" + + # 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="POST", url=_url, headers=_headers, **kwargs) + + +def build_pilots_update_pilot_metadata_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/metadata" + + # 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) + + +def build_pilots_search_request(*, page: int = 1, per_page: int = 100, **kwargs: Any) -> HttpRequest: + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = case_insensitive_dict(kwargs.pop("params", {}) or {}) + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + accept = _headers.pop("Accept", "application/json") + + # Construct URL + _url = "/api/pilots/search" + + # Construct parameters + if page is not None: + _params["page"] = _SERIALIZER.query("page", page, "int", minimum=1) + if per_page is not None: + _params["per_page"] = _SERIALIZER.query("per_page", per_page, "int", maximum=10000, minimum=1) + + # 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="POST", url=_url, params=_params, headers=_headers, **kwargs) + + +def build_pilots_summary_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/pilots/summary" + + # 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="POST", url=_url, headers=_headers, **kwargs) + + def build_rss_get_storage_status_request( *, if_modified_since: Optional[str] = None, @@ -2566,6 +2638,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: ~_generated.models.SearchParams :keyword page: Default value is 1. @@ -2606,6 +2684,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: IO[bytes] :keyword page: Default value is 1. @@ -2645,6 +2729,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. :type body: ~_generated.models.SearchParams or IO[bytes] :keyword page: Default value is 1. @@ -2920,7 +3010,7 @@ def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any) -> l response = pipeline_response.http_response - if response.status_code not in [200]: + if response.status_code not in [201]: map_error(status_code=response.status_code, response=response, error_map=error_map) raise HttpResponseError(response=response) @@ -2932,6 +3022,496 @@ def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any) -> l return deserialized # type: ignore +class PilotsOperations: + """ + .. warning:: + **DO NOT** instantiate this class directly. + + Instead, you should access the following operations through + :class:`~_generated.Dirac`'s + :attr:`pilots` attribute. + """ + + models = _models + + def __init__(self, *args, **kwargs) -> None: + input_args = list(args) + self._client: PipelineClient = input_args.pop(0) if input_args else kwargs.pop("client") + self._config: DiracConfiguration = input_args.pop(0) if input_args else kwargs.pop("config") + self._serialize: Serializer = input_args.pop(0) if input_args else kwargs.pop("serializer") + self._deserialize: Deserializer = input_args.pop(0) if input_args else kwargs.pop("deserializer") + + @overload + def register_pilot( + self, body: _models.PilotRegistrationParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Required. + :type body: ~_generated.models.PilotRegistrationParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def register_pilot(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def register_pilot(self, body: Union[_models.PilotRegistrationParams, IO[bytes]], **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Is either a PilotRegistrationParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.PilotRegistrationParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "PilotRegistrationParams") + + _request = build_pilots_register_pilot_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 [201]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + @overload + def update_pilot_metadata( + self, body: dict[str, _models.PilotMetadata], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Required. + :type body: dict[str, ~_generated.models.PilotMetadata] + :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 update_pilot_metadata(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :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 update_pilot_metadata( # pylint: disable=inconsistent-return-statements + self, body: Union[dict[str, _models.PilotMetadata], IO[bytes]], **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Is either a {str: PilotMetadata} type or a IO[bytes] type. Required. + :type body: dict[str, ~_generated.models.PilotMetadata] 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, "{PilotMetadata}") + + _request = build_pilots_update_pilot_metadata_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 + + @overload + def search( + self, + body: Optional[_models.SearchParams] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: ~_generated.models.SearchParams + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def search( + self, + body: Optional[IO[bytes]] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def search( + self, + body: Optional[Union[_models.SearchParams, IO[bytes]]] = None, + *, + page: int = 1, + per_page: int = 100, + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. + :type body: ~_generated.models.SearchParams or IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :return: list of dict mapping str to any + :rtype: list[dict[str, any]] + :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)) + content_type = content_type if body else None + cls: ClsType[list[dict[str, Any]]] = kwargs.pop("cls", None) + + content_type = content_type or "application/json" if body else None + _json = None + _content = None + if isinstance(body, (IOBase, bytes)): + _content = body + else: + if body is not None: + _json = self._serialize.body(body, "SearchParams") + else: + _json = None + + _request = build_pilots_search_request( + page=page, + per_page=per_page, + 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, 206]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + response_headers = {} + if response.status_code == 206: + response_headers["Content-Range"] = self._deserialize("str", response.headers.get("Content-Range")) + + deserialized = self._deserialize("[{object}]", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + def summary(self, body: _models.SummaryParams, *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Required. + :type body: ~_generated.models.SummaryParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def summary(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def summary(self, body: Union[_models.SummaryParams, IO[bytes]], **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Is either a SummaryParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.SummaryParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "SummaryParams") + + _request = build_pilots_summary_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("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + class RssOperations: """ .. warning:: diff --git a/diracx-client/src/diracx/client/_generated/operations/_patch.py b/diracx-client/src/diracx/client/_generated/operations/_patch.py index b7b8c67fa..dd8ba9d0b 100644 --- a/diracx-client/src/diracx/client/_generated/operations/_patch.py +++ b/diracx-client/src/diracx/client/_generated/operations/_patch.py @@ -6,15 +6,18 @@ Follow our quickstart for examples: https://aka.ms/azsdk/python/dpcodegen/python/customize """ + from __future__ import annotations __all__ = [ "AuthOperations", "JobsOperations", + "PilotsOperations", ] # 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 +from ...patches.pilots.sync import PilotsOperations def patch_sdk(): diff --git a/diracx-client/src/diracx/client/patches/pilots/aio.py b/diracx-client/src/diracx/client/patches/pilots/aio.py new file mode 100644 index 000000000..622ba2749 --- /dev/null +++ b/diracx-client/src/diracx/client/patches/pilots/aio.py @@ -0,0 +1,49 @@ +"""Patches for the autorest-generated async pilots client.""" + +from __future__ import annotations + +__all__ = [ + "PilotsOperations", +] + +from typing import Any, Unpack + +from azure.core.tracing.decorator_async import distributed_trace_async + +from ..._generated.aio.operations._operations import ( + PilotsOperations as _PilotsOperations, +) +from .common import ( + RegisterPilotsKwargs, + SearchKwargs, + SummaryKwargs, + make_register_pilots_body, + make_search_body, + make_summary_body, +) + +# We're intentionally ignoring overrides here because we want to change the interface. +# mypy: disable-error-code=override + + +class PilotsOperations(_PilotsOperations): + @distributed_trace_async + async def search( + self, **kwargs: Unpack[SearchKwargs] + ) -> list[dict[str, Any]]: + """Search for pilots matching the provided filters.""" + return await super().search(**make_search_body(**kwargs)) + + @distributed_trace_async + async def summary( + self, **kwargs: Unpack[SummaryKwargs] + ) -> list[dict[str, Any]]: + """Return pilot counts aggregated by the requested columns.""" + return await super().summary(**make_summary_body(**kwargs)) + + @distributed_trace_async + async def register_pilot( + self, **kwargs: Unpack[RegisterPilotsKwargs] + ) -> None: + """Register a pilot.""" + return await super().register_pilot(**make_register_pilots_body(**kwargs)) diff --git a/diracx-client/src/diracx/client/patches/pilots/common.py b/diracx-client/src/diracx/client/patches/pilots/common.py new file mode 100644 index 000000000..d0de5b852 --- /dev/null +++ b/diracx-client/src/diracx/client/patches/pilots/common.py @@ -0,0 +1,145 @@ +"""Utilities shared by the sync and async pilots operator patches.""" + +from __future__ import annotations + +__all__ = [ + "SearchKwargs", + "make_search_body", + "SummaryKwargs", + "make_summary_body", + "RegisterPilotsKwargs", + "make_register_pilots_body", +] + +import json +from io import BytesIO +from typing import IO, Any, Literal, TypedDict, Unpack, cast + +from diracx.core.models.pilot import PilotStatus +from diracx.core.models.search import SearchSpec + + +class ResponseExtra(TypedDict, total=False): + content_type: str + headers: dict[str, str] + params: dict[str, str] + cls: Any + + +# ------------------ Search ------------------ + + +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 operation expects IO[bytes] despite its + # signature advertising 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 + + +# ------------------ Summary ------------------ + + +class SummaryBody(TypedDict, total=False): + grouping: list[str] + search: list[SearchSpec] + + +class SummaryKwargs(SummaryBody, ResponseExtra): ... + + +class UnderlyingSummaryArgs(ResponseExtra, total=False): + # FIXME: The autorest-generated operation expects IO[bytes] despite its + # signature advertising 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 + + +# ------------------ Register pilots ------------------ + + +class RegisterPilotsBody(TypedDict, total=False): + pilot_stamp: str + vo: str + grid_type: str + grid_site: str + destination_site: str + pilot_reference: str + pilot_status: PilotStatus + + +class RegisterPilotsKwargs(RegisterPilotsBody, ResponseExtra): ... + + +class UnderlyingRegisterPilotsArgs(ResponseExtra, total=False): + # FIXME: The autorest-generated operation expects IO[bytes] despite its + # signature advertising IO[bytes] | bytes. + body: IO[bytes] + + +def make_register_pilots_body( + **kwargs: Unpack[RegisterPilotsKwargs], +) -> UnderlyingRegisterPilotsArgs: + body: RegisterPilotsBody = {} + for key in RegisterPilotsBody.__optional_keys__: + if key not in kwargs: + continue + key = cast( + Literal[ + "pilot_stamp", + "vo", + "grid_type", + "grid_site", + "destination_site", + "pilot_reference", + "pilot_status", + ], + key, + ) + value = kwargs.pop(key) + if value is not None: + body[key] = value + result: UnderlyingRegisterPilotsArgs = { + "body": BytesIO(json.dumps(body).encode("utf-8")) + } + result.update(cast(ResponseExtra, kwargs)) + return result diff --git a/diracx-client/src/diracx/client/patches/pilots/sync.py b/diracx-client/src/diracx/client/patches/pilots/sync.py new file mode 100644 index 000000000..4012e4cc5 --- /dev/null +++ b/diracx-client/src/diracx/client/patches/pilots/sync.py @@ -0,0 +1,41 @@ +"""Patches for the autorest-generated sync pilots client.""" + +from __future__ import annotations + +__all__ = [ + "PilotsOperations", +] + +from typing import Any, Unpack + +from azure.core.tracing.decorator import distributed_trace + +from ..._generated.operations._operations import PilotsOperations as _PilotsOperations +from .common import ( + RegisterPilotsKwargs, + SearchKwargs, + SummaryKwargs, + make_register_pilots_body, + make_search_body, + make_summary_body, +) + +# We're intentionally ignoring overrides here because we want to change the interface. +# mypy: disable-error-code=override + + +class PilotsOperations(_PilotsOperations): + @distributed_trace + def search(self, **kwargs: Unpack[SearchKwargs]) -> list[dict[str, Any]]: + """Search for pilots matching the provided filters.""" + return super().search(**make_search_body(**kwargs)) + + @distributed_trace + def summary(self, **kwargs: Unpack[SummaryKwargs]) -> list[dict[str, Any]]: + """Return pilot counts aggregated by the requested columns.""" + return super().summary(**make_summary_body(**kwargs)) + + @distributed_trace + def register_pilot(self, **kwargs: Unpack[RegisterPilotsKwargs]) -> None: + """Register a pilot.""" + return super().register_pilot(**make_register_pilots_body(**kwargs)) diff --git a/diracx-core/src/diracx/core/exceptions.py b/diracx-core/src/diracx/core/exceptions.py index 889752c50..4ac333787 100644 --- a/diracx-core/src/diracx/core/exceptions.py +++ b/diracx-core/src/diracx/core/exceptions.py @@ -9,6 +9,9 @@ "InvalidQueryError", "NotReadyError", "PendingAuthorizationError", + "PilotAlreadyAssociatedWithJobError", + "PilotAlreadyExistsError", + "PilotNotFoundError", "SandboxAlreadyAssignedError", "SandboxAlreadyInsertedError", "SandboxNotFoundError", @@ -19,6 +22,7 @@ class DiracError(RuntimeError): def __init__(self, detail: str = "Unknown"): self.detail = detail + super().__init__(detail) class AuthorizationError(DiracError): ... @@ -53,19 +57,19 @@ class InvalidQueryError(DiracError): class TokenNotFoundError(DiracError): - def __init__(self, jti: str, detail: str | None = None): + def __init__(self, jti: str, detail: str = ""): self.jti: str = jti super().__init__(f"Token {jti} not found" + (f" ({detail})" if detail else "")) class JobNotFoundError(DiracError): - def __init__(self, job_id: int, detail: str | None = None): + def __init__(self, job_id: int, detail: str = ""): self.job_id: int = job_id super().__init__(f"Job {job_id} not found" + (f" ({detail})" if detail else "")) class SandboxNotFoundError(DiracError): - def __init__(self, pfn: str, se_name: str, detail: str | None = None): + def __init__(self, pfn: str, se_name: str, detail: str = ""): self.pfn: str = pfn self.se_name: str = se_name super().__init__( @@ -81,7 +85,7 @@ def __init__(self, name: str, detail: str | None = None): class SandboxAlreadyAssignedError(DiracError): - def __init__(self, pfn: str, se_name: str, detail: str | None = None): + def __init__(self, pfn: str, se_name: str, detail: str = ""): self.pfn: str = pfn self.se_name: str = se_name super().__init__( @@ -91,7 +95,7 @@ def __init__(self, pfn: str, se_name: str, detail: str | None = None): class SandboxAlreadyInsertedError(DiracError): - def __init__(self, pfn: str, se_name: str, detail: str | None = None): + def __init__(self, pfn: str, se_name: str, detail: str = ""): self.pfn: str = pfn self.se_name: str = se_name super().__init__( @@ -101,7 +105,7 @@ def __init__(self, pfn: str, se_name: str, detail: str | None = None): class JobError(DiracError): - def __init__(self, job_id, detail: str | None = None): + def __init__(self, job_id, detail: str = ""): self.job_id: int = job_id super().__init__( f"Error concerning job {job_id}" + (f" ({detail})" if detail else "") @@ -110,3 +114,15 @@ def __init__(self, job_id, detail: str | None = None): class NotReadyError(DiracError): """Tried to access a value which is asynchronously loaded but not yet available.""" + + +class PilotNotFoundError(DiracError): + """At least one pilot is not found.""" + + +class PilotAlreadyExistsError(DiracError): + """At least one pilot already exists, we avoid collisions.""" + + +class PilotAlreadyAssociatedWithJobError(DiracError): + """We can't associate a pilot with the same job twice.""" diff --git a/diracx-core/src/diracx/core/models/pilot.py b/diracx-core/src/diracx/core/models/pilot.py new file mode 100644 index 000000000..a325ca044 --- /dev/null +++ b/diracx-core/src/diracx/core/models/pilot.py @@ -0,0 +1,93 @@ +"""Pilot-related models shared between client, logic, and services.""" + +from __future__ import annotations + +from enum import StrEnum + +from pydantic import BaseModel, Field + + +class PilotStatus(StrEnum): + SUBMITTED = "Submitted" + WAITING = "Waiting" + RUNNING = "Running" + DONE = "Done" + FAILED = "Failed" + DELETED = "Deleted" + ABORTED = "Aborted" + UNKNOWN = "Unknown" + + +class PilotRegistrationParams(BaseModel, extra="forbid"): + """Body of ``POST /api/pilots/`` to register a single pilot. + + The ``max_length`` constraints mirror the column sizes of the legacy + ``PilotAgents`` table so that oversized values are rejected with a 422 + instead of a backend-dependent database error. + """ + + pilot_stamp: str = Field( + min_length=1, max_length=32, description="Stamp of the pilot to create." + ) + vo: str = Field( + min_length=1, max_length=128, description="Pilot virtual organization." + ) + grid_type: str = Field( + "DIRAC", max_length=32, description="Grid type of the pilot." + ) + grid_site: str = Field("Unknown", max_length=128, description="Pilot grid site.") + destination_site: str = Field( + "NotAssigned", max_length=128, description="Pilot destination site." + ) + pilot_reference: str | None = Field( + None, + max_length=255, + description="CE job reference of the pilot; defaults to the stamp.", + ) + pilot_status: PilotStatus = Field( + PilotStatus.SUBMITTED, description="Initial status of the pilot." + ) + + +class PilotMetadata(BaseModel, populate_by_name=True, extra="forbid"): + """Mutable metadata attached to a pilot. + + The pilot is identified by its stamp, passed alongside this model + (e.g. as the mapping key on ``PATCH /api/pilots/metadata``). Every + field is optional; when absent it is left untouched by an update. + """ + + status_reason: str | None = Field( + None, + alias="StatusReason", + max_length=255, + description="Human-readable reason for the current status.", + ) + status: PilotStatus | None = Field( + None, alias="Status", description="Current pilot status." + ) + benchmark: float | None = Field( + None, alias="BenchMark", description="Pilot benchmark value." + ) + destination_site: str | None = Field( + None, alias="DestinationSite", max_length=128, description="Destination site." + ) + queue: str | None = Field( + None, alias="Queue", max_length=128, description="Batch queue name." + ) + grid_site: str | None = Field( + None, alias="GridSite", max_length=128, description="Grid site." + ) + grid_type: str | None = Field( + None, alias="GridType", max_length=32, description="Grid type." + ) + accounting_sent: bool | None = Field( + None, + alias="AccountingSent", + description="Whether accounting has been sent for this pilot.", + ) + current_job_id: int | None = Field( + None, + alias="CurrentJobID", + description="ID of the job currently running on this pilot.", + ) diff --git a/diracx-core/tests/test_config_source.py b/diracx-core/tests/test_config_source.py index e1a095aff..99c9265c2 100644 --- a/diracx-core/tests/test_config_source.py +++ b/diracx-core/tests/test_config_source.py @@ -41,12 +41,12 @@ def default_remote_conf_assertions(monkeypatch, repo_url): return hexsha -@pytest.mark.skipif(github_is_down(), reason="Github unavailble") +@pytest.mark.skipif(github_is_down(), reason="Github unavailable") def test_remote_git_config_source_default(monkeypatch): default_remote_conf_assertions(monkeypatch, TEST_REPO) -@pytest.mark.skipif(github_is_down(), reason="Github unavailble") +@pytest.mark.skipif(github_is_down(), reason="Github unavailable") def test_remote_git_config_source_branch(monkeypatch): monkeypatch.setattr( "diracx.core.config.sources.DEFAULT_GIT_BRANCH", @@ -55,7 +55,7 @@ def test_remote_git_config_source_branch(monkeypatch): default_remote_conf_assertions(monkeypatch, TEST_REPO_SPECIFIC_BRANCH) -@pytest.mark.skipif(github_is_down(), reason="Github unavailble") +@pytest.mark.skipif(github_is_down(), reason="Github unavailable") def test_remote_git_config_source_commit(monkeypatch): hexsha = default_remote_conf_assertions(monkeypatch, TEST_REPO_SPECIFIC_COMMIT_HASH) assert hexsha == COMMIT_HASH diff --git a/diracx-db/src/diracx/db/sql/__init__.py b/diracx-db/src/diracx/db/sql/__init__.py index 704dce3bc..de697c843 100644 --- a/diracx-db/src/diracx/db/sql/__init__.py +++ b/diracx-db/src/diracx/db/sql/__init__.py @@ -13,7 +13,7 @@ from .auth.db import AuthDB from .job.db import JobDB from .job_logging.db import JobLoggingDB -from .pilot_agents.db import PilotAgentsDB +from .pilots.db import PilotAgentsDB from .rss.db import ResourceStatusDB from .sandbox_metadata.db import SandboxMetadataDB from .task_queue.db import TaskQueueDB diff --git a/diracx-db/src/diracx/db/sql/dummy/db.py b/diracx-db/src/diracx/db/sql/dummy/db.py index 76e8db07b..e94f7074b 100644 --- a/diracx-db/src/diracx/db/sql/dummy/db.py +++ b/diracx-db/src/diracx/db/sql/dummy/db.py @@ -3,6 +3,7 @@ from sqlalchemy import insert from uuid_utils import UUID +from diracx.core.models.search import SearchSpec from diracx.db.sql.utils import BaseSQLDB from .schema import Base as DummyDBBase @@ -20,8 +21,11 @@ class DummyDB(BaseSQLDB): # This needs to be here for the BaseSQLDB to create the engine metadata = DummyDBBase.metadata - async def summary(self, group_by, search) -> list[dict[str, str | int]]: - return await self._summary(Cars, group_by, search) + async def summary( + self, group_by: list[str], search: list[SearchSpec] + ) -> list[dict[str, str | int]]: + """Get a summary of the cars.""" + return await self._summary(table=Cars, group_by=group_by, search=search) async def insert_owner(self, name: str) -> int: stmt = insert(Owners).values(name=name) diff --git a/diracx-db/src/diracx/db/sql/job/db.py b/diracx-db/src/diracx/db/sql/job/db.py index 94ba4df63..0c7d429a6 100644 --- a/diracx-db/src/diracx/db/sql/job/db.py +++ b/diracx-db/src/diracx/db/sql/job/db.py @@ -12,8 +12,7 @@ from diracx.core.exceptions import InvalidQueryError from diracx.core.models import JobCommand, SearchSpec, SortSpec -from ..utils import BaseSQLDB, _get_columns -from ..utils.functions import utcnow +from ..utils import BaseSQLDB, _get_columns, utcnow from .schema import ( HeartBeatLoggingInfo, InputData, diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/db.py b/diracx-db/src/diracx/db/sql/pilot_agents/db.py deleted file mode 100644 index 954f081b1..000000000 --- a/diracx-db/src/diracx/db/sql/pilot_agents/db.py +++ /dev/null @@ -1,45 +0,0 @@ -from __future__ import annotations - -from datetime import datetime, timezone - -from sqlalchemy import insert - -from ..utils import BaseSQLDB -from .schema import PilotAgents, PilotAgentsDBBase - - -class PilotAgentsDB(BaseSQLDB): - """PilotAgentsDB class is a front-end to the PilotAgents Database.""" - - metadata = PilotAgentsDBBase.metadata - - async def add_pilot_references( - self, - pilot_ref: list[str], - vo: str, - grid_type: str = "DIRAC", - pilot_stamps: dict | None = None, - ) -> None: - if pilot_stamps is None: - pilot_stamps = {} - - now = datetime.now(tz=timezone.utc) - - # Prepare the list of dictionaries for bulk insertion - values = [ - { - "PilotJobReference": ref, - "VO": vo, - "GridType": grid_type, - "SubmissionTime": now, - "LastUpdateTime": now, - "Status": "Submitted", - "PilotStamp": pilot_stamps.get(ref, ""), - } - for ref in pilot_ref - ] - - # Insert multiple rows in a single execute call - stmt = insert(PilotAgents).values(values) - await self.conn.execute(stmt) - return diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/__init__.py b/diracx-db/src/diracx/db/sql/pilots/__init__.py similarity index 100% rename from diracx-db/src/diracx/db/sql/pilot_agents/__init__.py rename to diracx-db/src/diracx/db/sql/pilots/__init__.py diff --git a/diracx-db/src/diracx/db/sql/pilots/db.py b/diracx-db/src/diracx/db/sql/pilots/db.py new file mode 100644 index 000000000..4e7deda97 --- /dev/null +++ b/diracx-db/src/diracx/db/sql/pilots/db.py @@ -0,0 +1,189 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from sqlalchemy import case, insert, literal, select, update +from sqlalchemy.exc import IntegrityError + +from diracx.core.exceptions import ( + PilotAlreadyAssociatedWithJobError, + PilotNotFoundError, +) +from diracx.core.models.pilot import PilotStatus +from diracx.core.models.search import SearchSpec, SortSpec + +from ..utils import BaseSQLDB +from .schema import ( + JobToPilotMapping, + PilotAgents, + PilotAgentsDBBase, +) + + +class PilotAgentsDB(BaseSQLDB): + """Front-end to the PilotAgents database.""" + + metadata = PilotAgentsDBBase.metadata + + async def register_pilots( + self, + pilot_stamps: list[str], + vo: str, + grid_type: str = "DIRAC", + grid_site: str = "Unknown", + destination_site: str = "NotAssigned", + pilot_references: dict[str, str] | None = None, + status: PilotStatus = PilotStatus.SUBMITTED, + ): + """Bulk-register pilots. + + If a stamp has no entry in `pilot_references` the stamp is used as + the reference. + """ + if pilot_references is None: + pilot_references = {} + + now = datetime.now(tz=timezone.utc) + + values = [ + { + "PilotJobReference": pilot_references.get(stamp, stamp), + "VO": vo, + "GridType": grid_type, + "GridSite": grid_site, + "DestinationSite": destination_site, + "SubmissionTime": now, + "LastUpdateTime": now, + "Status": status, + "PilotStamp": stamp, + } + for stamp in pilot_stamps + ] + + await self.conn.execute(insert(PilotAgents).values(values)) + + async def assign_jobs_to_pilot(self, job_to_pilot_mapping: list[dict[str, Any]]): + """Associate a pilot with jobs. + + Each entry has the shape `{"PilotID": ..., "JobID": ..., "StartTime": ...}`. + Raises PilotAlreadyAssociatedWithJobError on duplicates. The legacy + schema has no foreign key on JobToPilotMapping, so the caller must + ensure the pilots and jobs exist. + """ + stmt = insert(JobToPilotMapping).values(job_to_pilot_mapping) + + try: + await self.conn.execute(stmt) + except IntegrityError as e: + msg = str(e.orig).lower() + if "duplicate entry" in msg or "unique constraint" in msg: + raise PilotAlreadyAssociatedWithJobError( + detail=( + "at least one of these pilots is already associated " + "with a given job." + ) + ) from e + raise + + async def update_pilot_metadata(self, updates: dict[str, dict[str, Any]]): + """Bulk-update pilot fields. + + `updates` maps a pilot stamp to the column/value pairs to set for + that pilot; each entry may set a different subset of columns. Uses + a per-column CASE expression to support heterogeneous updates, + matching the pattern in JobDB.set_job_attributes. Raises + PilotNotFoundError if any of the pilot stamps is not found. + """ + if not updates: + return + + columns = {col for fields in updates.values() for col in fields} + + case_expressions = { + column: case( + *[ + ( + PilotAgents.__table__.c.PilotStamp == stamp, + literal( + fields[column], + type_=PilotAgents.__table__.c[column].type, + ), + ) + for stamp, fields in updates.items() + if column in fields + ], + else_=getattr(PilotAgents.__table__.c, column), + ) + for column in columns + } + + stmt = ( + update(PilotAgents) + .values(**case_expressions) + .where(PilotAgents.__table__.c.PilotStamp.in_(updates.keys())) + ) + result = await self.conn.execute(stmt) + + if result.rowcount != len(updates): + raise PilotNotFoundError("at least one of the given pilots does not exist.") + + async def search( + self, + parameters: list[str] | None, + search: list[SearchSpec], + sorts: list[SortSpec], + *, + distinct: bool = False, + per_page: int = 100, + page: int | None = None, + ) -> tuple[int, list[dict[str, Any]]]: + """Search for pilot information in the database.""" + return await self._search( + table=PilotAgents, + parameters=parameters, + search=search, + sorts=sorts, + distinct=distinct, + per_page=per_page, + page=page, + ) + + async def summary( + self, group_by: list[str], search: list[SearchSpec] + ) -> list[dict[str, str | int]]: + """Aggregate pilot counts by the requested columns.""" + return await self._summary(table=PilotAgents, group_by=group_by, search=search) + + async def get_job_ids_for_stamps(self, pilot_stamps: list[str]) -> list[int]: + """Return the IDs of jobs that have run on any of the given pilot stamps. + + Single round-trip SQL join over JobToPilotMapping and PilotAgents + (both live in the same metadata, so the join is legitimate at the + DB layer). + """ + if not pilot_stamps: + return [] + stmt = ( + select(JobToPilotMapping.job_id) + .join( + PilotAgents, + PilotAgents.pilot_id == JobToPilotMapping.pilot_id, + ) + .where(PilotAgents.pilot_stamp.in_(pilot_stamps)) + .distinct() + ) + result = await self.conn.execute(stmt) + return [row[0] for row in result] + + async def get_pilot_ids_for_job_ids(self, job_ids: list[int]) -> list[int]: + """Return the IDs of pilots that have run any of the given jobs.""" + if not job_ids: + return [] + stmt = ( + select(JobToPilotMapping.pilot_id) + .where(JobToPilotMapping.job_id.in_(job_ids)) + .distinct() + ) + result = await self.conn.execute(stmt) + return [row[0] for row in result] diff --git a/diracx-db/src/diracx/db/sql/pilot_agents/schema.py b/diracx-db/src/diracx/db/sql/pilots/schema.py similarity index 80% rename from diracx-db/src/diracx/db/sql/pilot_agents/schema.py rename to diracx-db/src/diracx/db/sql/pilots/schema.py index 770b62b79..5b1cd1713 100644 --- a/diracx-db/src/diracx/db/sql/pilot_agents/schema.py +++ b/diracx-db/src/diracx/db/sql/pilots/schema.py @@ -1,7 +1,6 @@ from __future__ import annotations from datetime import datetime -from typing import Optional from sqlalchemy import ( Double, @@ -11,6 +10,7 @@ ) from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column +from diracx.core.models.pilot import PilotStatus from diracx.db.sql.utils import ( EnumBackedBool, str32, @@ -48,22 +48,26 @@ class PilotAgents(PilotAgentsDBBase): vo: Mapped[str128] = mapped_column("VO") grid_type: Mapped[str32] = mapped_column("GridType", default="LCG") benchmark: Mapped[float] = mapped_column("BenchMark", Double, default=0.0) - submission_time: Mapped[Optional[datetime]] = mapped_column( + submission_time: Mapped[datetime | None] = mapped_column( "SubmissionTime", SmarterDateTime ) - last_update_time: Mapped[Optional[datetime]] = mapped_column( + last_update_time: Mapped[datetime | None] = mapped_column( "LastUpdateTime", SmarterDateTime ) - status: Mapped[str32] = mapped_column("Status", default="Unknown") + status: Mapped[str32] = mapped_column("Status", default=PilotStatus.UNKNOWN) status_reason: Mapped[str255] = mapped_column("StatusReason", default="Unknown") accounting_sent: Mapped[bool] = mapped_column( "AccountingSent", EnumBackedBool(), default=False ) - + # Index parity with the legacy DIRAC PilotAgentsDB schema: a PilotStamp + # index is deliberately NOT added here even though stamp-keyed lookups + # are the hot path — it must first be added on the DIRAC side (with an + # update-guide entry), like idx_dest_queue_status was in v9.1.0. __table_args__ = ( Index("PilotJobReference", "PilotJobReference"), Index("Status", "Status"), Index("Statuskey", "GridSite", "DestinationSite", "Status"), + Index("idx_dest_queue_status", "DestinationSite", "Queue", "Status"), ) diff --git a/diracx-db/tests/pilot_agents/test_pilot_agents_db.py b/diracx-db/tests/pilot_agents/test_pilot_agents_db.py deleted file mode 100644 index 3ca989885..000000000 --- a/diracx-db/tests/pilot_agents/test_pilot_agents_db.py +++ /dev/null @@ -1,30 +0,0 @@ -from __future__ import annotations - -import pytest - -from diracx.db.sql.pilot_agents.db import PilotAgentsDB - - -@pytest.fixture -async def pilot_agents_db(tmp_path) -> PilotAgentsDB: - agents_db = PilotAgentsDB("sqlite+aiosqlite:///:memory:") - async with agents_db.engine_context(): - async with agents_db.engine.begin() as conn: - await conn.run_sync(agents_db.metadata.create_all) - yield agents_db - - -async def test_insert_and_select(pilot_agents_db: PilotAgentsDB): - async with pilot_agents_db as pilot_agents_db: - # Add a pilot reference - refs = [f"ref_{i}" for i in range(10)] - stamps = [f"stamp_{i}" for i in range(10)] - stamp_dict = dict(zip(refs, stamps)) - - await pilot_agents_db.add_pilot_references( - refs, "test_vo", grid_type="DIRAC", pilot_stamps=stamp_dict - ) - - await pilot_agents_db.add_pilot_references( - refs, "test_vo", grid_type="DIRAC", pilot_stamps=None - ) diff --git a/diracx-db/tests/pilot_agents/__init__.py b/diracx-db/tests/pilots/__init__.py similarity index 100% rename from diracx-db/tests/pilot_agents/__init__.py rename to diracx-db/tests/pilots/__init__.py diff --git a/diracx-db/tests/pilots/test_pilot_db.py b/diracx-db/tests/pilots/test_pilot_db.py new file mode 100644 index 000000000..519383e11 --- /dev/null +++ b/diracx-db/tests/pilots/test_pilot_db.py @@ -0,0 +1,119 @@ +"""DB-level tests for PilotAgentsDB. + +Only functionality that has no HTTP route (the job-to-pilot mapping) is +tested here; everything else is covered by the router-level tests in +diracx-routers/tests/pilots/. +""" + +from __future__ import annotations + +import pytest + +from diracx.core.exceptions import PilotAlreadyAssociatedWithJobError +from diracx.core.models.search import ( + VectorSearchOperator, + VectorSearchSpec, +) +from diracx.db.sql.pilots.db import PilotAgentsDB + + +@pytest.fixture +async def pilot_db(): + db = PilotAgentsDB("sqlite+aiosqlite:///:memory:") + async with db.engine_context(): + async with db.engine.begin() as conn: + await conn.run_sync(db.metadata.create_all) + yield db + + +async def _get_pilot_id(pilot_db, stamp: str) -> int: + async with pilot_db as db: + _, rows = await db.search( + parameters=["PilotID"], + search=[ + VectorSearchSpec( + parameter="PilotStamp", + operator=VectorSearchOperator.IN, + values=[stamp], + ) + ], + sorts=[], + ) + return rows[0]["PilotID"] + + +async def test_assign_jobs_to_pilot(pilot_db): + async with pilot_db as db: + await db.register_pilots(pilot_stamps=["s1"], vo="lhcb") + + pilot_id = await _get_pilot_id(pilot_db, "s1") + + async with pilot_db as db: + await db.assign_jobs_to_pilot( + [ + {"PilotID": pilot_id, "JobID": 10, "StartTime": "2024-01-01T00:00:00Z"}, + {"PilotID": pilot_id, "JobID": 20, "StartTime": "2024-01-01T00:00:00Z"}, + ] + ) + + async with pilot_db as db: + job_ids = await db.get_job_ids_for_stamps(["s1"]) + assert sorted(job_ids) == [10, 20] + + +async def test_assign_jobs_to_pilot_duplicate_raises(pilot_db): + async with pilot_db as db: + await db.register_pilots(pilot_stamps=["s1"], vo="lhcb") + + pilot_id = await _get_pilot_id(pilot_db, "s1") + + async with pilot_db as db: + await db.assign_jobs_to_pilot( + [{"PilotID": pilot_id, "JobID": 10, "StartTime": "2024-01-01T00:00:00Z"}] + ) + + with pytest.raises(PilotAlreadyAssociatedWithJobError): + async with pilot_db as db: + await db.assign_jobs_to_pilot( + [ + { + "PilotID": pilot_id, + "JobID": 10, + "StartTime": "2024-01-01T00:00:00Z", + } + ] + ) + + +async def test_get_job_ids_for_stamps(pilot_db): + async with pilot_db as db: + await db.register_pilots(pilot_stamps=["s1", "s2"], vo="lhcb") + + pilot_id = await _get_pilot_id(pilot_db, "s1") + + async with pilot_db as db: + await db.assign_jobs_to_pilot( + [{"PilotID": pilot_id, "JobID": 100, "StartTime": "2024-01-01T00:00:00Z"}] + ) + + async with pilot_db as db: + assert await db.get_job_ids_for_stamps(["s1"]) == [100] + assert await db.get_job_ids_for_stamps(["s2"]) == [] + assert await db.get_job_ids_for_stamps([]) == [] + + +async def test_get_pilot_ids_for_job_ids(pilot_db): + async with pilot_db as db: + await db.register_pilots(pilot_stamps=["s1"], vo="lhcb") + + pilot_id = await _get_pilot_id(pilot_db, "s1") + + async with pilot_db as db: + await db.assign_jobs_to_pilot( + [{"PilotID": pilot_id, "JobID": 200, "StartTime": "2024-01-01T00:00:00Z"}] + ) + + async with pilot_db as db: + assert await db.get_pilot_ids_for_job_ids([200]) == [pilot_id] + assert await db.get_pilot_ids_for_job_ids([999]) == [] + assert await db.get_pilot_ids_for_job_ids([]) == [] diff --git a/diracx-db/tests/test_dummy_db.py b/diracx-db/tests/test_dummy_db.py index f94eda5b7..8e324a28e 100644 --- a/diracx-db/tests/test_dummy_db.py +++ b/diracx-db/tests/test_dummy_db.py @@ -149,6 +149,7 @@ async def test_failed_transaction(dummy_db): assert result # This will raise an exception and the transaction will be rolled back + result = await dummy_db.summary(["unexistingfieldraisinganerror"], []) assert result[0]["count"] == 10 diff --git a/diracx-logic/src/diracx/logic/jobs/query.py b/diracx-logic/src/diracx/logic/jobs/query.py index f03158771..f5d649964 100644 --- a/diracx-logic/src/diracx/logic/jobs/query.py +++ b/diracx-logic/src/diracx/logic/jobs/query.py @@ -4,33 +4,114 @@ from typing import Any from diracx.core.config import Config +from diracx.core.exceptions import InvalidQueryError from diracx.core.models import ( ScalarSearchOperator, SearchParams, SummaryParams, + VectorSearchOperator, + VectorSearchSpec, ) from diracx.db.os.job_parameters import JobParametersDB from diracx.db.sql.job.db import JobDB from diracx.db.sql.job_logging.db import JobLoggingDB +from diracx.db.sql.pilots.db import PilotAgentsDB +from diracx.logic.pilots.query import resolve_jobs_for_pilot_stamps logger = logging.getLogger(__name__) MAX_PER_PAGE = 10000 +# Pseudo-parameter accepted on POST /api/jobs/search. Resolves to a +# JobID IN (...) filter via JobToPilotMapping. +PILOT_STAMP_PSEUDO_PARAM = "PilotStamp" +# Real Jobs column that PilotStamp would collide with if both were +# accepted in the same request body. +JOB_ID_REAL_PARAM = "JobID" + + +async def _rewrite_pilot_stamp_pseudo_param( + pilot_db: PilotAgentsDB, body: SearchParams +) -> bool: + """Rewrite any `PilotStamp` pseudo-parameter in `body.search`. + + Collects every `PilotStamp` filter, resolves them through + `JobToPilotMapping`, removes the originals from `body.search`, and + appends a single `JobID IN (...)` vector filter. Returns `True` + if the resolution produced an empty list (the caller should + short-circuit to an empty result), `False` otherwise. + + Supports `eq` and `in` operators only; every other operator raises + `InvalidQueryError` because the join semantics are ambiguous. + Combining a `PilotStamp` pseudo-filter with a real `JobID` filter + in the same body is also refused. + """ + matches = [ + spec + for spec in body.search + if spec.get("parameter") == PILOT_STAMP_PSEUDO_PARAM + ] + if not matches: + return False + + if any(spec.get("parameter") == JOB_ID_REAL_PARAM for spec in body.search): + raise InvalidQueryError( + f"Cannot combine {PILOT_STAMP_PSEUDO_PARAM!r} pseudo-parameter " + f"with a real {JOB_ID_REAL_PARAM!r} filter in the same request." + ) + + stamps: list[str] = [] + for spec in matches: + operator = spec.get("operator") + if operator == ScalarSearchOperator.EQUAL: + stamps.append(str(spec["value"])) # type: ignore[typeddict-item] + elif operator == VectorSearchOperator.IN: + stamps.extend(str(v) for v in spec["values"]) # type: ignore[typeddict-item] + else: + raise InvalidQueryError( + f"Operator {operator!r} is not supported on the " + f"{PILOT_STAMP_PSEUDO_PARAM!r} pseudo-parameter; " + "use 'eq' or 'in'." + ) + + job_ids = await resolve_jobs_for_pilot_stamps(pilot_db, stamps) + body.search = [ + spec + for spec in body.search + if spec.get("parameter") != PILOT_STAMP_PSEUDO_PARAM + ] + if not job_ids: + return True + body.search.append( + VectorSearchSpec( + parameter=JOB_ID_REAL_PARAM, + operator=VectorSearchOperator.IN, + values=job_ids, + ) + ) + return False + async def search( config: Config, job_db: JobDB, job_parameters_db: JobParametersDB, job_logging_db: JobLoggingDB, + pilot_db: PilotAgentsDB, preferred_username: str | None, vo: str, page: int = 1, per_page: int = 100, body: SearchParams | None = None, ) -> tuple[int, list[dict[str, Any]]]: - """Retrieve information about jobs.""" + """Retrieve information about jobs. + + Accepts a `PilotStamp` pseudo-parameter in `body.search` + (`eq`/`in` only): it is resolved through `JobToPilotMapping` into + a concrete `JobID` vector filter before the main query runs. Mirrors + the `JobID` pseudo-parameter on `POST /api/pilots/search`. + """ # Apply a limit to per_page to prevent abuse of the API if per_page > MAX_PER_PAGE: per_page = MAX_PER_PAGE @@ -38,6 +119,10 @@ async def search( if body is None: body = SearchParams() + empty_after_rewrite = await _rewrite_pilot_stamp_pseudo_param(pilot_db, body) + if empty_after_rewrite: + return 0, [] + if query_logging_info := ("LoggingInfo" in (body.parameters or [])): if body.parameters: body.parameters.remove("LoggingInfo") diff --git a/diracx-logic/src/diracx/logic/pilots/__init__.py b/diracx-logic/src/diracx/logic/pilots/__init__.py new file mode 100644 index 000000000..2cda7e587 --- /dev/null +++ b/diracx-logic/src/diracx/logic/pilots/__init__.py @@ -0,0 +1,29 @@ +from __future__ import annotations + +__all__ = [ + "JOB_ID_PSEUDO_PARAM", + "MAX_PER_PAGE", + "PILOT_ID_REAL_PARAM", + "assign_jobs_to_pilot", + "get_pilots_by_stamp", + "register_new_pilot", + "resolve_jobs_for_pilot_stamps", + "search", + "summary", + "update_pilots_metadata", +] + +from .management import ( + assign_jobs_to_pilot, + register_new_pilot, + update_pilots_metadata, +) +from .query import ( + JOB_ID_PSEUDO_PARAM, + MAX_PER_PAGE, + PILOT_ID_REAL_PARAM, + get_pilots_by_stamp, + resolve_jobs_for_pilot_stamps, + search, + summary, +) diff --git a/diracx-logic/src/diracx/logic/pilots/management.py b/diracx-logic/src/diracx/logic/pilots/management.py new file mode 100644 index 000000000..6395fe00d --- /dev/null +++ b/diracx-logic/src/diracx/logic/pilots/management.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +from datetime import datetime, timezone +from typing import Any + +from diracx.core.config import Config +from diracx.core.exceptions import PilotAlreadyExistsError, PilotNotFoundError +from diracx.core.models.pilot import PilotMetadata, PilotRegistrationParams +from diracx.db.sql import PilotAgentsDB + +from .query import get_pilots_by_stamp + + +async def register_new_pilot( + config: Config, + pilot_db: PilotAgentsDB, + registration: PilotRegistrationParams, +): + """Register a new pilot. + + Raises `ValueError` if the VO is not in the registry, and + `PilotAlreadyExistsError` if the stamp already exists. + + Uniqueness is best-effort: the DIRAC `PilotAgents` schema has no unique + constraint on `PilotStamp`, so a concurrent registration of the same + stamp from two processes could race past this check. In practice pilot + stamps are cryptographically random UUIDs, making the collision window + negligible. + """ + # TODO: https://github.com/DIRACGrid/diracx/issues/1005 + # Also validate grid_type, grid_site and destination_site once the + # Resources section of the CS is modeled in the Config schema. + if registration.vo not in config.registry: + raise ValueError( + f"VO {registration.vo!r} is not registered in this installation." + ) + + stamp = registration.pilot_stamp + if await get_pilots_by_stamp(pilot_db=pilot_db, pilot_stamps=[stamp]): + raise PilotAlreadyExistsError(f"Pilot with stamp {stamp!r} already exists") + + await pilot_db.register_pilots( + pilot_stamps=[stamp], + vo=registration.vo, + grid_type=registration.grid_type, + grid_site=registration.grid_site, + destination_site=registration.destination_site, + pilot_references={stamp: registration.pilot_reference} + if registration.pilot_reference + else None, + status=registration.pilot_status, + ) + + +async def update_pilots_metadata( + pilot_db: PilotAgentsDB, + updates: dict[str, PilotMetadata], +): + """Bulk-update pilot metadata, keyed by pilot stamp. + + Unset fields (None) are preserved. `LastUpdateTime` is refreshed on + every updated pilot. + """ + fields_by_stamp = { + stamp: metadata.model_dump(by_alias=True, exclude_none=True) + for stamp, metadata in updates.items() + } + + if not any(fields_by_stamp.values()): + return + + now = datetime.now(tz=timezone.utc) + for fields in fields_by_stamp.values(): + fields["LastUpdateTime"] = now + + await pilot_db.update_pilot_metadata(fields_by_stamp) + + +async def assign_jobs_to_pilot( + pilot_db: PilotAgentsDB, pilot_stamp: str, job_ids: list[int] +): + """Associate jobs with a pilot identified by its stamp.""" + pilots = await get_pilots_by_stamp( + pilot_db=pilot_db, + pilot_stamps=[pilot_stamp], + parameters=["PilotID"], + ) + if not pilots: + raise PilotNotFoundError(detail=f"pilot {pilot_stamp!r} does not exist") + pilot_id = pilots[0]["PilotID"] + + job_to_pilot_mapping: list[dict[str, Any]] = [ + { + "PilotID": pilot_id, + "JobID": job_id, + "StartTime": datetime.now(tz=timezone.utc), + } + for job_id in job_ids + ] + + await pilot_db.assign_jobs_to_pilot(job_to_pilot_mapping=job_to_pilot_mapping) diff --git a/diracx-logic/src/diracx/logic/pilots/query.py b/diracx-logic/src/diracx/logic/pilots/query.py new file mode 100644 index 000000000..aa14f89ce --- /dev/null +++ b/diracx-logic/src/diracx/logic/pilots/query.py @@ -0,0 +1,206 @@ +from __future__ import annotations + +from typing import Any + +from diracx.core.exceptions import InvalidQueryError +from diracx.core.models.search import ( + ScalarSearchOperator, + ScalarSearchSpec, + SearchParams, + SummaryParams, + VectorSearchOperator, + VectorSearchSpec, +) +from diracx.db.sql import PilotAgentsDB + +MAX_PER_PAGE = 10000 + +# Pseudo-parameter accepted on POST /api/pilots/search. Resolves to a +# PilotID IN (...) filter via JobToPilotMapping. +JOB_ID_PSEUDO_PARAM = "JobID" +# Real column on PilotAgents that JobID would collide with if both +# were accepted in the same request body. +PILOT_ID_REAL_PARAM = "PilotID" + + +def _add_vo_constraint( + body: SearchParams | SummaryParams, vo_constraint: str | None +) -> None: + """Add a VO filter to the search body if a constraint is supplied. + + Admin callers pass `vo_constraint=None` to bypass the filter and query + across all VOs. Mirrors the intra-VO pattern of `logic/jobs/query.py`. + """ + if vo_constraint is None: + return + body.search.append( + ScalarSearchSpec( + parameter="VO", + operator=ScalarSearchOperator.EQUAL, + value=vo_constraint, + ) + ) + + +async def resolve_jobs_for_pilot_stamps( + pilot_db: PilotAgentsDB, pilot_stamps: list[str] +) -> list[int]: + """Resolve a batch of pilot stamps to the job IDs they have run. + + Used by `logic/jobs/query.py:search` to rewrite the `PilotStamp` + pseudo-parameter into a concrete `JobID` vector filter. + """ + return await pilot_db.get_job_ids_for_stamps(pilot_stamps) + + +async def _resolve_pilots_for_job_ids( + pilot_db: PilotAgentsDB, job_ids: list[int] +) -> list[int]: + """Resolve a batch of job IDs to the pilot IDs that have run them.""" + return await pilot_db.get_pilot_ids_for_job_ids(job_ids) + + +async def _rewrite_job_id_pseudo_param( + pilot_db: PilotAgentsDB, body: SearchParams +) -> bool: + """Rewrite any `JobID` pseudo-parameter in `body.search`. + + Collects every `JobID` filter, resolves them through + `JobToPilotMapping`, removes the originals from `body.search`, and + appends a single `PilotID IN (...)` vector filter. Returns `True` + if the resolution produced an empty list (in which case the caller + should short-circuit to an empty result), `False` otherwise. + + Supports `eq` and `in` operators only; every other operator raises + `InvalidQueryError` because the join semantics are ambiguous. + Combining a `JobID` pseudo-filter with a real `PilotID` filter in + the same body is also refused. + """ + matches = [ + spec for spec in body.search if spec.get("parameter") == JOB_ID_PSEUDO_PARAM + ] + if not matches: + return False + + if any(spec.get("parameter") == PILOT_ID_REAL_PARAM for spec in body.search): + raise InvalidQueryError( + f"Cannot combine {JOB_ID_PSEUDO_PARAM!r} pseudo-parameter with a " + f"real {PILOT_ID_REAL_PARAM!r} filter in the same request." + ) + + job_ids: list[int] = [] + for spec in matches: + operator = spec.get("operator") + try: + if operator == ScalarSearchOperator.EQUAL: + job_ids.append(int(spec["value"])) # type: ignore[typeddict-item] + elif operator == VectorSearchOperator.IN: + job_ids.extend(int(v) for v in spec["values"]) # type: ignore[typeddict-item] + else: + raise InvalidQueryError( + f"Operator {operator!r} is not supported on the " + f"{JOB_ID_PSEUDO_PARAM!r} pseudo-parameter; use 'eq' or 'in'." + ) + except (TypeError, ValueError) as e: + raise InvalidQueryError( + f"{JOB_ID_PSEUDO_PARAM!r} pseudo-parameter values must be integers." + ) from e + + pilot_ids = await _resolve_pilots_for_job_ids(pilot_db, job_ids) + body.search = [ + spec for spec in body.search if spec.get("parameter") != JOB_ID_PSEUDO_PARAM + ] + if not pilot_ids: + return True + body.search.append( + VectorSearchSpec( + parameter=PILOT_ID_REAL_PARAM, + operator=VectorSearchOperator.IN, + values=pilot_ids, + ) + ) + return False + + +async def search( + pilot_db: PilotAgentsDB, + vo_constraint: str | None, + page: int = 1, + per_page: int = 100, + body: SearchParams | None = None, +) -> tuple[int, list[dict[str, Any]]]: + """Retrieve information about pilots. + + `vo_constraint` restricts results to a single VO; pass `None` to + query across VOs (reserved for service administrators). + + Accepts a `JobID` pseudo-parameter in `body.search` (`eq`/`in` + only): it is resolved through `JobToPilotMapping` into a concrete + `PilotID` vector filter before the main query runs. Mirrors the + `PilotStamp` pseudo-parameter on `POST /api/jobs/search`. + """ + if per_page > MAX_PER_PAGE: + per_page = MAX_PER_PAGE + + if body is None: + body = SearchParams() + + empty_after_rewrite = await _rewrite_job_id_pseudo_param(pilot_db, body) + if empty_after_rewrite: + return 0, [] + + _add_vo_constraint(body, vo_constraint) + + return await pilot_db.search( + body.parameters, + body.search, + body.sort, + distinct=body.distinct, + page=page, + per_page=per_page, + ) + + +async def summary( + pilot_db: PilotAgentsDB, + body: SummaryParams, + vo_constraint: str | None, +): + """Aggregate pilot counts suitable for plotting.""" + _add_vo_constraint(body, vo_constraint) + return await pilot_db.summary(body.grouping, body.search) + + +async def get_pilots_by_stamp( + pilot_db: PilotAgentsDB, + pilot_stamps: list[str], + parameters: list[str] | None = None, +) -> list[dict[str, Any]]: + """Return the pilots whose stamp is in `pilot_stamps`. + + Missing stamps are silently omitted from the result. Callers that care + about completeness must compare the returned length to the input. + `PilotStamp` is always included in the returned parameters so callers + can identify which stamps were found. + """ + if parameters is None: + query_parameters: list[str] | None = None + else: + query_parameters = list(parameters) + if "PilotStamp" not in query_parameters: + query_parameters.append("PilotStamp") + + # No `page` is passed, so pagination is disabled and all matches are + # returned (`per_page` is ignored when `page` is None). + _, pilots = await pilot_db.search( + parameters=query_parameters, + search=[ + VectorSearchSpec( + parameter="PilotStamp", + operator=VectorSearchOperator.IN, + values=pilot_stamps, + ) + ], + sorts=[], + ) + return pilots diff --git a/diracx-routers/pyproject.toml b/diracx-routers/pyproject.toml index ac70f9cfe..cc57c49ae 100644 --- a/diracx-routers/pyproject.toml +++ b/diracx-routers/pyproject.toml @@ -6,11 +6,11 @@ requires-python = ">=3.11" keywords = [] license = { text = "GPL-3.0-only" } classifiers = [ - "Intended Audience :: Science/Research", - "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", - "Programming Language :: Python :: 3", - "Topic :: Scientific/Engineering", - "Topic :: System :: Distributed Computing", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", + "Programming Language :: Python :: 3", + "Topic :: Scientific/Engineering", + "Topic :: System :: Distributed Computing", ] dependencies = [ "cachetools", @@ -48,11 +48,13 @@ config = "diracx.routers.configuration:router" health = "diracx.routers.health:router" jobs = "diracx.routers.jobs:router" rss = "diracx.routers.rss:router" +pilots = "diracx.routers.pilots:router" [project.entry-points."diracx.access_policies"] wms = "diracx.routers.jobs.access_policies:WMSAccessPolicy" sandbox = "diracx.routers.jobs.access_policies:SandboxAccessPolicy" rss = "diracx.routers.rss.access_policies:RSSAccessPolicy" +pilot = "diracx.routers.pilots.access_policies:PilotManagementAccessPolicy" # Minimum version of the client supported [project.entry-points."diracx.min_client_version"] @@ -77,16 +79,16 @@ packages = ["src/diracx"] [tool.pytest.ini_options] testpaths = ["tests"] addopts = [ - "-v", - "--cov=diracx.routers", - "--cov-report=term-missing", - "-pdiracx.testing", - "-pdiracx.testing.osdb", - "--import-mode=importlib", + "-v", + "--cov=diracx.routers", + "--cov-report=term-missing", + "-pdiracx.testing", + "-pdiracx.testing.osdb", + "--import-mode=importlib", ] asyncio_mode = "auto" markers = [ - "enabled_dependencies: List of dependencies which should be available to the FastAPI test client", + "enabled_dependencies: List of dependencies which should be available to the FastAPI test client", ] asyncio_default_fixture_loop_scope = "function" diff --git a/diracx-routers/src/diracx/routers/jobs/query.py b/diracx-routers/src/diracx/routers/jobs/query.py index 2fdf0b1cc..37470f86f 100644 --- a/diracx-routers/src/diracx/routers/jobs/query.py +++ b/diracx-routers/src/diracx/routers/jobs/query.py @@ -11,7 +11,7 @@ ) from diracx.core.properties import JOB_ADMINISTRATOR from diracx.db.os import JobParametersDB -from diracx.db.sql import JobDB, JobLoggingDB +from diracx.db.sql import JobDB, JobLoggingDB, PilotAgentsDB from diracx.logic.jobs import MAX_PER_PAGE from diracx.logic.jobs import search as search_bl from diracx.logic.jobs import summary as summary_bl @@ -45,6 +45,20 @@ "sort": [{"parameter": "JobID", "direction": "asc"}], }, }, + "Jobs run on a given pilot": { + "summary": "Jobs run on a given pilot", + "description": ( + "Find all jobs that have run on a specific pilot. `PilotStamp` " + "is a pseudo-parameter resolved through `JobToPilotMapping` " + "into a `JobID` filter; only `eq` and `in` operators are " + "supported." + ), + "value": { + "search": [ + {"parameter": "PilotStamp", "operator": "eq", "value": "abc-123"} + ] + }, + }, } @@ -123,6 +137,7 @@ async def search( job_db: JobDB, job_parameters_db: JobParametersDB, job_logging_db: JobLoggingDB, + pilot_db: PilotAgentsDB, user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)], check_permissions: CheckWMSPolicyCallable, response: Response, @@ -143,6 +158,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + + A `PilotStamp` pseudo-parameter is also accepted in the `search` + filter list (operators `eq` / `in` only): it is transparently + resolved through `JobToPilotMapping` into a `JobID` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. """ await check_permissions(action=ActionType.QUERY, job_db=job_db) @@ -155,6 +176,7 @@ async def search( job_db=job_db, job_parameters_db=job_parameters_db, job_logging_db=job_logging_db, + pilot_db=pilot_db, preferred_username=preferred_username, vo=user_info.vo, page=page, diff --git a/diracx-routers/src/diracx/routers/jobs/submission.py b/diracx-routers/src/diracx/routers/jobs/submission.py index 62f215394..8bca33e73 100644 --- a/diracx-routers/src/diracx/routers/jobs/submission.py +++ b/diracx-routers/src/diracx/routers/jobs/submission.py @@ -51,7 +51,7 @@ class JobID(BaseModel): } -@router.post("/jdl") +@router.post("/jdl", status_code=HTTPStatus.CREATED) async def submit_jdl_jobs( job_definitions: Annotated[list[str], Body(openapi_examples=EXAMPLE_JDLS)], job_db: JobDB, diff --git a/diracx-routers/src/diracx/routers/pilots/__init__.py b/diracx-routers/src/diracx/routers/pilots/__init__.py new file mode 100644 index 000000000..08b53db29 --- /dev/null +++ b/diracx-routers/src/diracx/routers/pilots/__init__.py @@ -0,0 +1,20 @@ +from __future__ import annotations + +__all__ = [ + "ActionType", + "PilotManagementAccessPolicy", + "router", +] + +import logging + +from ..fastapi_classes import DiracxRouter +from .access_policies import ActionType, PilotManagementAccessPolicy +from .management import router as management_router +from .query import router as query_router + +logger = logging.getLogger(__name__) + +router = DiracxRouter() +router.include_router(management_router) +router.include_router(query_router) diff --git a/diracx-routers/src/diracx/routers/pilots/access_policies.py b/diracx-routers/src/diracx/routers/pilots/access_policies.py new file mode 100644 index 000000000..66e07fb40 --- /dev/null +++ b/diracx-routers/src/diracx/routers/pilots/access_policies.py @@ -0,0 +1,117 @@ +from __future__ import annotations + +from collections.abc import Callable +from enum import StrEnum, auto +from http import HTTPStatus +from typing import Annotated + +from fastapi import Depends, HTTPException + +from diracx.core.properties import GENERIC_PILOT, SERVICE_ADMINISTRATOR +from diracx.db.sql.pilots.db import PilotAgentsDB +from diracx.logic.pilots.query import get_pilots_by_stamp +from diracx.routers.access_policies import BaseAccessPolicy +from diracx.routers.utils.users import AuthorizedUserInfo + + +class ActionType(StrEnum): + # Change pilot metadata (status, fields, etc.). Admin-only by default; + # legacy pilot X.509 identities can be allowed via `allow_legacy_pilots`. + MANAGE_PILOTS = auto() + # Read pilot metadata. Normal users can read their own VO's pilots; + # `SERVICE_ADMINISTRATOR` can read across VOs. + READ_PILOT_METADATA = auto() + + +class PilotManagementAccessPolicy(BaseAccessPolicy): + """Pilot management access policy. + + * Every user can read pilots from their own VO. + * Service administrators can read across VOs and manage pilots. + * Legacy X.509 pilot identities may be allowed to manage themselves when + `allow_legacy_pilots=True` is passed by the route. + """ + + @staticmethod + async def policy( + policy_name: str, + user_info: AuthorizedUserInfo, + /, + *, + action: ActionType | None = None, + pilot_db: PilotAgentsDB | None = None, + pilot_stamps: list[str] | None = None, + target_vo: str | None = None, + allow_legacy_pilots: bool = False, + ): + # Authorization is VO-scoped, not bound to the caller's + # own pilot stamp. This mirrors DIRAC's PilotManagerHandler, which has + # no ownership check either. + if action is None: + raise ValueError("action is a mandatory parameter") + + # Service administrators read and manage pilots across VOs; every + # other identity is confined to its own VO. + is_admin = SERVICE_ADMINISTRATOR in user_info.properties + + if action == ActionType.MANAGE_PILOTS: + is_legacy_pilot = ( + allow_legacy_pilots and GENERIC_PILOT in user_info.properties + ) + if not is_admin and not is_legacy_pilot: + raise HTTPException( + status_code=HTTPStatus.FORBIDDEN, + detail="Insufficient permissions to manage pilots.", + ) + # Limit the damage a stolen legacy pilot credential can do: a + # pilot identity may only act on a single pilot per call. Note + # this bounds the rate of abuse, not its scope. + if not is_admin and pilot_stamps and len(set(pilot_stamps)) > 1: + raise HTTPException( + status_code=HTTPStatus.FORBIDDEN, + detail="Insufficient permissions to modify more than one pilot.", + ) + # `target_vo` is the VO that pilots are being registered into + # (pilots that already exist are checked against the DB through + # `pilot_stamps` below). Only service administrators may + # register pilots outside their own VO. + if target_vo is not None and not is_admin and target_vo != user_info.vo: + raise HTTPException( + status_code=HTTPStatus.FORBIDDEN, + detail="Pilots can only be registered for your own VO.", + ) + + if action == ActionType.READ_PILOT_METADATA: + if GENERIC_PILOT in user_info.properties: + raise HTTPException( + status_code=HTTPStatus.FORBIDDEN, + detail="Pilots cannot read other pilots' metadata.", + ) + + # If pilot stamps are provided, verify they all exist and (unless the + # caller is a service administrator) belong to the user's VO. + if pilot_db is not None and pilot_stamps: + pilots = await get_pilots_by_stamp( + pilot_db=pilot_db, + pilot_stamps=pilot_stamps, + parameters=["VO"], + ) + if len(pilots) != len(set(pilot_stamps)): + raise HTTPException( + status_code=HTTPStatus.NOT_FOUND, + detail="At least one pilot does not exist.", + ) + if not is_admin and not all( + pilot["VO"] == user_info.vo for pilot in pilots + ): + raise HTTPException( + status_code=HTTPStatus.FORBIDDEN, + detail=( + "Insufficient permissions to access all of the provided pilots." + ), + ) + + +CheckPilotManagementPolicyCallable = Annotated[ + Callable, Depends(PilotManagementAccessPolicy.check) +] diff --git a/diracx-routers/src/diracx/routers/pilots/management.py b/diracx-routers/src/diracx/routers/pilots/management.py new file mode 100644 index 000000000..3d92a32ad --- /dev/null +++ b/diracx-routers/src/diracx/routers/pilots/management.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +from http import HTTPStatus +from typing import Annotated + +from fastapi import Body, HTTPException + +from diracx.core.exceptions import PilotAlreadyExistsError, PilotNotFoundError +from diracx.core.models.pilot import PilotMetadata, PilotRegistrationParams +from diracx.db.sql import PilotAgentsDB +from diracx.logic.pilots import ( + register_new_pilot, + update_pilots_metadata, +) + +from ..dependencies import Config +from ..fastapi_classes import DiracxRouter +from .access_policies import ( + ActionType, + CheckPilotManagementPolicyCallable, +) + +router = DiracxRouter() + + +@router.post("/", status_code=HTTPStatus.CREATED) +async def register_pilot( + config: Config, + pilot_db: PilotAgentsDB, + registration: PilotRegistrationParams, + check_permissions: CheckPilotManagementPolicyCallable, +): + """Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + """ + # Legacy (X.509 / GENERIC_PILOT) pilot identities may self-register: + # pilots started in the vacuum have no SiteDirector to register them. + # This mirrors dirac-admin-add-pilot in legacy DIRAC. The route takes a + # single stamp per call, which bounds what a stolen credential can do. + await check_permissions( + action=ActionType.MANAGE_PILOTS, + target_vo=registration.vo, + allow_legacy_pilots=True, + ) + + try: + await register_new_pilot( + config=config, + pilot_db=pilot_db, + registration=registration, + ) + except ValueError as e: + raise HTTPException(status_code=HTTPStatus.BAD_REQUEST, detail=str(e)) from e + except PilotAlreadyExistsError as e: + raise HTTPException(status_code=HTTPStatus.CONFLICT, detail=str(e)) from e + + +EXAMPLE_UPDATE_METADATA = { + "Update the BenchMark field": { + "summary": "Update BenchMark", + "description": "Update only the BenchMark for one pilot.", + "value": {"the_pilot_stamp": {"BenchMark": 1.0}}, + }, + "Update multiple statuses": { + "summary": "Update multiple pilots", + "description": "Update statuses for multiple pilots at once.", + "value": { + "first_stamp": {"Status": "Waiting"}, + "second_stamp": {"Status": "Waiting"}, + }, + }, +} + + +@router.patch("/metadata", status_code=HTTPStatus.NO_CONTENT) +async def update_pilot_metadata( + updates: Annotated[ + dict[str, PilotMetadata], + Body( + description="Mapping from pilot stamp to the metadata to apply.", + openapi_examples=EXAMPLE_UPDATE_METADATA, # type: ignore + ), + ], + pilot_db: PilotAgentsDB, + check_permissions: CheckPilotManagementPolicyCallable, +): + """Update pilot metadata (status, benchmark, etc.). + + Only fields defined in `PilotMetadata` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + """ + # Legacy pilot identities may self-update (dirac-admin-add-pilot + # --status); the policy caps them to a single pilot stamp per call. + await check_permissions( + action=ActionType.MANAGE_PILOTS, + pilot_db=pilot_db, + pilot_stamps=list(updates), + allow_legacy_pilots=True, + ) + + try: + await update_pilots_metadata( + pilot_db=pilot_db, + updates=updates, + ) + except PilotNotFoundError as e: + raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail=str(e)) from e diff --git a/diracx-routers/src/diracx/routers/pilots/query.py b/diracx-routers/src/diracx/routers/pilots/query.py new file mode 100644 index 000000000..538f637a9 --- /dev/null +++ b/diracx-routers/src/diracx/routers/pilots/query.py @@ -0,0 +1,238 @@ +from __future__ import annotations + +from http import HTTPStatus +from typing import Annotated, Any + +from fastapi import Body, Depends, Query, Response + +from diracx.core.models.search import SearchParams, SummaryParams +from diracx.core.properties import SERVICE_ADMINISTRATOR +from diracx.db.sql import PilotAgentsDB +from diracx.logic.pilots import MAX_PER_PAGE +from diracx.logic.pilots import search as search_bl +from diracx.logic.pilots import summary as summary_bl + +from ..fastapi_classes import DiracxRouter +from ..utils.users import AuthorizedUserInfo, verify_dirac_access_token +from .access_policies import ( + ActionType, + CheckPilotManagementPolicyCallable, +) + +router = DiracxRouter() + + +def _vo_constraint_for(user_info: AuthorizedUserInfo) -> str | None: + """Return the VO filter to apply for this user, or None for admins.""" + if SERVICE_ADMINISTRATOR in user_info.properties: + return None + return user_info.vo + + +EXAMPLE_SEARCHES = { + "Show all": { + "summary": "Show all", + "description": "Shows all pilots the current user has access to.", + "value": {}, + }, + "A specific pilot": { + "summary": "A specific pilot", + "description": "Search for a specific pilot by ID", + "value": {"search": [{"parameter": "PilotID", "operator": "eq", "value": "5"}]}, + }, + "Get ordered pilot statuses": { + "summary": "Get ordered pilot statuses", + "description": "Get only pilot statuses for specific pilots, ordered by status", + "value": { + "parameters": ["PilotID", "Status"], + "search": [ + {"parameter": "PilotID", "operator": "in", "values": ["6", "2", "3"]} + ], + "sort": [{"parameter": "PilotID", "direction": "asc"}], + }, + }, + "Pilots that ran a given job": { + "summary": "Pilots that ran a given job", + "description": ( + "Find all pilots that have run a specific job. `JobID` is a " + "pseudo-parameter resolved through `JobToPilotMapping` into a " + "`PilotID` filter; only `eq` and `in` operators are supported." + ), + "value": {"search": [{"parameter": "JobID", "operator": "eq", "value": 42}]}, + }, +} + + +EXAMPLE_RESPONSES: dict[int | str, dict[str, Any]] = { + 200: { + "description": "List of matching results", + "content": { + "application/json": { + "example": [ + { + "PilotID": 3, + "SubmissionTime": "2023-05-25T07:03:35.602654", + "LastUpdateTime": "2023-05-25T07:03:35.602656", + "Status": "Running", + "GridType": "Dirac", + "BenchMark": 1.0, + }, + ] + } + }, + }, + 206: { + "description": "Partial Content. Only a part of the requested range could be served.", + "headers": { + "Content-Range": { + "description": "The range of pilots returned in this response", + "schema": {"type": "string", "example": "pilots 0-1/4"}, + } + }, + "model": list[dict[str, Any]], + }, +} + + +@router.post("/search", responses=EXAMPLE_RESPONSES) +async def search( + pilot_db: PilotAgentsDB, + check_permissions: CheckPilotManagementPolicyCallable, + response: Response, + user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)], + page: Annotated[int, Query(ge=1)] = 1, + per_page: Annotated[int, Query(ge=1, le=MAX_PER_PAGE)] = 100, + body: Annotated[ + SearchParams | None, Body(openapi_examples=EXAMPLE_SEARCHES) # type: ignore + ] = None, +) -> list[dict[str, Any]]: + """Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A `JobID` pseudo-parameter is also accepted in the `search` filter + list (operators `eq` / `in` only): it is transparently resolved + through `JobToPilotMapping` into a `PilotID` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + """ + await check_permissions(action=ActionType.READ_PILOT_METADATA) + + total, pilots = await search_bl( + pilot_db=pilot_db, + vo_constraint=_vo_constraint_for(user_info), + page=page, + per_page=per_page, + body=body, + ) + + # RFC 7233 Content-Range handling, matching /api/jobs/search + if len(pilots) == 0 and total > 0: + response.headers["Content-Range"] = f"pilots */{total}" + response.status_code = HTTPStatus.REQUESTED_RANGE_NOT_SATISFIABLE + elif len(pilots) < total: + first_idx = per_page * (page - 1) + last_idx = min(first_idx + len(pilots), total) - 1 if total > 0 else 0 + response.headers["Content-Range"] = f"pilots {first_idx}-{last_idx}/{total}" + response.status_code = HTTPStatus.PARTIAL_CONTENT + return pilots + + +EXAMPLE_SUMMARY = { + "Show all": { + "summary": "Show all", + "description": "Shows all pilots the current user has access to.", + "value": {"grouping": [], "search": []}, + }, + "Group by Status": { + "summary": "Group the pilots by Status.", + "description": "We get all the unique pilot statuses and the number of pilots in each one.", + "value": {"grouping": ["Status"]}, + }, + "Group by GridSite with Search": { + "summary": "Group the pilots by GridSite and filter by status.", + "description": "We get all the unique GridSites where the status is 'Running'. We also get the number " + "of pilots in each site.", + "value": { + "grouping": ["GridSite"], + "search": [{"parameter": "Status", "operator": "eq", "value": "Running"}], + }, + }, +} + +EXAMPLE_SUMMARY_RESPONSES = { + 200: { + "description": "Results of the request", + "content": { + "application/json": { + "examples": { + "without_grouping": { + "summary": "Results without grouping parameters", + "description": "Shows all pilots when no grouping is specified", + "value": [ + { + "PilotID": 1, + "Status": "Running", + "GridType": "Dirac", + "GridSite": "Site_1", + "DestinationSite": "Site_1", + "VO": "diracAdmin", + "SubmissionTime": "2025-07-15T07:15:57", + "LastUpdateTime": "2025-07-15T07:15:57", + "BenchMark": 1.0, + "count": 1, + }, + { + "PilotID": 2, + "Status": "Done", + "GridType": "Dirac", + "GridSite": "Site_2", + "DestinationSite": "Site_2", + "VO": "diracAdmin", + "SubmissionTime": "2025-07-15T07:15:57", + "LastUpdateTime": "2025-07-15T07:15:57", + "BenchMark": 1.0, + "count": 1, + }, + ], + }, + "with_grouping": { + "summary": "Results with grouping parameters", + "description": "Shows grouped results when grouping parameters are specified", + "value": [ + {"Status": "Running", "count": 4000}, + {"Status": "Done", "count": 2000}, + {"Status": "Failed", "count": 1000}, + {"Status": "Aborted", "count": 1000}, + {"Status": "Submitted", "count": 20}, + {"Status": "Waiting", "count": 20}, + ], + }, + } + } + }, + } +} + + +@router.post("/summary", responses=EXAMPLE_SUMMARY_RESPONSES) +async def summary( + pilot_db: PilotAgentsDB, + user_info: Annotated[AuthorizedUserInfo, Depends(verify_dirac_access_token)], + check_permissions: CheckPilotManagementPolicyCallable, + body: Annotated[ + SummaryParams, Body(openapi_examples=EXAMPLE_SUMMARY) # type: ignore + ], +): + """Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + """ + await check_permissions(action=ActionType.READ_PILOT_METADATA) + + return await summary_bl( + pilot_db=pilot_db, + body=body, + vo_constraint=_vo_constraint_for(user_info), + ) diff --git a/diracx-routers/tests/jobs/conftest.py b/diracx-routers/tests/jobs/conftest.py index bb430756b..8d733b235 100644 --- a/diracx-routers/tests/jobs/conftest.py +++ b/diracx-routers/tests/jobs/conftest.py @@ -75,7 +75,7 @@ def admin_user_client(client_factory): def valid_job_id(normal_user_client: TestClient): job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == 1 return r.json()[0]["JobID"] @@ -84,7 +84,7 @@ def valid_job_id(normal_user_client: TestClient): def valid_job_ids(normal_user_client: TestClient): job_definitions = [TEST_PARAMETRIC_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == 3 return sorted([job_dict["JobID"] for job_dict in r.json()]) diff --git a/diracx-routers/tests/jobs/test_heartbeat_commands.py b/diracx-routers/tests/jobs/test_heartbeat_commands.py index 6a53dabdd..0804e15c2 100644 --- a/diracx-routers/tests/jobs/test_heartbeat_commands.py +++ b/diracx-routers/tests/jobs/test_heartbeat_commands.py @@ -18,6 +18,8 @@ "WMSAccessPolicy", "DevelopmentSettings", "JobParametersDB", + "PilotAgentsDB", + "PilotManagementAccessPolicy", ] ) diff --git a/diracx-routers/tests/jobs/test_query.py b/diracx-routers/tests/jobs/test_query.py index cb1097f2d..56b699ac6 100644 --- a/diracx-routers/tests/jobs/test_query.py +++ b/diracx-routers/tests/jobs/test_query.py @@ -37,6 +37,8 @@ "WMSAccessPolicy", "DevelopmentSettings", "JobParametersDB", + "PilotAgentsDB", + "PilotManagementAccessPolicy", ] ) @@ -44,7 +46,7 @@ def test_insert_and_list_parametric_jobs(normal_user_client): job_definitions = [TEST_PARAMETRIC_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == 3 # Parameters.JOB_ID is 3 submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()]) @@ -71,7 +73,7 @@ def test_insert_and_list_parametric_jobs(normal_user_client): ) def test_insert_and_list_bulk_jobs(job_definitions, normal_user_client): r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()]) @@ -94,7 +96,7 @@ def test_insert_and_search(normal_user_client): job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) listed_jobs = r.json() - assert r.status_code == 200, listed_jobs + assert r.status_code == 201, listed_jobs assert len(listed_jobs) == len(job_definitions) submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()]) @@ -214,7 +216,7 @@ def test_insert_and_search_by_datetime(normal_user_client): job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) listed_jobs = r.json() - assert r.status_code == 200, listed_jobs + assert r.status_code == 201, listed_jobs assert len(listed_jobs) == len(job_definitions) r = normal_user_client.post("/api/jobs/search") assert len(r.json()) == 1, "No jobs submitted" @@ -451,7 +453,7 @@ def test_search_distinct(normal_user_client): job_definitions = [TEST_JDL, TEST_JDL, TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) listed_jobs = r.json() - assert r.status_code == 200, listed_jobs + assert r.status_code == 201, listed_jobs assert len(listed_jobs) == len(job_definitions) # Check that distinct collapses identical records when true @@ -479,7 +481,7 @@ def test_search_pagination(normal_user_client): job_definitions = [TEST_JDL] * 20 r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) listed_jobs = r.json() - assert r.status_code == 200, listed_jobs + assert r.status_code == 201, listed_jobs assert len(listed_jobs) == len(job_definitions) # Get the first 20 jobs (all of them) @@ -917,3 +919,125 @@ def test_summary_doc_example(normal_user_client: TestClient, valid_job_id: int): assert r.status_code == 200, r.json() assert len(r.json()) == 1 + + +# --------------------------------------------------------------------------- +# Cross-table search: PilotStamp pseudo-parameter on POST /api/jobs/search +# --------------------------------------------------------------------------- + + +async def _assign_pilot_to_jobs(client, stamp: str, job_ids: list[int]) -> None: + """Insert JobToPilotMapping rows directly. + + The router does not expose a public endpoint for pilot-job association + (deliberately — it waits for the DiracX pilot token story). Tests reach + into the app's dependency override to insert the rows via the DB layer. + """ + from diracx.db.sql import PilotAgentsDB + from diracx.logic.pilots.management import assign_jobs_to_pilot + + db = client.app.dependency_overrides[PilotAgentsDB.transaction].args[0] + async with db: + await assign_jobs_to_pilot(pilot_db=db, pilot_stamp=stamp, job_ids=job_ids) + + +async def _register_pilot(client, stamp: str) -> None: + r = client.post( + "/api/pilots/", + json={"pilot_stamp": stamp, "vo": "lhcb"}, + ) + assert r.status_code == 201, r.json() + + +async def test_jobs_search_by_pilot_stamp_eq(normal_user_client): + """A ``PilotStamp`` eq filter on /jobs/search returns the jobs that ran on that pilot.""" + r = normal_user_client.post("/api/jobs/jdl", json=[TEST_JDL for _ in range(3)]) + assert r.status_code == 201, r.json() + job_ids = [j["JobID"] for j in r.json()] + + await _register_pilot(normal_user_client, "stamp-eq") + await _assign_pilot_to_jobs( + normal_user_client, "stamp-eq", [job_ids[0], job_ids[1]] + ) + + r = normal_user_client.post( + "/api/jobs/search", + json={ + "search": [ + {"parameter": "PilotStamp", "operator": "eq", "value": "stamp-eq"} + ] + }, + ) + assert r.status_code == 200, r.json() + returned = sorted(j["JobID"] for j in r.json()) + assert returned == sorted([job_ids[0], job_ids[1]]) + + +async def test_jobs_search_by_pilot_stamp_in_multiple(normal_user_client): + """An ``in`` filter over several stamps returns the union of their jobs.""" + r = normal_user_client.post("/api/jobs/jdl", json=[TEST_JDL for _ in range(4)]) + assert r.status_code == 201, r.json() + job_ids = [j["JobID"] for j in r.json()] + + await _register_pilot(normal_user_client, "stamp-in-a") + await _register_pilot(normal_user_client, "stamp-in-b") + await _assign_pilot_to_jobs(normal_user_client, "stamp-in-a", [job_ids[0]]) + await _assign_pilot_to_jobs( + normal_user_client, "stamp-in-b", [job_ids[1], job_ids[2]] + ) + + r = normal_user_client.post( + "/api/jobs/search", + json={ + "search": [ + { + "parameter": "PilotStamp", + "operator": "in", + "values": ["stamp-in-a", "stamp-in-b"], + } + ] + }, + ) + assert r.status_code == 200, r.json() + returned = sorted(j["JobID"] for j in r.json()) + assert returned == sorted([job_ids[0], job_ids[1], job_ids[2]]) + + +def test_jobs_search_by_unknown_pilot_stamp_returns_empty(normal_user_client): + """An unknown stamp resolves to an empty job list; the caller gets ``[]``.""" + r = normal_user_client.post("/api/jobs/jdl", json=[TEST_JDL]) + assert r.status_code == 201, r.json() + + r = normal_user_client.post( + "/api/jobs/search", + json={ + "search": [{"parameter": "PilotStamp", "operator": "eq", "value": "nope"}] + }, + ) + assert r.status_code == 200 + assert r.json() == [] + + +def test_jobs_search_combining_pilot_stamp_and_job_id_raises(normal_user_client): + """Combining a ``PilotStamp`` pseudo-filter with a real ``JobID`` filter is refused.""" + r = normal_user_client.post( + "/api/jobs/search", + json={ + "search": [ + {"parameter": "PilotStamp", "operator": "eq", "value": "any"}, + {"parameter": "JobID", "operator": "eq", "value": 1}, + ] + }, + ) + assert r.status_code in (400, 422), r.json() + + +def test_jobs_search_pilot_stamp_unsupported_operator_raises(normal_user_client): + """Operators other than ``eq`` / ``in`` on ``PilotStamp`` are refused.""" + r = normal_user_client.post( + "/api/jobs/search", + json={ + "search": [{"parameter": "PilotStamp", "operator": "neq", "value": "any"}] + }, + ) + assert r.status_code in (400, 422), r.json() diff --git a/diracx-routers/tests/jobs/test_sandboxes.py b/diracx-routers/tests/jobs/test_sandboxes.py index 452556903..fa4286ff6 100644 --- a/diracx-routers/tests/jobs/test_sandboxes.py +++ b/diracx-routers/tests/jobs/test_sandboxes.py @@ -126,7 +126,7 @@ def test_assign_then_unassign_sandboxes_to_jobs(normal_user_client: TestClient): # Submit a job: job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) job_id = r.json()[0]["JobID"] @@ -167,7 +167,7 @@ def test_assign_then_unassign_sandboxes_to_jobs(normal_user_client: TestClient): r = normal_user_client.post("/api/jobs/sandbox/unassign", json={"job_ids": job_ids}) assert r.status_code == 204 # No content expected - # Get the sb again, it should'nt be there anymore: + # Get the sb again, it shouldn't be there anymore: short_pfn = sandbox_pfn.split("|", 1)[-1] r = normal_user_client.get(f"/api/jobs/{job_id}/sandbox") assert r.status_code == 200 @@ -221,7 +221,7 @@ def test_malformed_request_to_get_job_sandbox(normal_user_client: TestClient): # Submit a job: job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) job_id = r.json()[0]["JobID"] @@ -236,7 +236,7 @@ def test_get_empty_job_sandboxes(normal_user_client: TestClient): # Submit a job: job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) job_id = r.json()[0]["JobID"] @@ -251,7 +251,7 @@ def test_assign_nonexisting_sb_to_job(normal_user_client: TestClient): # Submit a job: job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) job_id = r.json()[0]["JobID"] @@ -288,7 +288,7 @@ def test_assign_sb_to_job_twice(normal_user_client: TestClient): # Submit a job: job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) job_id = r.json()[0]["JobID"] diff --git a/diracx-routers/tests/jobs/test_status.py b/diracx-routers/tests/jobs/test_status.py index e3ecaf3a7..d84a8868a 100644 --- a/diracx-routers/tests/jobs/test_status.py +++ b/diracx-routers/tests/jobs/test_status.py @@ -26,6 +26,8 @@ "WMSAccessPolicy", "DevelopmentSettings", "JobParametersDB", + "PilotAgentsDB", + "PilotManagementAccessPolicy", ] ) @@ -356,7 +358,7 @@ def test_set_job_status_with_invalid_job_id( def test_insert_and_reschedule(normal_user_client: TestClient): job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()]) @@ -406,7 +408,7 @@ def test_reschedule_job_attr_update(normal_user_client: TestClient): job_definitions = [TEST_JDL] * 15 r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == len(job_definitions) submitted_job_ids = sorted([job_dict["JobID"] for job_dict in r.json()]) diff --git a/diracx-routers/tests/pilots/__init__.py b/diracx-routers/tests/pilots/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/diracx-routers/tests/pilots/test_access_policy.py b/diracx-routers/tests/pilots/test_access_policy.py new file mode 100644 index 000000000..eb9692623 --- /dev/null +++ b/diracx-routers/tests/pilots/test_access_policy.py @@ -0,0 +1,203 @@ +"""Unit tests for `PilotManagementAccessPolicy`. + +These tests bypass the FastAPI test harness (which stubs the real policy +with `AlwaysAllowAccessPolicy`) and invoke the policy coroutine +directly, mirroring how it is called from a real request. +""" + +from __future__ import annotations + +from uuid import uuid4 + +import pytest +from fastapi import HTTPException + +from diracx.core.properties import GENERIC_PILOT, NORMAL_USER, SERVICE_ADMINISTRATOR +from diracx.db.sql.pilots.db import PilotAgentsDB +from diracx.routers.pilots.access_policies import ( + ActionType, + PilotManagementAccessPolicy, +) +from diracx.routers.utils.users import AuthorizedUserInfo + +MAIN_VO = "lhcb" + + +def _user(*properties, vo: str = MAIN_VO) -> AuthorizedUserInfo: + """Build a minimal AuthorizedUserInfo for policy tests.""" + return AuthorizedUserInfo( + bearer_token="", + token_id=str(uuid4()), + properties=list(properties), + sub="testingVO:sub", + preferred_username="test-user", + dirac_group="test_group", + vo=vo, + policies={}, + ) + + +@pytest.fixture +async def pilot_db_with_pilots(): + """Yield a pilot DB seeded with two pilots, both in MAIN_VO.""" + db = PilotAgentsDB("sqlite+aiosqlite:///:memory:") + async with db.engine_context(): + async with db.engine.begin() as conn: + await conn.run_sync(db.metadata.create_all) + async with db as pdb: + await pdb.register_pilots(pilot_stamps=["stamp-a", "stamp-b"], vo=MAIN_VO) + yield db + + +async def test_manage_requires_service_administrator(): + """A normal user cannot manage pilots.""" + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(NORMAL_USER), + action=ActionType.MANAGE_PILOTS, + ) + assert exc_info.value.status_code == 403 + + +async def test_manage_allows_service_administrator(): + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(SERVICE_ADMINISTRATOR), + action=ActionType.MANAGE_PILOTS, + ) + + +async def test_manage_allows_legacy_pilot_when_opted_in(): + """`allow_legacy_pilots=True` lets GENERIC_PILOT identities manage.""" + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.MANAGE_PILOTS, + allow_legacy_pilots=True, + ) + + +async def test_manage_rejects_legacy_pilot_when_not_opted_in(): + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.MANAGE_PILOTS, + ) + assert exc_info.value.status_code == 403 + + +async def test_manage_legacy_pilot_capped_to_single_stamp(): + """A legacy pilot identity may only act on one pilot stamp per call.""" + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.MANAGE_PILOTS, + pilot_stamps=["stamp-a", "stamp-b"], + allow_legacy_pilots=True, + ) + assert exc_info.value.status_code == 403 + + +async def test_manage_admin_not_capped_to_single_stamp(pilot_db_with_pilots): + """Service administrators may act on several pilots in one call.""" + async with pilot_db_with_pilots as db: + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(SERVICE_ADMINISTRATOR), + action=ActionType.MANAGE_PILOTS, + pilot_db=db, + pilot_stamps=["stamp-a", "stamp-b"], + ) + + +async def test_manage_target_vo_own_vo_allowed(): + """A legacy pilot may register a pilot into its own VO.""" + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.MANAGE_PILOTS, + target_vo=MAIN_VO, + allow_legacy_pilots=True, + ) + + +async def test_manage_target_vo_cross_vo_rejected_for_legacy_pilot(): + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.MANAGE_PILOTS, + target_vo="other-vo", + allow_legacy_pilots=True, + ) + assert exc_info.value.status_code == 403 + + +async def test_manage_target_vo_cross_vo_allowed_for_admin(): + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(SERVICE_ADMINISTRATOR), + action=ActionType.MANAGE_PILOTS, + target_vo="other-vo", + ) + + +async def test_read_denies_generic_pilots(): + """A pilot identity is not allowed to read other pilots' metadata.""" + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT), + action=ActionType.READ_PILOT_METADATA, + ) + assert exc_info.value.status_code == 403 + + +async def test_pilot_stamp_check_raises_404_on_unknown(pilot_db_with_pilots): + """Supplying an unknown pilot stamp must surface as 404.""" + async with pilot_db_with_pilots as db: + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(SERVICE_ADMINISTRATOR), + action=ActionType.MANAGE_PILOTS, + pilot_db=db, + pilot_stamps=["stamp-a", "nope"], + ) + assert exc_info.value.status_code == 404 + + +async def test_pilot_stamp_check_raises_403_on_cross_vo(pilot_db_with_pilots): + """A non-admin from another VO must not be able to act on this VO's pilots.""" + async with pilot_db_with_pilots as db: + with pytest.raises(HTTPException) as exc_info: + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(GENERIC_PILOT, vo="other-vo"), + action=ActionType.MANAGE_PILOTS, + pilot_db=db, + pilot_stamps=["stamp-a"], + allow_legacy_pilots=True, + ) + assert exc_info.value.status_code == 403 + + +async def test_pilot_stamp_check_cross_vo_allowed_for_admin(pilot_db_with_pilots): + """Service administrators manage pilots across VOs, like they read them.""" + async with pilot_db_with_pilots as db: + # Must not raise + await PilotManagementAccessPolicy.policy( + "PilotManagementAccessPolicy", + _user(SERVICE_ADMINISTRATOR, vo="other-vo"), + action=ActionType.MANAGE_PILOTS, + pilot_db=db, + pilot_stamps=["stamp-a"], + ) diff --git a/diracx-routers/tests/pilots/test_management.py b/diracx-routers/tests/pilots/test_management.py new file mode 100644 index 000000000..0b4e447a3 --- /dev/null +++ b/diracx-routers/tests/pilots/test_management.py @@ -0,0 +1,161 @@ +"""Router-level tests for pilot register / update.""" + +from __future__ import annotations + +import pytest + +from diracx.core.models.pilot import PilotMetadata, PilotStatus + +pytestmark = pytest.mark.enabled_dependencies( + [ + "DevelopmentSettings", + "AuthDB", + "AuthSettings", + "ConfigSource", + "PilotAgentsDB", + "PilotManagementAccessPolicy", + "JobDB", + ] +) + +MAIN_VO = "lhcb" + + +@pytest.fixture +def normal_test_client(client_factory): + with client_factory.normal_user() as client: + yield client + + +async def test_register_then_duplicate_then_success(normal_test_client): + """Registering an existing stamp is rejected with 409; a fresh one still succeeds.""" + pilot_stamps = [f"stamps_{i}" for i in range(5)] + + for stamp in pilot_stamps: + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": stamp, "vo": MAIN_VO} + ) + assert r.status_code == 201, r.json() + + # Existing stamp is rejected + r = normal_test_client.post( + "/api/pilots/", + json={ + "pilot_stamp": pilot_stamps[0], + "vo": MAIN_VO, + }, + ) + assert r.status_code == 409, r.json() + + # The new stamp alone was NOT committed by the failing call above + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": "stamps_new", "vo": MAIN_VO} + ) + assert r.status_code == 201, r.json() + + +async def test_register_pilot_unknown_vo_returns_400(normal_test_client): + """Registering into a VO absent from the registry is rejected.""" + r = normal_test_client.post( + "/api/pilots/", + json={"pilot_stamp": "stamp_vo", "vo": "not-a-registered-vo"}, + ) + assert r.status_code == 400, r.json() + + +async def test_update_pilot_metadata_applies_partial_fields(normal_test_client): + """PATCH /pilots/metadata supports heterogeneous field subsets per row.""" + stamps = ["stamp_m1", "stamp_m2"] + for stamp in stamps: + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": stamp, "vo": MAIN_VO} + ) + assert r.status_code == 201 + + # stamp_m1 updates only BenchMark; stamp_m2 only Status + r = normal_test_client.patch( + "/api/pilots/metadata", + json={ + "stamp_m1": PilotMetadata(BenchMark=1.0).model_dump( + by_alias=True, exclude_unset=True + ), + "stamp_m2": PilotMetadata(Status=PilotStatus.WAITING).model_dump( + by_alias=True, exclude_unset=True + ), + }, + ) + assert r.status_code == 204, r.json() + + r = normal_test_client.post("/api/pilots/search", json={}) + assert r.status_code == 200 + by_stamp = {p["PilotStamp"]: p for p in r.json()} + assert by_stamp["stamp_m1"]["BenchMark"] == 1.0 + assert by_stamp["stamp_m1"]["Status"] == PilotStatus.SUBMITTED # untouched + assert by_stamp["stamp_m2"]["Status"] == PilotStatus.WAITING + assert by_stamp["stamp_m2"]["BenchMark"] == 0.0 # untouched + + +async def test_update_pilot_metadata_unknown_stamp_returns_404(normal_test_client): + r = normal_test_client.patch( + "/api/pilots/metadata", + json={"nonexistent": {"Status": PilotStatus.DONE.value}}, + ) + assert r.status_code == 404, r.json() + + +async def test_update_pilot_metadata_refreshes_last_update_time(normal_test_client): + """Any metadata update also bumps LastUpdateTime.""" + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": "stamp_t", "vo": MAIN_VO} + ) + assert r.status_code == 201 + + r = normal_test_client.post( + "/api/pilots/search", json={"parameters": ["LastUpdateTime"]} + ) + before = r.json()[0]["LastUpdateTime"] + + r = normal_test_client.patch( + "/api/pilots/metadata", + json={"stamp_t": {"BenchMark": 1.0}}, + ) + assert r.status_code == 204 + + r = normal_test_client.post( + "/api/pilots/search", json={"parameters": ["LastUpdateTime"]} + ) + assert r.json()[0]["LastUpdateTime"] > before + + +async def test_register_pilot_reference(normal_test_client): + """An explicit reference is stored; without one the stamp is used.""" + r = normal_test_client.post( + "/api/pilots/", + json={"pilot_stamp": "stamp_r1", "vo": MAIN_VO, "pilot_reference": "ref-1"}, + ) + assert r.status_code == 201, r.json() + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": "stamp_r2", "vo": MAIN_VO} + ) + assert r.status_code == 201, r.json() + + r = normal_test_client.post( + "/api/pilots/search", + json={"parameters": ["PilotStamp", "PilotJobReference"]}, + ) + assert r.status_code == 200 + by_stamp = {p["PilotStamp"]: p["PilotJobReference"] for p in r.json()} + assert by_stamp["stamp_r1"] == "ref-1" + assert by_stamp["stamp_r2"] == "stamp_r2" + + +async def test_register_pilot_oversized_stamp_returns_422(normal_test_client): + """A stamp longer than the PilotStamp column (32 chars) is rejected upfront. + + Without the model-level max_length this would only fail inside the + database, and only on backends with strict length enforcement. + """ + r = normal_test_client.post( + "/api/pilots/", json={"pilot_stamp": "s" * 33, "vo": MAIN_VO} + ) + assert r.status_code == 422, r.json() diff --git a/diracx-routers/tests/pilots/test_query.py b/diracx-routers/tests/pilots/test_query.py new file mode 100644 index 000000000..25c9bb01a --- /dev/null +++ b/diracx-routers/tests/pilots/test_query.py @@ -0,0 +1,259 @@ +"""Router-level tests for pilots search/summary and JobID pseudo-parameter.""" + +from __future__ import annotations + +import pytest + +from diracx.core.models.pilot import PilotMetadata, PilotStatus +from diracx.db.sql import PilotAgentsDB +from diracx.logic.pilots.management import assign_jobs_to_pilot + +pytestmark = pytest.mark.enabled_dependencies( + [ + "AuthSettings", + "ConfigSource", + "DevelopmentSettings", + "PilotAgentsDB", + "PilotManagementAccessPolicy", + "JobDB", + ] +) + + +MAIN_VO = "lhcb" +N = 20 + +PILOT_STATUSES = list(PilotStatus) + + +@pytest.fixture +def normal_test_client(client_factory): + with client_factory.normal_user() as client: + yield client + + +@pytest.fixture +async def populated_pilot_client(normal_test_client): + """Client with N pilots registered and metadata patched.""" + pilot_stamps = [f"stamp_{i}" for i in range(1, N + 1)] + + for stamp in pilot_stamps: + r = normal_test_client.post( + "/api/pilots/", + json={"vo": MAIN_VO, "pilot_stamp": stamp}, + ) + assert r.status_code == 201, r.json() + + r = normal_test_client.patch( + "/api/pilots/metadata", + json={ + stamp: PilotMetadata( + BenchMark=float(i), + Status=PILOT_STATUSES[i % len(PILOT_STATUSES)], + Queue=f"queue_{i}", + ).model_dump(by_alias=True, exclude_unset=True) + for i, stamp in enumerate(pilot_stamps) + }, + ) + assert r.status_code == 204, r.json() + return normal_test_client + + +def test_search_returns_pilots_from_own_vo(populated_pilot_client): + r = populated_pilot_client.post("/api/pilots/search", json={}) + assert r.status_code == 200, r.json() + pilots = r.json() + assert len(pilots) == N + assert all(p["VO"] == MAIN_VO for p in pilots) + + +def test_search_filter_by_status(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/search", + json={ + "parameters": ["PilotStamp", "Status"], + "search": [ + { + "parameter": "Status", + "operator": "eq", + "value": PilotStatus.WAITING.value, + } + ], + }, + ) + assert r.status_code == 200, r.json() + pilots = r.json() + assert all(p["Status"] == PilotStatus.WAITING for p in pilots) + + +def test_search_pagination_content_range(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/search?per_page=5&page=1", + json={}, + ) + assert r.status_code == 206 + assert "Content-Range" in r.headers + assert r.headers["Content-Range"] == f"pilots 0-4/{N}" + assert len(r.json()) == 5 + + +def test_summary_groups_by_status(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/summary", json={"grouping": ["Status"]} + ) + assert r.status_code == 200, r.json() + totals = {row["Status"]: row["count"] for row in r.json()} + assert sum(totals.values()) == N + + +# --------------------------------------------------------------------------- +# Cross-table search: JobID pseudo-parameter on POST /api/pilots/search +# --------------------------------------------------------------------------- + + +async def _assign(client, stamp: str, job_ids: list[int]) -> None: + """Insert JobToPilotMapping rows directly via the DB dependency override.""" + db = client.app.dependency_overrides[PilotAgentsDB.transaction].args[0] + async with db: + await assign_jobs_to_pilot(pilot_db=db, pilot_stamp=stamp, job_ids=job_ids) + + +async def test_pilots_search_by_job_id_eq(populated_pilot_client): + """A `JobID` eq filter returns only the pilots that ran that job.""" + await _assign(populated_pilot_client, "stamp_1", [100]) + await _assign(populated_pilot_client, "stamp_2", [100]) + await _assign(populated_pilot_client, "stamp_3", [200]) + + r = populated_pilot_client.post( + "/api/pilots/search", + json={ + "parameters": ["PilotStamp"], + "search": [{"parameter": "JobID", "operator": "eq", "value": 100}], + }, + ) + assert r.status_code == 200, r.json() + stamps = sorted(p["PilotStamp"] for p in r.json()) + assert stamps == ["stamp_1", "stamp_2"] + + +async def test_pilots_search_by_job_id_in(populated_pilot_client): + """An `in` filter over several job IDs returns the union of their pilots.""" + await _assign(populated_pilot_client, "stamp_4", [300]) + await _assign(populated_pilot_client, "stamp_5", [301]) + + r = populated_pilot_client.post( + "/api/pilots/search", + json={ + "parameters": ["PilotStamp"], + "search": [ + { + "parameter": "JobID", + "operator": "in", + "values": [300, 301], + } + ], + }, + ) + assert r.status_code == 200, r.json() + stamps = sorted(p["PilotStamp"] for p in r.json()) + assert stamps == ["stamp_4", "stamp_5"] + + +def test_pilots_search_by_unknown_job_id_returns_empty(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/search", + json={"search": [{"parameter": "JobID", "operator": "eq", "value": 999999}]}, + ) + assert r.status_code == 200 + assert r.json() == [] + + +def test_pilots_search_job_id_unsupported_operator_raises(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/search", + json={"search": [{"parameter": "JobID", "operator": "neq", "value": 1}]}, + ) + assert r.status_code in (400, 422), r.json() + + +def test_pilots_search_job_id_non_numeric_value_raises(populated_pilot_client): + """A non-numeric `JobID` value is a 400, not a 500.""" + r = populated_pilot_client.post( + "/api/pilots/search", + json={"search": [{"parameter": "JobID", "operator": "eq", "value": "abc"}]}, + ) + assert r.status_code in (400, 422), r.json() + + +def test_pilots_search_combining_job_id_and_pilot_id_raises(populated_pilot_client): + r = populated_pilot_client.post( + "/api/pilots/search", + json={ + "search": [ + {"parameter": "JobID", "operator": "eq", "value": 1}, + {"parameter": "PilotID", "operator": "eq", "value": 1}, + ] + }, + ) + assert r.status_code in (400, 422), r.json() + + +# --------------------------------------------------------------------------- +# VO isolation +# +# HTTP-level tests stub the access policy with AlwaysAllowAccessPolicy, so +# the read scoping cannot be fully exercised through the routes here. The +# logic layer is tested directly instead: it must apply the VO constraint +# that the route derives from the caller's identity. +# --------------------------------------------------------------------------- + +OTHER_VO = "gridpp" + + +def _pilot_db(client) -> PilotAgentsDB: + return client.app.dependency_overrides[PilotAgentsDB.transaction].args[0] + + +async def test_search_vo_constraint_hides_other_vos(populated_pilot_client): + """A VO-constrained search must not leak pilots from another VO.""" + from diracx.logic.pilots import search as search_bl + + db = _pilot_db(populated_pilot_client) + async with db: + await db.register_pilots(pilot_stamps=["other-vo-stamp"], vo=OTHER_VO) + + async with db: + total, pilots = await search_bl(pilot_db=db, vo_constraint=MAIN_VO) + assert total == N + assert {p["VO"] for p in pilots} == {MAIN_VO} + + # Unconstrained (service administrator) searches see every VO + async with db: + total, pilots = await search_bl(pilot_db=db, vo_constraint=None) + assert total == N + 1 + assert {p["VO"] for p in pilots} == {MAIN_VO, OTHER_VO} + + +async def test_summary_vo_constraint_hides_other_vos(populated_pilot_client): + """A VO-constrained summary must not count pilots from another VO.""" + from diracx.core.models.search import SummaryParams + from diracx.logic.pilots import summary as summary_bl + + db = _pilot_db(populated_pilot_client) + async with db: + await db.register_pilots(pilot_stamps=["other-vo-stamp"], vo=OTHER_VO) + + async with db: + rows = await summary_bl( + pilot_db=db, body=SummaryParams(grouping=["VO"]), vo_constraint=MAIN_VO + ) + assert rows == [{"VO": MAIN_VO, "count": N}] + + async with db: + rows = await summary_bl( + pilot_db=db, body=SummaryParams(grouping=["VO"]), vo_constraint=None + ) + assert sorted(rows, key=lambda r: r["VO"]) == [ + {"VO": OTHER_VO, "count": 1}, + {"VO": MAIN_VO, "count": N}, + ] diff --git a/diracx-tasks/tests/test_worker_integration.py b/diracx-tasks/tests/test_worker_integration.py index 37ac9c1dc..26232595e 100644 --- a/diracx-tasks/tests/test_worker_integration.py +++ b/diracx-tasks/tests/test_worker_integration.py @@ -313,7 +313,7 @@ async def test_process_message_acks_on_unknown_task( async def test_process_message_acks_on_parse_error( broker, task_class_registry, wrapped_registry ): - """Unparseable messages should be acked and logged, not crash the worker.""" + """Unparsable messages should be acked and logged, not crash the worker.""" worker = Worker( broker=broker, task_registry=wrapped_registry, diff --git a/diracx-testing/src/diracx/testing/utils.py b/diracx-testing/src/diracx/testing/utils.py index 9816457f7..d851cb662 100644 --- a/diracx-testing/src/diracx/testing/utils.py +++ b/diracx-testing/src/diracx/testing/utils.py @@ -576,7 +576,7 @@ def demo_urls(demo_dir): yield helm_values["developer"]["urls"] except KeyError: # If we're testing an extension they will have an umbrella chart so the - # URLs are under a "diracx" section corrosponding to the vanilla chart + # URLs are under a "diracx" section corresponding to the vanilla chart yield helm_values["diracx"]["developer"]["urls"] diff --git a/docs/admin/explanations/pilots.md b/docs/admin/explanations/pilots.md new file mode 100644 index 000000000..b3a7783ef --- /dev/null +++ b/docs/admin/explanations/pilots.md @@ -0,0 +1,11 @@ +# Pilots from the operator's perspective + +This page is for operators running a DiracX instance. For the developer view (identity model, lifecycle, cross-table search, extension points) see the [developer explanation](../../dev/explanations/pilots.md). + +## VO scoping + +Pilots are partitioned by VO. By default a user only sees and acts on pilots belonging to their own VO. Users holding the +`SERVICE_ADMINISTRATOR` security property bypass that filter and can read pilots across every VO through the same search and summary endpoints. + +Management actions (register, patch metadata) are restricted to `SERVICE_ADMINISTRATOR`, who — as for reads — may act across VOs. Legacy X.509 pilot identities (`GENERIC_PILOT` property) may additionally self-register and self-update within their own VO — pilots started in the vacuum have no SiteDirector to register them, mirroring `dirac-admin-add-pilot` in legacy DIRAC. Those routes opt in via `allow_legacy_pilots=True` in the access policy, which caps each call to a single pilot stamp to limit the blast radius of a stolen credential (note this bounds the rate of abuse, not its scope: a legacy pilot identity is not bound to its own stamp and can act on any pilot in its VO). +See [authentication with DiracX](auth-with-diracx.md) for the full pilot auth story (X.509, DiracX pilot tokens). diff --git a/docs/admin/reference/env-variables.md b/docs/admin/reference/env-variables.md index e2b2b6285..55792886c 100644 --- a/docs/admin/reference/env-variables.md +++ b/docs/admin/reference/env-variables.md @@ -62,6 +62,12 @@ Enable the HEALTH router Enable the JOBS router +#### `DIRACX_SERVICE_PILOTS_ENABLED` + +*Optional*, default value: `True` + +Enable the PILOTS router + #### `DIRACX_SERVICE_RSS_ENABLED` *Optional*, default value: `True` diff --git a/docs/dev/explanations/pilots.md b/docs/dev/explanations/pilots.md new file mode 100644 index 000000000..cc143c1ba --- /dev/null +++ b/docs/dev/explanations/pilots.md @@ -0,0 +1,55 @@ +# Pilots + +## What is a pilot + +A pilot is a small piece of software that runs on a *worker node* and pulls user payloads (jobs). Two authentication modes are relevant to this service: + +- **X.509 proxy** (legacy): the pilot presents a proxy and exchanges it for a DiracX token. Callers authenticated this way carry the `GENERIC_PILOT` property and are handled by the "legacy pilot" code paths in the access policy. +- **Pre-issued secret**: the pilot is provisioned with a secret that it exchanges for a DiracX token. Pilots authenticated this way are identified by their unique *stamp* rather than by a set of security properties. + +## Identity model + +Three identifiers appear throughout the code and are easy to confuse: + +- `PilotStamp`: immutable string chosen by the pilot factory. Primary user-facing key; never changes for the lifetime of a pilot. +- `PilotID`: auto-incrementing database primary key. It appears in search filters and results, but management routes always key on the stamp. +- `PilotJobReference`: the CE job reference (batch-system identifier) + that submitted the pilot process. Defaults to the stamp when not known. + +## Lifecycle + +```mermaid +stateDiagram-v2 + [*] --> Submitted + Submitted --> Waiting + Waiting --> Running + Running --> Done + Running --> Failed + Submitted --> Aborted + Waiting --> Aborted + Running --> Aborted + Done --> [*] + Failed --> [*] + Aborted --> [*] + [*] --> Unknown + Unknown --> [*] +``` + +## Relationship to jobs + +A pilot can execute zero or more jobs over its lifetime. The association is tracked in the `JobToPilotMapping` table and is append-only: once a job has run on a pilot, the link is preserved until the pilot row is deleted. + +Both directions of the lookup are exposed as *pseudo-parameters* on the respective search endpoints. This keeps every pilot and job attribute addressable through a single `POST /search` per resource type, matching the UI's one-search-bar-per-resource mental model. The pattern mirrors the existing `LoggingInfo` pseudo-parameter on `POST /api/jobs/search`: the filter is intercepted in the logic layer, resolved against `JobToPilotMapping`, and rewritten into a normal vector filter before hitting the DB. + +- `POST /api/jobs/search` accepts a `PilotStamp` filter, resolved to a `JobID` filter via `JobToPilotMapping`. +- `POST /api/pilots/search` accepts a `JobID` filter, resolved to a `PilotID` filter. + +Concrete request bodies for both are provided as OpenAPI examples on the respective search routes; open the Swagger UI at `/api/docs` to see them. + +Only `eq` and `in` operators are supported on the pseudo-parameter; other operators (`neq`, `not in`, `lt`, ...) are refused with `InvalidQueryError` because their semantics across the join are ambiguous. Combining a `PilotStamp` filter with a `JobID` filter in the same request body is likewise refused; clients that want the intersection should compute it themselves. + +## VO scoping and authorization + +Pilots are partitioned by VO. By default a normal user sees and acts on pilots belonging to their own VO only. `SERVICE_ADMINISTRATOR` can read pilots across VOs via `/search` and `/summary`. + +Management actions (register, patch metadata) require `SERVICE_ADMINISTRATOR`, who — as for reads — may act across VOs. Legacy X.509 pilot identities may additionally self-register and self-update within their own VO — pilots started in the vacuum have no SiteDirector to register them, mirroring `dirac-admin-add-pilot` in legacy DIRAC. Those paths opt in via `allow_legacy_pilots=True` in the access policy, which caps each call to a single pilot stamp to limit the blast radius of a stolen credential (this bounds the rate of abuse, not its scope: a legacy pilot identity is not bound to its own stamp and can act on any pilot in its VO). diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/_client.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/_client.py index 04db73e5f..606ad38aa 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/_client.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/_client.py @@ -21,6 +21,7 @@ JobsOperations, LollygagOperations, MyOperations, + PilotsOperations, RssOperations, WellKnownOperations, ) @@ -41,6 +42,8 @@ class Dirac: # pylint: disable=client-accepts-api-version-keyword,too-many-inst :vartype lollygag: _generated.operations.LollygagOperations :ivar my: MyOperations operations :vartype my: _generated.operations.MyOperations + :ivar pilots: PilotsOperations operations + :vartype pilots: _generated.operations.PilotsOperations :ivar rss: RssOperations operations :vartype rss: _generated.operations.RssOperations :keyword endpoint: Service URL. Required. Default value is "". @@ -81,6 +84,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self.jobs = JobsOperations(self._client, self._config, self._serialize, self._deserialize) self.lollygag = LollygagOperations(self._client, self._config, self._serialize, self._deserialize) self.my = MyOperations(self._client, self._config, self._serialize, self._deserialize) + self.pilots = PilotsOperations(self._client, self._config, self._serialize, self._deserialize) self.rss = RssOperations(self._client, self._config, self._serialize, self._deserialize) def send_request(self, request: HttpRequest, *, stream: bool = False, **kwargs: Any) -> HttpResponse: diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/_client.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/_client.py index 1adec809d..4a6004c25 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/_client.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/_client.py @@ -21,6 +21,7 @@ JobsOperations, LollygagOperations, MyOperations, + PilotsOperations, RssOperations, WellKnownOperations, ) @@ -41,6 +42,8 @@ class Dirac: # pylint: disable=client-accepts-api-version-keyword,too-many-inst :vartype lollygag: _generated.aio.operations.LollygagOperations :ivar my: MyOperations operations :vartype my: _generated.aio.operations.MyOperations + :ivar pilots: PilotsOperations operations + :vartype pilots: _generated.aio.operations.PilotsOperations :ivar rss: RssOperations operations :vartype rss: _generated.aio.operations.RssOperations :keyword endpoint: Service URL. Required. Default value is "". @@ -81,6 +84,7 @@ def __init__( # pylint: disable=missing-client-constructor-parameter-credential self.jobs = JobsOperations(self._client, self._config, self._serialize, self._deserialize) self.lollygag = LollygagOperations(self._client, self._config, self._serialize, self._deserialize) self.my = MyOperations(self._client, self._config, self._serialize, self._deserialize) + self.pilots = PilotsOperations(self._client, self._config, self._serialize, self._deserialize) self.rss = RssOperations(self._client, self._config, self._serialize, self._deserialize) def send_request( diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/__init__.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/__init__.py index d7d250107..889422f07 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/__init__.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/__init__.py @@ -16,6 +16,7 @@ from ._operations import JobsOperations # type: ignore from ._operations import LollygagOperations # type: ignore from ._operations import MyOperations # type: ignore +from ._operations import PilotsOperations # type: ignore from ._operations import RssOperations # type: ignore from ._patch import __all__ as _patch_all @@ -29,6 +30,7 @@ "JobsOperations", "LollygagOperations", "MyOperations", + "PilotsOperations", "RssOperations", ] __all__.extend([p for p in _patch_all if p not in __all__]) # pyright: ignore diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/_operations.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/_operations.py index fb42de5d0..2d0b0cfde 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/_operations.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/aio/operations/_operations.py @@ -56,6 +56,10 @@ build_lollygag_insert_owner_object_request, build_my_pilots_get_pilot_summary_request, build_my_pilots_submit_pilot_request, + build_pilots_register_pilot_request, + build_pilots_search_request, + build_pilots_summary_request, + build_pilots_update_pilot_metadata_request, build_rss_get_compute_status_request, build_rss_get_fts_status_request, build_rss_get_site_status_request, @@ -1962,6 +1966,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: ~_generated.models.SearchParams :keyword page: Default value is 1. @@ -2002,6 +2012,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: IO[bytes] :keyword page: Default value is 1. @@ -2041,6 +2057,12 @@ async def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. :type body: ~_generated.models.SearchParams or IO[bytes] :keyword page: Default value is 1. @@ -2318,7 +2340,7 @@ async def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any response = pipeline_response.http_response - if response.status_code not in [200]: + if response.status_code not in [201]: map_error(status_code=response.status_code, response=response, error_map=error_map) raise HttpResponseError(response=response) @@ -2611,6 +2633,500 @@ async def pilots_get_pilot_summary(self, **kwargs: Any) -> dict[str, int]: return deserialized # type: ignore +class PilotsOperations: + """ + .. warning:: + **DO NOT** instantiate this class directly. + + Instead, you should access the following operations through + :class:`~_generated.aio.Dirac`'s + :attr:`pilots` attribute. + """ + + models = _models + + def __init__(self, *args, **kwargs) -> None: + input_args = list(args) + self._client: AsyncPipelineClient = input_args.pop(0) if input_args else kwargs.pop("client") + self._config: DiracConfiguration = input_args.pop(0) if input_args else kwargs.pop("config") + self._serialize: Serializer = input_args.pop(0) if input_args else kwargs.pop("serializer") + self._deserialize: Deserializer = input_args.pop(0) if input_args else kwargs.pop("deserializer") + + @overload + async def register_pilot( + self, body: _models.PilotRegistrationParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Required. + :type body: ~_generated.models.PilotRegistrationParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def register_pilot(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def register_pilot(self, body: Union[_models.PilotRegistrationParams, IO[bytes]], **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Is either a PilotRegistrationParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.PilotRegistrationParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "PilotRegistrationParams") + + _request = build_pilots_register_pilot_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 [201]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + @overload + async def update_pilot_metadata( + self, body: dict[str, _models.PilotMetadata], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Required. + :type body: dict[str, ~_generated.models.PilotMetadata] + :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 update_pilot_metadata( + self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :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 update_pilot_metadata( + self, body: Union[dict[str, _models.PilotMetadata], IO[bytes]], **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Is either a {str: PilotMetadata} type or a IO[bytes] type. Required. + :type body: dict[str, ~_generated.models.PilotMetadata] 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, "{PilotMetadata}") + + _request = build_pilots_update_pilot_metadata_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 + + @overload + async def search( + self, + body: Optional[_models.SearchParams] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: ~_generated.models.SearchParams + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def search( + self, + body: Optional[IO[bytes]] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def search( + self, + body: Optional[Union[_models.SearchParams, IO[bytes]]] = None, + *, + page: int = 1, + per_page: int = 100, + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. + :type body: ~_generated.models.SearchParams or IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :return: list of dict mapping str to any + :rtype: list[dict[str, any]] + :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)) + content_type = content_type if body else None + cls: ClsType[list[dict[str, Any]]] = kwargs.pop("cls", None) + + content_type = content_type or "application/json" if body else None + _json = None + _content = None + if isinstance(body, (IOBase, bytes)): + _content = body + else: + if body is not None: + _json = self._serialize.body(body, "SearchParams") + else: + _json = None + + _request = build_pilots_search_request( + page=page, + per_page=per_page, + 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, 206]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + response_headers = {} + if response.status_code == 206: + response_headers["Content-Range"] = self._deserialize("str", response.headers.get("Content-Range")) + + deserialized = self._deserialize("[{object}]", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + async def summary( + self, body: _models.SummaryParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Required. + :type body: ~_generated.models.SummaryParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + async def summary(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace_async + async def summary(self, body: Union[_models.SummaryParams, IO[bytes]], **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Is either a SummaryParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.SummaryParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "SummaryParams") + + _request = build_pilots_summary_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 [200]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + class RssOperations: """ .. warning:: diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/__init__.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/__init__.py index 6684af567..22b9d7249 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/__init__.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/__init__.py @@ -32,6 +32,8 @@ JobMetaData, JobStatusUpdate, OpenIDConfiguration, + PilotMetadata, + PilotRegistrationParams, SandboxDownloadResponse, SandboxInfo, SandboxUploadResponse, @@ -61,6 +63,7 @@ from ._enums import ( # type: ignore ChecksumAlgorithm, JobStatus, + PilotStatus, SandboxFormat, SandboxType, ScalarSearchOperator, @@ -92,6 +95,8 @@ "JobMetaData", "JobStatusUpdate", "OpenIDConfiguration", + "PilotMetadata", + "PilotRegistrationParams", "SandboxDownloadResponse", "SandboxInfo", "SandboxUploadResponse", @@ -118,6 +123,7 @@ "VectorSearchSpec", "ChecksumAlgorithm", "JobStatus", + "PilotStatus", "SandboxFormat", "SandboxType", "ScalarSearchOperator", diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_enums.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_enums.py index b83473639..849d3252a 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_enums.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_enums.py @@ -34,6 +34,19 @@ class JobStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta): RESCHEDULED = "Rescheduled" +class PilotStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta): + """PilotStatus.""" + + SUBMITTED = "Submitted" + WAITING = "Waiting" + RUNNING = "Running" + DONE = "Done" + FAILED = "Failed" + DELETED = "Deleted" + ABORTED = "Aborted" + UNKNOWN = "Unknown" + + class SandboxFormat(str, Enum, metaclass=CaseInsensitiveEnumMeta): """SandboxFormat.""" diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_models.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_models.py index 6953a050f..a29eb56ad 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_models.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/models/_models.py @@ -1074,6 +1074,185 @@ def __init__( self.code_challenge_methods_supported = code_challenge_methods_supported +class PilotMetadata(_serialization.Model): + """Mutable metadata attached to a pilot. + + The pilot is identified by its stamp, passed alongside this model + (e.g. as the mapping key on ``PATCH /api/pilots/metadata``\\ ). Every + field is optional; when absent it is left untouched by an update. + + :ivar status_reason: Human-readable reason for the current status. + :vartype status_reason: str + :ivar status: Current pilot status. Known values are: "Submitted", "Waiting", "Running", + "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :vartype status: str or ~_generated.models.PilotStatus + :ivar bench_mark: Pilot benchmark value. + :vartype bench_mark: float + :ivar destination_site: Destination site. + :vartype destination_site: str + :ivar queue: Batch queue name. + :vartype queue: str + :ivar grid_site: Grid site. + :vartype grid_site: str + :ivar grid_type: Grid type. + :vartype grid_type: str + :ivar accounting_sent: Whether accounting has been sent for this pilot. + :vartype accounting_sent: bool + :ivar current_job_id: ID of the job currently running on this pilot. + :vartype current_job_id: int + """ + + _validation = { + "status_reason": {"max_length": 255}, + "destination_site": {"max_length": 128}, + "queue": {"max_length": 128}, + "grid_site": {"max_length": 128}, + "grid_type": {"max_length": 32}, + } + + _attribute_map = { + "status_reason": {"key": "StatusReason", "type": "str"}, + "status": {"key": "Status", "type": "str"}, + "bench_mark": {"key": "BenchMark", "type": "float"}, + "destination_site": {"key": "DestinationSite", "type": "str"}, + "queue": {"key": "Queue", "type": "str"}, + "grid_site": {"key": "GridSite", "type": "str"}, + "grid_type": {"key": "GridType", "type": "str"}, + "accounting_sent": {"key": "AccountingSent", "type": "bool"}, + "current_job_id": {"key": "CurrentJobID", "type": "int"}, + } + + def __init__( + self, + *, + status_reason: Optional[str] = None, + status: Optional[Union[str, "_models.PilotStatus"]] = None, + bench_mark: Optional[float] = None, + destination_site: Optional[str] = None, + queue: Optional[str] = None, + grid_site: Optional[str] = None, + grid_type: Optional[str] = None, + accounting_sent: Optional[bool] = None, + current_job_id: Optional[int] = None, + **kwargs: Any + ) -> None: + """ + :keyword status_reason: Human-readable reason for the current status. + :paramtype status_reason: str + :keyword status: Current pilot status. Known values are: "Submitted", "Waiting", "Running", + "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :paramtype status: str or ~_generated.models.PilotStatus + :keyword bench_mark: Pilot benchmark value. + :paramtype bench_mark: float + :keyword destination_site: Destination site. + :paramtype destination_site: str + :keyword queue: Batch queue name. + :paramtype queue: str + :keyword grid_site: Grid site. + :paramtype grid_site: str + :keyword grid_type: Grid type. + :paramtype grid_type: str + :keyword accounting_sent: Whether accounting has been sent for this pilot. + :paramtype accounting_sent: bool + :keyword current_job_id: ID of the job currently running on this pilot. + :paramtype current_job_id: int + """ + super().__init__(**kwargs) + self.status_reason = status_reason + self.status = status + self.bench_mark = bench_mark + self.destination_site = destination_site + self.queue = queue + self.grid_site = grid_site + self.grid_type = grid_type + self.accounting_sent = accounting_sent + self.current_job_id = current_job_id + + +class PilotRegistrationParams(_serialization.Model): + """Body of ``POST /api/pilots/`` to register a single pilot. + + The ``max_length`` constraints mirror the column sizes of the legacy + ``PilotAgents`` table so that oversized values are rejected with a 422 + instead of a backend-dependent database error. + + All required parameters must be populated in order to send to server. + + :ivar pilot_stamp: Stamp of the pilot to create. Required. + :vartype pilot_stamp: str + :ivar vo: Pilot virtual organization. Required. + :vartype vo: str + :ivar grid_type: Grid type of the pilot. + :vartype grid_type: str + :ivar grid_site: Pilot grid site. + :vartype grid_site: str + :ivar destination_site: Pilot destination site. + :vartype destination_site: str + :ivar pilot_reference: CE job reference of the pilot; defaults to the stamp. + :vartype pilot_reference: str + :ivar pilot_status: Initial status of the pilot. Known values are: "Submitted", "Waiting", + "Running", "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :vartype pilot_status: str or ~_generated.models.PilotStatus + """ + + _validation = { + "pilot_stamp": {"required": True, "max_length": 32, "min_length": 1}, + "vo": {"required": True, "max_length": 128, "min_length": 1}, + "grid_type": {"max_length": 32}, + "grid_site": {"max_length": 128}, + "destination_site": {"max_length": 128}, + "pilot_reference": {"max_length": 255}, + } + + _attribute_map = { + "pilot_stamp": {"key": "pilot_stamp", "type": "str"}, + "vo": {"key": "vo", "type": "str"}, + "grid_type": {"key": "grid_type", "type": "str"}, + "grid_site": {"key": "grid_site", "type": "str"}, + "destination_site": {"key": "destination_site", "type": "str"}, + "pilot_reference": {"key": "pilot_reference", "type": "str"}, + "pilot_status": {"key": "pilot_status", "type": "str"}, + } + + def __init__( + self, + *, + pilot_stamp: str, + vo: str, + grid_type: str = "DIRAC", + grid_site: str = "Unknown", + destination_site: str = "NotAssigned", + pilot_reference: Optional[str] = None, + pilot_status: Optional[Union[str, "_models.PilotStatus"]] = None, + **kwargs: Any + ) -> None: + """ + :keyword pilot_stamp: Stamp of the pilot to create. Required. + :paramtype pilot_stamp: str + :keyword vo: Pilot virtual organization. Required. + :paramtype vo: str + :keyword grid_type: Grid type of the pilot. + :paramtype grid_type: str + :keyword grid_site: Pilot grid site. + :paramtype grid_site: str + :keyword destination_site: Pilot destination site. + :paramtype destination_site: str + :keyword pilot_reference: CE job reference of the pilot; defaults to the stamp. + :paramtype pilot_reference: str + :keyword pilot_status: Initial status of the pilot. Known values are: "Submitted", "Waiting", + "Running", "Done", "Failed", "Deleted", "Aborted", and "Unknown". + :paramtype pilot_status: str or ~_generated.models.PilotStatus + """ + super().__init__(**kwargs) + self.pilot_stamp = pilot_stamp + self.vo = vo + self.grid_type = grid_type + self.grid_site = grid_site + self.destination_site = destination_site + self.pilot_reference = pilot_reference + self.pilot_status = pilot_status + + class SandboxDownloadResponse(_serialization.Model): """SandboxDownloadResponse. diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/__init__.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/__init__.py index d7d250107..889422f07 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/__init__.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/__init__.py @@ -16,6 +16,7 @@ from ._operations import JobsOperations # type: ignore from ._operations import LollygagOperations # type: ignore from ._operations import MyOperations # type: ignore +from ._operations import PilotsOperations # type: ignore from ._operations import RssOperations # type: ignore from ._patch import __all__ as _patch_all @@ -29,6 +30,7 @@ "JobsOperations", "LollygagOperations", "MyOperations", + "PilotsOperations", "RssOperations", ] __all__.extend([p for p in _patch_all if p not in __all__]) # pyright: ignore diff --git a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/_operations.py b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/_operations.py index 7866a140a..bb92cb8b2 100644 --- a/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/_operations.py +++ b/extensions/gubbins/gubbins-client/src/gubbins/client/_generated/operations/_operations.py @@ -647,6 +647,78 @@ def build_my_pilots_get_pilot_summary_request(**kwargs: Any) -> HttpRequest: # return HttpRequest(method="GET", url=_url, headers=_headers, **kwargs) +def build_pilots_register_pilot_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/pilots/" + + # 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="POST", url=_url, headers=_headers, **kwargs) + + +def build_pilots_update_pilot_metadata_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/metadata" + + # 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) + + +def build_pilots_search_request(*, page: int = 1, per_page: int = 100, **kwargs: Any) -> HttpRequest: + _headers = case_insensitive_dict(kwargs.pop("headers", {}) or {}) + _params = case_insensitive_dict(kwargs.pop("params", {}) or {}) + + content_type: Optional[str] = kwargs.pop("content_type", _headers.pop("Content-Type", None)) + accept = _headers.pop("Accept", "application/json") + + # Construct URL + _url = "/api/pilots/search" + + # Construct parameters + if page is not None: + _params["page"] = _SERIALIZER.query("page", page, "int", minimum=1) + if per_page is not None: + _params["per_page"] = _SERIALIZER.query("per_page", per_page, "int", maximum=10000, minimum=1) + + # 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="POST", url=_url, params=_params, headers=_headers, **kwargs) + + +def build_pilots_summary_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/pilots/summary" + + # 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="POST", url=_url, headers=_headers, **kwargs) + + def build_rss_get_storage_status_request( *, if_modified_since: Optional[str] = None, @@ -2648,6 +2720,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: ~_generated.models.SearchParams :keyword page: Default value is 1. @@ -2688,6 +2766,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Default value is None. :type body: IO[bytes] :keyword page: Default value is 1. @@ -2727,6 +2811,12 @@ def search( By default, the search will return all jobs the user has access to, and all the fields of the job will be returned. + A ``PilotStamp`` pseudo-parameter is also accepted in the ``search`` + filter list (operators ``eq`` / ``in`` only): it is transparently + resolved through ``JobToPilotMapping`` into a ``JobID`` filter, + allowing callers to ask "jobs run by this pilot" through the same + endpoint. + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. :type body: ~_generated.models.SearchParams or IO[bytes] :keyword page: Default value is 1. @@ -3002,7 +3092,7 @@ def submit_jdl_jobs(self, body: Union[list[str], IO[bytes]], **kwargs: Any) -> l response = pipeline_response.http_response - if response.status_code not in [200]: + if response.status_code not in [201]: map_error(status_code=response.status_code, response=response, error_map=error_map) raise HttpResponseError(response=response) @@ -3295,6 +3385,496 @@ def pilots_get_pilot_summary(self, **kwargs: Any) -> dict[str, int]: return deserialized # type: ignore +class PilotsOperations: + """ + .. warning:: + **DO NOT** instantiate this class directly. + + Instead, you should access the following operations through + :class:`~_generated.Dirac`'s + :attr:`pilots` attribute. + """ + + models = _models + + def __init__(self, *args, **kwargs) -> None: + input_args = list(args) + self._client: PipelineClient = input_args.pop(0) if input_args else kwargs.pop("client") + self._config: DiracConfiguration = input_args.pop(0) if input_args else kwargs.pop("config") + self._serialize: Serializer = input_args.pop(0) if input_args else kwargs.pop("serializer") + self._deserialize: Deserializer = input_args.pop(0) if input_args else kwargs.pop("deserializer") + + @overload + def register_pilot( + self, body: _models.PilotRegistrationParams, *, content_type: str = "application/json", **kwargs: Any + ) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Required. + :type body: ~_generated.models.PilotRegistrationParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def register_pilot(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def register_pilot(self, body: Union[_models.PilotRegistrationParams, IO[bytes]], **kwargs: Any) -> Any: + """Register Pilot. + + Register a pilot with its reference. + + If the stamp already exists, the registration is rejected with a 409. + Registering into a VO that is not in the registry is rejected with a 400. + + :param body: Is either a PilotRegistrationParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.PilotRegistrationParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "PilotRegistrationParams") + + _request = build_pilots_register_pilot_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 [201]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + deserialized = self._deserialize("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + @overload + def update_pilot_metadata( + self, body: dict[str, _models.PilotMetadata], *, content_type: str = "application/json", **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Required. + :type body: dict[str, ~_generated.models.PilotMetadata] + :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 update_pilot_metadata(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :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 update_pilot_metadata( # pylint: disable=inconsistent-return-statements + self, body: Union[dict[str, _models.PilotMetadata], IO[bytes]], **kwargs: Any + ) -> None: + """Update Pilot Metadata. + + Update pilot metadata (status, benchmark, etc.). + + Only fields defined in ``PilotMetadata`` are mutable. The pilot stamp + (the mapping key) identifies the pilot and cannot be changed. + + :param body: Is either a {str: PilotMetadata} type or a IO[bytes] type. Required. + :type body: dict[str, ~_generated.models.PilotMetadata] 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, "{PilotMetadata}") + + _request = build_pilots_update_pilot_metadata_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 + + @overload + def search( + self, + body: Optional[_models.SearchParams] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: ~_generated.models.SearchParams + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def search( + self, + body: Optional[IO[bytes]] = None, + *, + page: int = 1, + per_page: int = 100, + content_type: str = "application/json", + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Default value is None. + :type body: IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :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 dict mapping str to any + :rtype: list[dict[str, any]] + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def search( + self, + body: Optional[Union[_models.SearchParams, IO[bytes]]] = None, + *, + page: int = 1, + per_page: int = 100, + **kwargs: Any + ) -> list[dict[str, Any]]: + """Search. + + Retrieve information about pilots. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + A ``JobID`` pseudo-parameter is also accepted in the ``search`` filter + list (operators ``eq`` / ``in`` only): it is transparently resolved + through ``JobToPilotMapping`` into a ``PilotID`` filter, allowing + callers to ask "pilots that ran this job" through the same endpoint. + + :param body: Is either a SearchParams type or a IO[bytes] type. Default value is None. + :type body: ~_generated.models.SearchParams or IO[bytes] + :keyword page: Default value is 1. + :paramtype page: int + :keyword per_page: Default value is 100. + :paramtype per_page: int + :return: list of dict mapping str to any + :rtype: list[dict[str, any]] + :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)) + content_type = content_type if body else None + cls: ClsType[list[dict[str, Any]]] = kwargs.pop("cls", None) + + content_type = content_type or "application/json" if body else None + _json = None + _content = None + if isinstance(body, (IOBase, bytes)): + _content = body + else: + if body is not None: + _json = self._serialize.body(body, "SearchParams") + else: + _json = None + + _request = build_pilots_search_request( + page=page, + per_page=per_page, + 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, 206]: + map_error(status_code=response.status_code, response=response, error_map=error_map) + raise HttpResponseError(response=response) + + response_headers = {} + if response.status_code == 206: + response_headers["Content-Range"] = self._deserialize("str", response.headers.get("Content-Range")) + + deserialized = self._deserialize("[{object}]", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, response_headers) # type: ignore + + return deserialized # type: ignore + + @overload + def summary(self, body: _models.SummaryParams, *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Required. + :type body: ~_generated.models.SummaryParams + :keyword content_type: Body Parameter content-type. Content type parameter for JSON body. + Default value is "application/json". + :paramtype content_type: str + :return: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @overload + def summary(self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :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: any + :rtype: any + :raises ~azure.core.exceptions.HttpResponseError: + """ + + @distributed_trace + def summary(self, body: Union[_models.SummaryParams, IO[bytes]], **kwargs: Any) -> Any: + """Summary. + + Aggregate pilot counts suitable for plotting. + + Normal users see only their own VO's pilots. Service administrators see + pilots from all VOs. + + :param body: Is either a SummaryParams type or a IO[bytes] type. Required. + :type body: ~_generated.models.SummaryParams or IO[bytes] + :return: any + :rtype: any + :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[Any] = 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, "SummaryParams") + + _request = build_pilots_summary_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("object", pipeline_response.http_response) + + if cls: + return cls(pipeline_response, deserialized, {}) # type: ignore + + return deserialized # type: ignore + + class RssOperations: """ .. warning:: diff --git a/extensions/gubbins/gubbins-core/tests/test_config.py b/extensions/gubbins/gubbins-core/tests/test_config.py index 9833decb3..9e7209946 100644 --- a/extensions/gubbins/gubbins-core/tests/test_config.py +++ b/extensions/gubbins/gubbins-core/tests/test_config.py @@ -22,7 +22,7 @@ def github_is_down(): return True -@pytest.mark.skipif(github_is_down(), reason="Github unavailble") +@pytest.mark.skipif(github_is_down(), reason="Github unavailable") def test_remote_git_config_source(monkeypatch): monkeypatch.setattr( "diracx.core.config.sources.DEFAULT_CONFIG_FILE", diff --git a/extensions/gubbins/gubbins-routers/tests/test_gubbins_job_manager.py b/extensions/gubbins/gubbins-routers/tests/test_gubbins_job_manager.py index 11caaa145..dde97955d 100644 --- a/extensions/gubbins/gubbins-routers/tests/test_gubbins_job_manager.py +++ b/extensions/gubbins/gubbins-routers/tests/test_gubbins_job_manager.py @@ -23,6 +23,8 @@ "ConfigSource", "TaskQueueDB", "DevelopmentSettings", + "PilotAgentsDB", + "PilotManagementAccessPolicy", ] ) @@ -62,7 +64,7 @@ def valid_job_id(normal_user_client: TestClient): """ job_definitions = [TEST_JDL] r = normal_user_client.post("/api/jobs/jdl", json=job_definitions) - assert r.status_code == 200, r.json() + assert r.status_code == 201, r.json() assert len(r.json()) == 1 return r.json()[0]["JobID"] diff --git a/extensions/gubbins/gubbins-routers/tests/test_wellknown.py b/extensions/gubbins/gubbins-routers/tests/test_wellknown.py index 40c642437..aefdd11ac 100644 --- a/extensions/gubbins/gubbins-routers/tests/test_wellknown.py +++ b/extensions/gubbins/gubbins-routers/tests/test_wellknown.py @@ -19,7 +19,7 @@ def test_client(client_factory): async def test_dirac_metadata_is_overwriten(test_client): """ - Makes sure that the dirac-metadata endpoint is properly overwriten + Makes sure that the dirac-metadata endpoint is properly overwritten """ r = test_client.get( "/.well-known/dirac-metadata", diff --git a/pyproject.toml b/pyproject.toml index 99c0efb25..21ebff84a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,11 +41,11 @@ select = [ "PLE", # pylint errors "D", # pydocstyle # "UP", # pyUpgrade - "FLY", # flynt - "DTZ", # flake8-datetimez - "S", # flake8-bandit - "N", # pep8-naming - "INP", # flake8-no-implicit-namespace-packages + "FLY", # flynt + "DTZ", # flake8-datetimez + "S", # flake8-bandit + "N", # pep8-naming + "INP", # flake8-no-implicit-namespace-packages "RUF022", "TID251", # flake8-tidy-imports.banned-api ] @@ -113,8 +113,9 @@ skip = [ "diracx-testing/*", "extensions/gubbins/gubbins-client/src/gubbins/client/_generated/*", "extensions/gubbins/gubbins-*/tests/*", + "pyproject.toml", ] -ignore-words-list = ["CheckIn", "dependant"] +ignore-words-list = ["CheckIn", "dependent", "dependant"] [tool.mypy]