diff --git a/airflow/api_connexion/endpoints/pool_endpoint.py b/airflow/api_connexion/endpoints/pool_endpoint.py index 497f31c21c22..a6ccd3a4aa9b 100644 --- a/airflow/api_connexion/endpoints/pool_endpoint.py +++ b/airflow/api_connexion/endpoints/pool_endpoint.py @@ -147,6 +147,7 @@ def patch_pool( return pool_schema.dump(pool) +@mark_fastapi_migration_done @security.requires_access_pool("POST") @action_logging @provide_session diff --git a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml index 3a3afbab95dd..76c28214100f 100644 --- a/airflow/api_fastapi/core_api/openapi/v1-generated.yaml +++ b/airflow/api_fastapi/core_api/openapi/v1-generated.yaml @@ -1191,7 +1191,7 @@ paths: content: application/json: schema: - $ref: '#/components/schemas/PoolBody' + $ref: '#/components/schemas/PoolPatchBody' responses: '200': description: Successful Response @@ -1289,6 +1289,43 @@ paths: application/json: schema: $ref: '#/components/schemas/HTTPValidationError' + post: + tags: + - Pool + summary: Post Pool + description: Create a Pool. + operationId: post_pool + requestBody: + required: true + content: + application/json: + schema: + $ref: '#/components/schemas/PoolPostBody' + responses: + '201': + description: Successful Response + content: + application/json: + schema: + $ref: '#/components/schemas/PoolResponse' + '401': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Unauthorized + '403': + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPExceptionResponse' + description: Forbidden + '422': + description: Validation Error + content: + application/json: + schema: + $ref: '#/components/schemas/HTTPValidationError' /public/providers/: get: tags: @@ -2290,7 +2327,23 @@ components: - timetables title: PluginResponse description: Plugin serializer. - PoolBody: + PoolCollectionResponse: + properties: + pools: + items: + $ref: '#/components/schemas/PoolResponse' + type: array + title: Pools + total_entries: + type: integer + title: Total Entries + type: object + required: + - pools + - total_entries + title: PoolCollectionResponse + description: Pool Collection serializer for responses. + PoolPatchBody: properties: pool: anyOf: @@ -2313,24 +2366,31 @@ components: - type: 'null' title: Include Deferred type: object - title: PoolBody - description: Pool serializer for bodies. - PoolCollectionResponse: + title: PoolPatchBody + description: Pool serializer for patch bodies. + PoolPostBody: properties: - pools: - items: - $ref: '#/components/schemas/PoolResponse' - type: array - title: Pools - total_entries: + name: + type: string + title: Name + slots: type: integer - title: Total Entries + title: Slots + description: + anyOf: + - type: string + - type: 'null' + title: Description + include_deferred: + type: boolean + title: Include Deferred + default: false type: object required: - - pools - - total_entries - title: PoolCollectionResponse - description: Pool Collection serializer for responses. + - name + - slots + title: PoolPostBody + description: Pool serializer for post bodies. PoolResponse: properties: name: diff --git a/airflow/api_fastapi/core_api/routes/public/pools.py b/airflow/api_fastapi/core_api/routes/public/pools.py index c9e30a7e2504..5690196e850a 100644 --- a/airflow/api_fastapi/core_api/routes/public/pools.py +++ b/airflow/api_fastapi/core_api/routes/public/pools.py @@ -29,8 +29,9 @@ from airflow.api_fastapi.core_api.openapi.exceptions import create_openapi_http_exception_doc from airflow.api_fastapi.core_api.serializers.pools import ( BasePool, - PoolBody, PoolCollectionResponse, + PoolPatchBody, + PoolPostBody, PoolResponse, ) from airflow.models.pool import Pool @@ -107,7 +108,7 @@ async def get_pools( @pools_router.patch("/{pool_name}", responses=create_openapi_http_exception_doc([400, 401, 403, 404])) async def patch_pool( pool_name: str, - patch_body: PoolBody, + patch_body: PoolPatchBody, session: Annotated[Session, Depends(get_session)], update_mask: list[str] | None = Query(None), ) -> PoolResponse: @@ -136,3 +137,16 @@ async def patch_pool( setattr(pool, key, value) return PoolResponse.model_validate(pool, from_attributes=True) + + +@pools_router.post("/", status_code=201, responses=create_openapi_http_exception_doc([401, 403])) +async def post_pool( + post_body: PoolPostBody, + session: Annotated[Session, Depends(get_session)], +) -> PoolResponse: + """Create a Pool.""" + pool = Pool(**post_body.model_dump()) + + session.add(pool) + + return PoolResponse.model_validate(pool, from_attributes=True) diff --git a/airflow/api_fastapi/core_api/serializers/pools.py b/airflow/api_fastapi/core_api/serializers/pools.py index dd1d6df884cc..ef3676a8afec 100644 --- a/airflow/api_fastapi/core_api/serializers/pools.py +++ b/airflow/api_fastapi/core_api/serializers/pools.py @@ -58,8 +58,8 @@ class PoolCollectionResponse(BaseModel): total_entries: int -class PoolBody(BaseModel): - """Pool serializer for bodies.""" +class PoolPatchBody(BaseModel): + """Pool serializer for patch bodies.""" model_config = ConfigDict(populate_by_name=True) @@ -67,3 +67,11 @@ class PoolBody(BaseModel): slots: int | None = None description: str | None = None include_deferred: bool | None = None + + +class PoolPostBody(BasePool): + """Pool serializer for post bodies.""" + + pool: str = Field(alias="name") + description: str | None = None + include_deferred: bool = False diff --git a/airflow/ui/openapi-gen/queries/common.ts b/airflow/ui/openapi-gen/queries/common.ts index b5e730822ffd..b12c133b6c2e 100644 --- a/airflow/ui/openapi-gen/queries/common.ts +++ b/airflow/ui/openapi-gen/queries/common.ts @@ -361,6 +361,9 @@ export const UseVersionServiceGetVersionKeyFn = (queryKey?: Array) => [ export type VariableServicePostVariableMutationResult = Awaited< ReturnType >; +export type PoolServicePostPoolMutationResult = Awaited< + ReturnType +>; export type DagServicePatchDagsMutationResult = Awaited< ReturnType >; diff --git a/airflow/ui/openapi-gen/queries/queries.ts b/airflow/ui/openapi-gen/queries/queries.ts index 31d9e94d6172..e3942ad84e08 100644 --- a/airflow/ui/openapi-gen/queries/queries.ts +++ b/airflow/ui/openapi-gen/queries/queries.ts @@ -22,7 +22,8 @@ import { import { DAGPatchBody, DagRunState, - PoolBody, + PoolPatchBody, + PoolPostBody, VariableBody, } from "../requests/types.gen"; import * as Common from "./common"; @@ -621,6 +622,43 @@ export const useVariableServicePostVariable = < }) as unknown as Promise, ...options, }); +/** + * Post Pool + * Create a Pool. + * @param data The data for the request. + * @param data.requestBody + * @returns PoolResponse Successful Response + * @throws ApiError + */ +export const usePoolServicePostPool = < + TData = Common.PoolServicePostPoolMutationResult, + TError = unknown, + TContext = unknown, +>( + options?: Omit< + UseMutationOptions< + TData, + TError, + { + requestBody: PoolPostBody; + }, + TContext + >, + "mutationFn" + >, +) => + useMutation< + TData, + TError, + { + requestBody: PoolPostBody; + }, + TContext + >({ + mutationFn: ({ requestBody }) => + PoolService.postPool({ requestBody }) as unknown as Promise, + ...options, + }); /** * Patch Dags * Patch multiple DAGs. @@ -822,7 +860,7 @@ export const usePoolServicePatchPool = < TError, { poolName: string; - requestBody: PoolBody; + requestBody: PoolPatchBody; updateMask?: string[]; }, TContext @@ -835,7 +873,7 @@ export const usePoolServicePatchPool = < TError, { poolName: string; - requestBody: PoolBody; + requestBody: PoolPatchBody; updateMask?: string[]; }, TContext diff --git a/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow/ui/openapi-gen/requests/schemas.gen.ts index 2f3eb390b485..b940646a69a7 100644 --- a/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -1420,7 +1420,27 @@ export const $PluginResponse = { description: "Plugin serializer.", } as const; -export const $PoolBody = { +export const $PoolCollectionResponse = { + properties: { + pools: { + items: { + $ref: "#/components/schemas/PoolResponse", + }, + type: "array", + title: "Pools", + }, + total_entries: { + type: "integer", + title: "Total Entries", + }, + }, + type: "object", + required: ["pools", "total_entries"], + title: "PoolCollectionResponse", + description: "Pool Collection serializer for responses.", +} as const; + +export const $PoolPatchBody = { properties: { pool: { anyOf: [ @@ -1468,28 +1488,41 @@ export const $PoolBody = { }, }, type: "object", - title: "PoolBody", - description: "Pool serializer for bodies.", + title: "PoolPatchBody", + description: "Pool serializer for patch bodies.", } as const; -export const $PoolCollectionResponse = { +export const $PoolPostBody = { properties: { - pools: { - items: { - $ref: "#/components/schemas/PoolResponse", - }, - type: "array", - title: "Pools", + name: { + type: "string", + title: "Name", }, - total_entries: { + slots: { type: "integer", - title: "Total Entries", + title: "Slots", + }, + description: { + anyOf: [ + { + type: "string", + }, + { + type: "null", + }, + ], + title: "Description", + }, + include_deferred: { + type: "boolean", + title: "Include Deferred", + default: false, }, }, type: "object", - required: ["pools", "total_entries"], - title: "PoolCollectionResponse", - description: "Pool Collection serializer for responses.", + required: ["name", "slots"], + title: "PoolPostBody", + description: "Pool serializer for post bodies.", } as const; export const $PoolResponse = { diff --git a/airflow/ui/openapi-gen/requests/services.gen.ts b/airflow/ui/openapi-gen/requests/services.gen.ts index b9d9f5265510..bfe1d2e39d36 100644 --- a/airflow/ui/openapi-gen/requests/services.gen.ts +++ b/airflow/ui/openapi-gen/requests/services.gen.ts @@ -50,6 +50,8 @@ import type { PatchPoolResponse, GetPoolsData, GetPoolsResponse, + PostPoolData, + PostPoolResponse, GetProvidersData, GetProvidersResponse, GetPluginsData, @@ -754,6 +756,30 @@ export class PoolService { }, }); } + + /** + * Post Pool + * Create a Pool. + * @param data The data for the request. + * @param data.requestBody + * @returns PoolResponse Successful Response + * @throws ApiError + */ + public static postPool( + data: PostPoolData, + ): CancelablePromise { + return __request(OpenAPI, { + method: "POST", + url: "/public/pools/", + body: data.requestBody, + mediaType: "application/json", + errors: { + 401: "Unauthorized", + 403: "Forbidden", + 422: "Validation Error", + }, + }); + } } export class ProviderService { diff --git a/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow/ui/openapi-gen/requests/types.gen.ts index 66f4437cec10..1bca12c04e77 100644 --- a/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow/ui/openapi-gen/requests/types.gen.ts @@ -324,9 +324,17 @@ export type PluginResponse = { }; /** - * Pool serializer for bodies. + * Pool Collection serializer for responses. + */ +export type PoolCollectionResponse = { + pools: Array; + total_entries: number; +}; + +/** + * Pool serializer for patch bodies. */ -export type PoolBody = { +export type PoolPatchBody = { pool?: string | null; slots?: number | null; description?: string | null; @@ -334,11 +342,13 @@ export type PoolBody = { }; /** - * Pool Collection serializer for responses. + * Pool serializer for post bodies. */ -export type PoolCollectionResponse = { - pools: Array; - total_entries: number; +export type PoolPostBody = { + name: string; + slots: number; + description?: string | null; + include_deferred?: boolean; }; /** @@ -613,7 +623,7 @@ export type GetPoolResponse = PoolResponse; export type PatchPoolData = { poolName: string; - requestBody: PoolBody; + requestBody: PoolPatchBody; updateMask?: Array | null; }; @@ -627,6 +637,12 @@ export type GetPoolsData = { export type GetPoolsResponse = PoolCollectionResponse; +export type PostPoolData = { + requestBody: PoolPostBody; +}; + +export type PostPoolResponse = PoolResponse; + export type GetProvidersData = { limit?: number; offset?: number; @@ -1248,6 +1264,27 @@ export type $OpenApiTs = { 422: HTTPValidationError; }; }; + post: { + req: PostPoolData; + res: { + /** + * Successful Response + */ + 201: PoolResponse; + /** + * Unauthorized + */ + 401: HTTPExceptionResponse; + /** + * Forbidden + */ + 403: HTTPExceptionResponse; + /** + * Validation Error + */ + 422: HTTPValidationError; + }; + }; }; "/public/providers/": { get: { diff --git a/tests/api_fastapi/core_api/routes/public/test_pools.py b/tests/api_fastapi/core_api/routes/public/test_pools.py index e2a61fedd202..da9ac0c046e8 100644 --- a/tests/api_fastapi/core_api/routes/public/test_pools.py +++ b/tests/api_fastapi/core_api/routes/public/test_pools.py @@ -281,3 +281,52 @@ def test_should_respond_200( del error["url"] assert body == expected_response + + +class TestPostPool(TestPoolsEndpoint): + @pytest.mark.parametrize( + "body, expected_status_code, expected_response", + [ + ( + {"name": "my_pool", "slots": 11}, + 201, + { + "name": "my_pool", + "slots": 11, + "description": None, + "include_deferred": False, + "occupied_slots": 0, + "running_slots": 0, + "queued_slots": 0, + "scheduled_slots": 0, + "open_slots": 11, + "deferred_slots": 0, + }, + ), + ( + {"name": "my_pool", "slots": 11, "include_deferred": True, "description": "Some description"}, + 201, + { + "name": "my_pool", + "slots": 11, + "description": "Some description", + "include_deferred": True, + "occupied_slots": 0, + "running_slots": 0, + "queued_slots": 0, + "scheduled_slots": 0, + "open_slots": 11, + "deferred_slots": 0, + }, + ), + ], + ) + def test_should_respond_200(self, test_client, session, body, expected_status_code, expected_response): + self.create_pools() + n_pools = session.query(Pool).count() + response = test_client.post("/public/pools/", json=body) + assert response.status_code == expected_status_code + + body = response.json() + assert response.json() == expected_response + assert session.query(Pool).count() == n_pools + 1