Skip to content

Commit 77a59b0

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

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_jwks_request,
6365
build_well_known_get_openid_configuration_request,
@@ -2774,6 +2776,107 @@ async def register_new_pilots_to_db(
27742776

27752777
return deserialized # type: ignore
27762778

2779+
@overload
2780+
async def patch_pilot_data(
2781+
self, body: List[str], *, content_type: str = "application/json", **kwargs: Any
2782+
) -> None:
2783+
"""Patch Pilot Data.
2784+
2785+
Endpoint to delete a pilot.
2786+
2787+
:param body: Required.
2788+
:type body: list[str]
2789+
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
2790+
Default value is "application/json".
2791+
:paramtype content_type: str
2792+
:return: None
2793+
:rtype: None
2794+
:raises ~azure.core.exceptions.HttpResponseError:
2795+
"""
2796+
2797+
@overload
2798+
async def patch_pilot_data(
2799+
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
2800+
) -> None:
2801+
"""Patch Pilot Data.
2802+
2803+
Endpoint to delete a pilot.
2804+
2805+
:param body: Required.
2806+
:type body: IO[bytes]
2807+
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
2808+
Default value is "application/json".
2809+
:paramtype content_type: str
2810+
:return: None
2811+
:rtype: None
2812+
:raises ~azure.core.exceptions.HttpResponseError:
2813+
"""
2814+
2815+
@distributed_trace_async
2816+
async def patch_pilot_data(
2817+
self, body: Union[List[str], IO[bytes]], **kwargs: Any
2818+
) -> None:
2819+
"""Patch Pilot Data.
2820+
2821+
Endpoint to delete a pilot.
2822+
2823+
:param body: Is either a [str] type or a IO[bytes] type. Required.
2824+
:type body: list[str] or IO[bytes]
2825+
:return: None
2826+
:rtype: None
2827+
:raises ~azure.core.exceptions.HttpResponseError:
2828+
"""
2829+
error_map: MutableMapping = {
2830+
401: ClientAuthenticationError,
2831+
404: ResourceNotFoundError,
2832+
409: ResourceExistsError,
2833+
304: ResourceNotModifiedError,
2834+
}
2835+
error_map.update(kwargs.pop("error_map", {}) or {})
2836+
2837+
_headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
2838+
_params = kwargs.pop("params", {}) or {}
2839+
2840+
content_type: Optional[str] = kwargs.pop(
2841+
"content_type", _headers.pop("Content-Type", None)
2842+
)
2843+
cls: ClsType[None] = kwargs.pop("cls", None)
2844+
2845+
content_type = content_type or "application/json"
2846+
_json = None
2847+
_content = None
2848+
if isinstance(body, (IOBase, bytes)):
2849+
_content = body
2850+
else:
2851+
_json = self._serialize.body(body, "[str]")
2852+
2853+
_request = build_pilots_patch_pilot_data_request(
2854+
content_type=content_type,
2855+
json=_json,
2856+
content=_content,
2857+
headers=_headers,
2858+
params=_params,
2859+
)
2860+
_request.url = self._client.format_url(_request.url)
2861+
2862+
_stream = False
2863+
pipeline_response: PipelineResponse = (
2864+
await self._client._pipeline.run( # pylint: disable=protected-access
2865+
_request, stream=_stream, **kwargs
2866+
)
2867+
)
2868+
2869+
response = pipeline_response.http_response
2870+
2871+
if response.status_code not in [204]:
2872+
map_error(
2873+
status_code=response.status_code, response=response, error_map=error_map
2874+
)
2875+
raise HttpResponseError(response=response)
2876+
2877+
if cls:
2878+
return cls(pipeline_response, None, {}) # type: ignore
2879+
27772880
@overload
27782881
async def create_pilot_secrets(
27792882
self,
@@ -2894,36 +2997,36 @@ async def associate_pilots_with_secrets(
28942997
*,
28952998
content_type: str = "application/json",
28962999
**kwargs: Any,
2897-
) -> Any:
3000+
) -> None:
28983001
"""Associate Pilots With Secrets.
28993002
2900-
Associate Pilots With Secrets.
3003+
Endpoint to associate pilots with secrets.
29013004
29023005
:param body: Required.
29033006
:type body: ~_generated.models.BodyPilotsAssociatePilotsWithSecrets
29043007
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
29053008
Default value is "application/json".
29063009
:paramtype content_type: str
2907-
:return: any
2908-
:rtype: any
3010+
:return: None
3011+
:rtype: None
29093012
:raises ~azure.core.exceptions.HttpResponseError:
29103013
"""
29113014

29123015
@overload
29133016
async def associate_pilots_with_secrets(
29143017
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
2915-
) -> Any:
3018+
) -> None:
29163019
"""Associate Pilots With Secrets.
29173020
2918-
Associate Pilots With Secrets.
3021+
Endpoint to associate pilots with secrets.
29193022
29203023
:param body: Required.
29213024
:type body: IO[bytes]
29223025
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
29233026
Default value is "application/json".
29243027
:paramtype content_type: str
2925-
:return: any
2926-
:rtype: any
3028+
:return: None
3029+
:rtype: None
29273030
:raises ~azure.core.exceptions.HttpResponseError:
29283031
"""
29293032

@@ -2932,16 +3035,16 @@ async def associate_pilots_with_secrets(
29323035
self,
29333036
body: Union[_models.BodyPilotsAssociatePilotsWithSecrets, IO[bytes]],
29343037
**kwargs: Any,
2935-
) -> Any:
3038+
) -> None:
29363039
"""Associate Pilots With Secrets.
29373040
2938-
Associate Pilots With Secrets.
3041+
Endpoint to associate pilots with secrets.
29393042
29403043
:param body: Is either a BodyPilotsAssociatePilotsWithSecrets type or a IO[bytes] type.
29413044
Required.
29423045
:type body: ~_generated.models.BodyPilotsAssociatePilotsWithSecrets or IO[bytes]
2943-
:return: any
2944-
:rtype: any
3046+
:return: None
3047+
:rtype: None
29453048
:raises ~azure.core.exceptions.HttpResponseError:
29463049
"""
29473050
error_map: MutableMapping = {
@@ -2958,7 +3061,7 @@ async def associate_pilots_with_secrets(
29583061
content_type: Optional[str] = kwargs.pop(
29593062
"content_type", _headers.pop("Content-Type", None)
29603063
)
2961-
cls: ClsType[Any] = kwargs.pop("cls", None)
3064+
cls: ClsType[None] = kwargs.pop("cls", None)
29623065

29633066
content_type = content_type or "application/json"
29643067
_json = None
@@ -2986,15 +3089,116 @@ async def associate_pilots_with_secrets(
29863089

29873090
response = pipeline_response.http_response
29883091

2989-
if response.status_code not in [200]:
3092+
if response.status_code not in [204]:
29903093
map_error(
29913094
status_code=response.status_code, response=response, error_map=error_map
29923095
)
29933096
raise HttpResponseError(response=response)
29943097

2995-
deserialized = self._deserialize("object", pipeline_response.http_response)
2996-
29973098
if cls:
2998-
return cls(pipeline_response, deserialized, {}) # type: ignore
3099+
return cls(pipeline_response, None, {}) # type: ignore
29993100

3000-
return deserialized # type: ignore
3101+
@overload
3102+
async def update_pilot_fields(
3103+
self,
3104+
body: List[_models.PilotFieldsMapping],
3105+
*,
3106+
content_type: str = "application/json",
3107+
**kwargs: Any,
3108+
) -> None:
3109+
"""Update Pilot Fields.
3110+
3111+
Update Pilot Fields.
3112+
3113+
:param body: Required.
3114+
:type body: list[~_generated.models.PilotFieldsMapping]
3115+
:keyword content_type: Body Parameter content-type. Content type parameter for JSON body.
3116+
Default value is "application/json".
3117+
:paramtype content_type: str
3118+
:return: None
3119+
:rtype: None
3120+
:raises ~azure.core.exceptions.HttpResponseError:
3121+
"""
3122+
3123+
@overload
3124+
async def update_pilot_fields(
3125+
self, body: IO[bytes], *, content_type: str = "application/json", **kwargs: Any
3126+
) -> None:
3127+
"""Update Pilot Fields.
3128+
3129+
Update Pilot Fields.
3130+
3131+
:param body: Required.
3132+
:type body: IO[bytes]
3133+
:keyword content_type: Body Parameter content-type. Content type parameter for binary body.
3134+
Default value is "application/json".
3135+
:paramtype content_type: str
3136+
:return: None
3137+
:rtype: None
3138+
:raises ~azure.core.exceptions.HttpResponseError:
3139+
"""
3140+
3141+
@distributed_trace_async
3142+
async def update_pilot_fields(
3143+
self, body: Union[List[_models.PilotFieldsMapping], IO[bytes]], **kwargs: Any
3144+
) -> None:
3145+
"""Update Pilot Fields.
3146+
3147+
Update Pilot Fields.
3148+
3149+
:param body: Is either a [PilotFieldsMapping] type or a IO[bytes] type. Required.
3150+
:type body: list[~_generated.models.PilotFieldsMapping] or IO[bytes]
3151+
:return: None
3152+
:rtype: None
3153+
:raises ~azure.core.exceptions.HttpResponseError:
3154+
"""
3155+
error_map: MutableMapping = {
3156+
401: ClientAuthenticationError,
3157+
404: ResourceNotFoundError,
3158+
409: ResourceExistsError,
3159+
304: ResourceNotModifiedError,
3160+
}
3161+
error_map.update(kwargs.pop("error_map", {}) or {})
3162+
3163+
_headers = case_insensitive_dict(kwargs.pop("headers", {}) or {})
3164+
_params = kwargs.pop("params", {}) or {}
3165+
3166+
content_type: Optional[str] = kwargs.pop(
3167+
"content_type", _headers.pop("Content-Type", None)
3168+
)
3169+
cls: ClsType[None] = kwargs.pop("cls", None)
3170+
3171+
content_type = content_type or "application/json"
3172+
_json = None
3173+
_content = None
3174+
if isinstance(body, (IOBase, bytes)):
3175+
_content = body
3176+
else:
3177+
_json = self._serialize.body(body, "[PilotFieldsMapping]")
3178+
3179+
_request = build_pilots_update_pilot_fields_request(
3180+
content_type=content_type,
3181+
json=_json,
3182+
content=_content,
3183+
headers=_headers,
3184+
params=_params,
3185+
)
3186+
_request.url = self._client.format_url(_request.url)
3187+
3188+
_stream = False
3189+
pipeline_response: PipelineResponse = (
3190+
await self._client._pipeline.run( # pylint: disable=protected-access
3191+
_request, stream=_stream, **kwargs
3192+
)
3193+
)
3194+
3195+
response = pipeline_response.http_response
3196+
3197+
if response.status_code not in [204]:
3198+
map_error(
3199+
status_code=response.status_code, response=response, error_map=error_map
3200+
)
3201+
raise HttpResponseError(response=response)
3202+
3203+
if cls:
3204+
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)