diff --git a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py index fc5344f7a7a85..0e56684081a73 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/datamodels/pools.py @@ -52,6 +52,7 @@ class PoolResponse(BasePool): scheduled_slots: Annotated[int, BeforeValidator(_call_function)] open_slots: Annotated[int, BeforeValidator(_call_function)] deferred_slots: Annotated[int, BeforeValidator(_call_function)] + team_name: str | None class PoolCollectionResponse(BaseModel): @@ -68,6 +69,7 @@ class PoolPatchBody(StrictBaseModel): slots: int | None = None description: str | None = None include_deferred: bool | None = None + team_name: str | None = Field(max_length=50, default=None) class PoolBody(BasePool, StrictBaseModel): @@ -76,3 +78,4 @@ class PoolBody(BasePool, StrictBaseModel): pool: str = Field(alias="name", max_length=256) description: str | None = None include_deferred: bool = False + team_name: str | None = Field(max_length=50, default=None) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml index 1715c2444a68d..9993e05efc688 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml +++ b/airflow-core/src/airflow/api_fastapi/core_api/openapi/v2-rest-api-generated.yaml @@ -12051,6 +12051,12 @@ components: type: boolean title: Include Deferred default: false + team_name: + anyOf: + - type: string + maxLength: 50 + - type: 'null' + title: Team Name additionalProperties: false type: object required: @@ -12096,6 +12102,12 @@ components: - type: boolean - type: 'null' title: Include Deferred + team_name: + anyOf: + - type: string + maxLength: 50 + - type: 'null' + title: Team Name additionalProperties: false type: object title: PoolPatchBody @@ -12134,6 +12146,11 @@ components: deferred_slots: type: integer title: Deferred Slots + team_name: + anyOf: + - type: string + - type: 'null' + title: Team Name type: object required: - name @@ -12145,6 +12162,7 @@ components: - scheduled_slots - open_slots - deferred_slots + - team_name title: PoolResponse description: Pool serializer for responses. ProviderCollectionResponse: diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts index de34ac2f9be2a..660742644e619 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/schemas.gen.ts @@ -4614,6 +4614,18 @@ export const $PoolBody = { type: 'boolean', title: 'Include Deferred', default: false + }, + team_name: { + anyOf: [ + { + type: 'string', + maxLength: 50 + }, + { + type: 'null' + } + ], + title: 'Team Name' } }, additionalProperties: false, @@ -4688,6 +4700,18 @@ export const $PoolPatchBody = { } ], title: 'Include Deferred' + }, + team_name: { + anyOf: [ + { + type: 'string', + maxLength: 50 + }, + { + type: 'null' + } + ], + title: 'Team Name' } }, additionalProperties: false, @@ -4744,10 +4768,21 @@ export const $PoolResponse = { deferred_slots: { type: 'integer', title: 'Deferred Slots' + }, + team_name: { + anyOf: [ + { + type: 'string' + }, + { + type: 'null' + } + ], + title: 'Team Name' } }, type: 'object', - required: ['name', 'slots', 'include_deferred', 'occupied_slots', 'running_slots', 'queued_slots', 'scheduled_slots', 'open_slots', 'deferred_slots'], + required: ['name', 'slots', 'include_deferred', 'occupied_slots', 'running_slots', 'queued_slots', 'scheduled_slots', 'open_slots', 'deferred_slots', 'team_name'], title: 'PoolResponse', description: 'Pool serializer for responses.' } as const; diff --git a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts index ef22addbc4f45..99da929a0508c 100644 --- a/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts +++ b/airflow-core/src/airflow/ui/openapi-gen/requests/types.gen.ts @@ -1196,6 +1196,7 @@ export type PoolBody = { slots: number; description?: string | null; include_deferred?: boolean; + team_name?: string | null; }; /** @@ -1214,6 +1215,7 @@ export type PoolPatchBody = { slots?: number | null; description?: string | null; include_deferred?: boolean | null; + team_name?: string | null; }; /** @@ -1230,6 +1232,7 @@ export type PoolResponse = { scheduled_slots: number; open_slots: number; deferred_slots: number; + team_name: string | null; }; /** diff --git a/airflow-core/src/airflow/ui/src/utils/slots.tsx b/airflow-core/src/airflow/ui/src/utils/slots.tsx index 66a055ebc34b9..8211e990c7cef 100644 --- a/airflow-core/src/airflow/ui/src/utils/slots.tsx +++ b/airflow-core/src/airflow/ui/src/utils/slots.tsx @@ -23,7 +23,7 @@ import { StateIcon } from "src/components/StateIcon"; export type Slots = Omit< PoolResponse, - "description" | "include_deferred" | "name" | "occupied_slots" | "slots" + "description" | "include_deferred" | "name" | "occupied_slots" | "slots" | "team_name" >; export type SlotConfig = { color: string; diff --git a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py index b74bd6bfab94a..6d8d924ae2b78 100644 --- a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py +++ b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py @@ -22,9 +22,10 @@ from sqlalchemy import func, select from airflow.models.pool import Pool +from airflow.models.team import Team from airflow.utils.session import provide_session -from tests_common.test_utils.db import clear_db_pools +from tests_common.test_utils.db import clear_db_pools, clear_db_teams from tests_common.test_utils.logs import check_last_log pytestmark = pytest.mark.db_test @@ -48,7 +49,7 @@ @provide_session def _create_pools(session) -> None: - pool1 = Pool(pool=POOL1_NAME, slots=POOL1_SLOT, include_deferred=POOL1_INCLUDE_DEFERRED) + pool1 = Pool(pool=POOL1_NAME, slots=POOL1_SLOT, include_deferred=POOL1_INCLUDE_DEFERRED, team_name="test") pool2 = Pool(pool=POOL2_NAME, slots=POOL2_SLOT, include_deferred=POOL2_INCLUDE_DEFERRED) pool3 = Pool( pool=POOL3_NAME, @@ -59,15 +60,23 @@ def _create_pools(session) -> None: session.add_all([pool1, pool2, pool3]) +@provide_session +def _create_team(session) -> None: + session.add(Team(name="test")) + session.commit() + + class TestPoolsEndpoint: @pytest.fixture(autouse=True) def setup(self) -> None: clear_db_pools() + clear_db_teams() def teardown_method(self) -> None: clear_db_pools() def create_pools(self): + _create_team() _create_pools() @@ -130,6 +139,7 @@ def test_get_should_respond_200(self, test_client, session): "running_slots": 0, "scheduled_slots": 0, "slots": 3, + "team_name": "test", } def test_get_should_respond_401(self, unauthenticated_test_client): @@ -162,6 +172,7 @@ def test_get_pool3_should_respond_200(self, test_client, session): "running_slots": 0, "scheduled_slots": 0, "slots": 5, + "team_name": None, } @@ -291,6 +302,7 @@ class TestPatchPool(TestPoolsEndpoint): "running_slots": 0, "scheduled_slots": 0, "slots": 150, + "team_name": None, }, ), # Partial body on default_pool alternate @@ -310,6 +322,7 @@ class TestPatchPool(TestPoolsEndpoint): "running_slots": 0, "scheduled_slots": 0, "slots": 150, + "team_name": None, }, ), # Full body @@ -334,6 +347,7 @@ class TestPatchPool(TestPoolsEndpoint): "running_slots": 0, "scheduled_slots": 0, "slots": 8, + "team_name": "test", }, ), ], @@ -386,6 +400,7 @@ def test_patch_pool3_should_respond_200(self, test_client, session): "running_slots": 0, "scheduled_slots": 0, "slots": 10, + "team_name": None, } assert response.json() == expected_response check_last_log(session, dag_id=None, event="patch_pool", logical_date=None) @@ -409,6 +424,7 @@ class TestPostPool(TestPoolsEndpoint): "scheduled_slots": 0, "open_slots": 11, "deferred_slots": 0, + "team_name": None, }, ), ( @@ -425,6 +441,30 @@ class TestPostPool(TestPoolsEndpoint): "scheduled_slots": 0, "open_slots": 11, "deferred_slots": 0, + "team_name": None, + }, + ), + ( + { + "name": "my_pool", + "slots": 11, + "include_deferred": True, + "description": "Some description", + "team_name": "test", + }, + 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, + "team_name": "test", }, ), ], @@ -470,6 +510,7 @@ def test_should_respond_403(self, unauthorized_test_client): "scheduled_slots": 0, "open_slots": 11, "deferred_slots": 0, + "team_name": None, }, 409, None, @@ -547,7 +588,12 @@ class TestBulkPools(TestPoolsEndpoint): { "action": "create", "entities": [ - {"name": "pool3", "slots": 10, "description": "New Description"}, + { + "name": "pool3", + "slots": 10, + "description": "New Description", + "team_name": "test", + }, {"name": "pool2", "slots": 20, "description": "New Description"}, ], "action_on_existence": "overwrite", diff --git a/airflow-ctl/src/airflowctl/api/datamodels/generated.py b/airflow-ctl/src/airflowctl/api/datamodels/generated.py index 9b37a181d73cc..ed5c0e4784383 100644 --- a/airflow-ctl/src/airflowctl/api/datamodels/generated.py +++ b/airflow-ctl/src/airflowctl/api/datamodels/generated.py @@ -636,6 +636,7 @@ class PoolBody(BaseModel): slots: Annotated[int, Field(title="Slots")] description: Annotated[str | None, Field(title="Description")] = None include_deferred: Annotated[bool | None, Field(title="Include Deferred")] = False + team_name: Annotated[TeamName | None, Field(title="Team Name")] = None class PoolPatchBody(BaseModel): @@ -650,6 +651,7 @@ class PoolPatchBody(BaseModel): slots: Annotated[int | None, Field(title="Slots")] = None description: Annotated[str | None, Field(title="Description")] = None include_deferred: Annotated[bool | None, Field(title="Include Deferred")] = None + team_name: Annotated[TeamName | None, Field(title="Team Name")] = None class PoolResponse(BaseModel): @@ -667,6 +669,7 @@ class PoolResponse(BaseModel): scheduled_slots: Annotated[int, Field(title="Scheduled Slots")] open_slots: Annotated[int, Field(title="Open Slots")] deferred_slots: Annotated[int, Field(title="Deferred Slots")] + team_name: Annotated[str | None, Field(title="Team Name")] = None class ProviderResponse(BaseModel):