Add team_name to Pool APIs

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)

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:

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;

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;
 };
 
 /**

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;

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",

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):