Skip to content

Commit cf2dd5b

Browse files
Robin-Van-de-MerghelRobin-Van-de-Merghel
committed
feat: We can modify pilot fields (statuses, benchmarks, ...)
1 parent be789f9 commit cf2dd5b

File tree

15 files changed

+896
-87
lines changed

15 files changed

+896
-87
lines changed

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

Lines changed: 222 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,9 @@
5757
build_jobs_unassign_job_sandboxes_request,
5858
build_pilots_associate_pilots_with_secrets_request,
5959
build_pilots_create_pilot_secrets_request,
60+
build_pilots_patch_pilot_data_request,
6061
build_pilots_register_new_pilots_to_db_request,
62+
build_pilots_update_pilot_fields_request,
6163
build_well_known_get_installation_metadata_request,
6264
build_well_known_get_openid_configuration_request,
6365
)
@@ -2722,6 +2724,107 @@ async def register_new_pilots_to_db(
27222724

27232725
return deserialized # type: ignore
27242726

2727+
@overload
2728+
async def patch_pilot_data(
2729+
self, body: List[str], *, content_type: str = "application/json", **kwargs: Any
2730+
) -> None:
2731+
"""Patch Pilot Data.
2732+
2733+
Endpoint to delete a pilot.
2734+
2735+
:param body: Required.
2736+
:type body: list[str]
2737+
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
2738+
Default value is "application/json".
2739+
:paramtype content_type: str
2740+
:return: None
2741+
:rtype: None
2742+
:raises ~azure.core.exceptions.HttpResponseError:
2743+
"""
2744+
2745+
@overload
2746+
async def patch_pilot_data(
2747+
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
2748+
) -> None:
2749+
"""Patch Pilot Data.
2750+
2751+
Endpoint to delete a pilot.
2752+
2753+
:param body: Required.
2754+
:type body: IO[bytes]
2755+
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
2756+
Default value is "application/json".
2757+
:paramtype content_type: str
2758+
:return: None
2759+
:rtype: None
2760+
:raises ~azure.core.exceptions.HttpResponseError:
2761+
"""
2762+
2763+
@distributed_trace_async
2764+
async def patch_pilot_data(
2765+
self, body: Union[List[str], IO[bytes]], **kwargs: Any
2766+
) -> None:
2767+
"""Patch Pilot Data.
2768+
2769+
Endpoint to delete a pilot.
2770+
2771+
:param body: Is either a [str] type or a IO[bytes] type. Required.
2772+
:type body: list[str] or IO[bytes]
2773+
:return: None
2774+
:rtype: None
2775+
:raises ~azure.core.exceptions.HttpResponseError:
2776+
"""
2777+
error_map: MutableMapping = {
2778+
401: ClientAuthenticationError,
2779+
404: ResourceNotFoundError,
2780+
409: ResourceExistsError,
2781+
304: ResourceNotModifiedError,
2782+
}
2783+
error_map.update(kwargs.pop("error_map", {}) or {})
2784+
2785+
_headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
2786+
_params = kwargs.pop("params", {}) or {}
2787+
2788+
content_type: Optional[str] = kwargs.pop(
2789+
"content_type", _headers.pop("Content-Type", None)
2790+
)
2791+
cls: ClsType[None] = kwargs.pop("cls", None)
2792+
2793+
content_type = content_type or "application/json"
2794+
_json = None
2795+
_content = None
2796+
if isinstance(body, (IOBase, bytes)):
2797+
_content = body
2798+
else:
2799+
_json = self._serialize.body(body, "[str]")
2800+
2801+
_request = build_pilots_patch_pilot_data_request(
2802+
content_type=content_type,
2803+
json=_json,
2804+
content=_content,
2805+
headers=_headers,
2806+
params=_params,
2807+
)
2808+
_request.url = self._client.format_url(_request.url)
2809+
2810+
_stream = False
2811+
pipeline_response: PipelineResponse = (
2812+
await self._client._pipeline.run( # pylint: disable=protected-access
2813+
_request, stream=_stream, **kwargs
2814+
)
2815+
)
2816+
2817+
response = pipeline_response.http_response
2818+
2819+
if response.status_code not in [204]:
2820+
map_error(
2821+
status_code=response.status_code, response=response, error_map=error_map
2822+
)
2823+
raise HttpResponseError(response=response)
2824+
2825+
if cls:
2826+
return cls(pipeline_response, None, {}) # type: ignore
2827+
27252828
@overload
27262829
async def create_pilot_secrets(
27272830
self,
@@ -2842,36 +2945,36 @@ async def associate_pilots_with_secrets(
28422945
*,
28432946
content_type: str = "application/json",
28442947
**kwargs: Any,
2845-
) -> Any:
2948+
) -> None:
28462949
"""Associate Pilots With Secrets.
28472950
2848-
Associate Pilots With Secrets.
2951+
Endpoint to associate pilots with secrets.
28492952
28502953
:param body: Required.
28512954
:type body: ~_generated.models.BodyPilotsAssociatePilotsWithSecrets
28522955
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
28532956
Default value is "application/json".
28542957
:paramtype content_type: str
2855-
:return: any
2856-
:rtype: any
2958+
:return: None
2959+
:rtype: None
28572960
:raises ~azure.core.exceptions.HttpResponseError:
28582961
"""
28592962

28602963
@overload
28612964
async def associate_pilots_with_secrets(
28622965
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
2863-
) -> Any:
2966+
) -> None:
28642967
"""Associate Pilots With Secrets.
28652968
2866-
Associate Pilots With Secrets.
2969+
Endpoint to associate pilots with secrets.
28672970
28682971
:param body: Required.
28692972
:type body: IO[bytes]
28702973
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
28712974
Default value is "application/json".
28722975
:paramtype content_type: str
2873-
:return: any
2874-
:rtype: any
2976+
:return: None
2977+
:rtype: None
28752978
:raises ~azure.core.exceptions.HttpResponseError:
28762979
"""
28772980

@@ -2880,16 +2983,16 @@ async def associate_pilots_with_secrets(
28802983
self,
28812984
body: Union[_models.BodyPilotsAssociatePilotsWithSecrets, IO[bytes]],
28822985
**kwargs: Any,
2883-
) -> Any:
2986+
) -> None:
28842987
"""Associate Pilots With Secrets.
28852988
2886-
Associate Pilots With Secrets.
2989+
Endpoint to associate pilots with secrets.
28872990
28882991
:param body: Is either a BodyPilotsAssociatePilotsWithSecrets type or a IO[bytes] type.
28892992
Required.
28902993
:type body: ~_generated.models.BodyPilotsAssociatePilotsWithSecrets or IO[bytes]
2891-
:return: any
2892-
:rtype: any
2994+
:return: None
2995+
:rtype: None
28932996
:raises ~azure.core.exceptions.HttpResponseError:
28942997
"""
28952998
error_map: MutableMapping = {
@@ -2906,7 +3009,7 @@ async def associate_pilots_with_secrets(
29063009
content_type: Optional[str] = kwargs.pop(
29073010
"content_type", _headers.pop("Content-Type", None)
29083011
)
2909-
cls: ClsType[Any] = kwargs.pop("cls", None)
3012+
cls: ClsType[None] = kwargs.pop("cls", None)
29103013

29113014
content_type = content_type or "application/json"
29123015
_json = None
@@ -2934,15 +3037,116 @@ async def associate_pilots_with_secrets(
29343037

29353038
response = pipeline_response.http_response
29363039

2937-
if response.status_code not in [200]:
3040+
if response.status_code not in [204]:
29383041
map_error(
29393042
status_code=response.status_code, response=response, error_map=error_map
29403043
)
29413044
raise HttpResponseError(response=response)
29423045

2943-
deserialized = self._deserialize("object", pipeline_response.http_response)
2944-
29453046
if cls:
2946-
return cls(pipeline_response, deserialized, {}) # type: ignore
3047+
return cls(pipeline_response, None, {}) # type: ignore
29473048

2948-
return deserialized # type: ignore
3049+
@overload
3050+
async def update_pilot_fields(
3051+
self,
3052+
body: List[_models.PilotFieldsMapping],
3053+
*,
3054+
content_type: str = "application/json",
3055+
**kwargs: Any,
3056+
) -> None:
3057+
"""Update Pilot Fields.
3058+
3059+
Update Pilot Fields.
3060+
3061+
:param body: Required.
3062+
:type body: list[~_generated.models.PilotFieldsMapping]
3063+
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
3064+
Default value is "application/json".
3065+
:paramtype content_type: str
3066+
:return: None
3067+
:rtype: None
3068+
:raises ~azure.core.exceptions.HttpResponseError:
3069+
"""
3070+
3071+
@overload
3072+
async def update_pilot_fields(
3073+
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
3074+
) -> None:
3075+
"""Update Pilot Fields.
3076+
3077+
Update Pilot Fields.
3078+
3079+
:param body: Required.
3080+
:type body: IO[bytes]
3081+
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
3082+
Default value is "application/json".
3083+
:paramtype content_type: str
3084+
:return: None
3085+
:rtype: None
3086+
:raises ~azure.core.exceptions.HttpResponseError:
3087+
"""
3088+
3089+
@distributed_trace_async
3090+
async def update_pilot_fields(
3091+
self, body: Union[List[_models.PilotFieldsMapping], IO[bytes]], **kwargs: Any
3092+
) -> None:
3093+
"""Update Pilot Fields.
3094+
3095+
Update Pilot Fields.
3096+
3097+
:param body: Is either a [PilotFieldsMapping] type or a IO[bytes] type. Required.
3098+
:type body: list[~_generated.models.PilotFieldsMapping] or IO[bytes]
3099+
:return: None
3100+
:rtype: None
3101+
:raises ~azure.core.exceptions.HttpResponseError:
3102+
"""
3103+
error_map: MutableMapping = {
3104+
401: ClientAuthenticationError,
3105+
404: ResourceNotFoundError,
3106+
409: ResourceExistsError,
3107+
304: ResourceNotModifiedError,
3108+
}
3109+
error_map.update(kwargs.pop("error_map", {}) or {})
3110+
3111+
_headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
3112+
_params = kwargs.pop("params", {}) or {}
3113+
3114+
content_type: Optional[str] = kwargs.pop(
3115+
"content_type", _headers.pop("Content-Type", None)
3116+
)
3117+
cls: ClsType[None] = kwargs.pop("cls", None)
3118+
3119+
content_type = content_type or "application/json"
3120+
_json = None
3121+
_content = None
3122+
if isinstance(body, (IOBase, bytes)):
3123+
_content = body
3124+
else:
3125+
_json = self._serialize.body(body, "[PilotFieldsMapping]")
3126+
3127+
_request = build_pilots_update_pilot_fields_request(
3128+
content_type=content_type,
3129+
json=_json,
3130+
content=_content,
3131+
headers=_headers,
3132+
params=_params,
3133+
)
3134+
_request.url = self._client.format_url(_request.url)
3135+
3136+
_stream = False
3137+
pipeline_response: PipelineResponse = (
3138+
await self._client._pipeline.run( # pylint: disable=protected-access
3139+
_request, stream=_stream, **kwargs
3140+
)
3141+
)
3142+
3143+
response = pipeline_response.http_response
3144+
3145+
if response.status_code not in [204]:
3146+
map_error(
3147+
status_code=response.status_code, response=response, error_map=error_map
3148+
)
3149+
raise HttpResponseError(response=response)
3150+
3151+
if cls:
3152+
return cls(pipeline_response, None, {}) # type: ignore

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
Metadata,
3434
OpenIDConfiguration,
3535
PilotCredentialsInfo,
36+
PilotFieldsMapping,
3637
PilotSecretsInfo,
3738
PilotStampInfo,
3839
ResponsePilotsRegisterNewPilotsToDb,
@@ -57,6 +58,7 @@
5758
from ._enums import ( # type: ignore
5859
ChecksumAlgorithm,
5960
JobStatus,
61+
PilotStatus,
6062
SandboxFormat,
6163
SandboxType,
6264
ScalarSearchOperator,
@@ -89,6 +91,7 @@
8991
"Metadata",
9092
"OpenIDConfiguration",
9193
"PilotCredentialsInfo",
94+
"PilotFieldsMapping",
9295
"PilotSecretsInfo",
9396
"PilotStampInfo",
9497
"ResponsePilotsRegisterNewPilotsToDb",
@@ -110,6 +113,7 @@
110113
"VectorSearchSpecValues",
111114
"ChecksumAlgorithm",
112115
"JobStatus",
116+
"PilotStatus",
113117
"SandboxFormat",
114118
"SandboxType",
115119
"ScalarSearchOperator",

‎diracx-client/src/diracx/client/_generated/models/_enums.py‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,19 @@ class JobStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta):
3434
RESCHEDULED = "Rescheduled"
3535

3636

37+
class PilotStatus(str, Enum, metaclass=CaseInsensitiveEnumMeta):
38+
"""Pilot statuses from Dirac."""
39+
40+
SUBMITTED = "Submitted"
41+
WAITING = "Waiting"
42+
RUNNING = "Running"
43+
DONE = "Done"
44+
FAILED = "Failed"
45+
DELETED = "Deleted"
46+
ABORTED = "Aborted"
47+
UNKNOWN = "Unknown"
48+
49+
3750
class SandboxFormat(str, Enum, metaclass=CaseInsensitiveEnumMeta):
3851
"""SandboxFormat."""
3952

0 commit comments

Comments
 (0)