Fix N+1 query in bulk update for Variables and Pools

airflow-core/src/airflow/api_fastapi/core_api/routes/public/pools.py

@@ -150,8 +150,14 @@ def patch_pool(
             "Invalid body, pool name from request body doesn't match uri parameter",
         )
 
-    pool = update_orm_from_pydantic(pool_name, patch_body, update_mask, session)
-    return pool
+    pool = session.scalar(select(Pool).where(Pool.pool == pool_name).limit(1))
+    if not pool:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND, detail=f"The Pool with name: `{pool_name}` was not found"
+        )
+
+    updated_pool = update_orm_from_pydantic(pool, patch_body, update_mask)
+    return updated_pool
 
 
 @pools_router.post(

airflow-core/src/airflow/api_fastapi/core_api/routes/public/variables.py

@@ -149,7 +149,16 @@ def patch_variable(
     update_mask: list[str] | None = Query(None),
 ) -> VariableResponse:
     """Update a variable by key."""
-    variable = update_orm_from_pydantic(variable_key, patch_body, update_mask, session)
+    if patch_body.key != variable_key:
+        raise HTTPException(
+            status.HTTP_400_BAD_REQUEST, "Invalid body, key from request body doesn't match uri parameter"
+        )
+    old_variable = session.scalar(select(Variable).filter_by(key=variable_key).limit(1))
+    if not old_variable:
+        raise HTTPException(
+            status.HTTP_404_NOT_FOUND, f"The Variable with key: `{variable_key}` was not found"
+        )
+    variable = update_orm_from_pydantic(old_variable, patch_body, update_mask)
     return variable
 
 

airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py

@@ -24,7 +24,6 @@
 from pydantic import ValidationError
 from sqlalchemy import select
 
-from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.core_api.datamodels.common import (
     BulkActionNotOnExistence,
     BulkActionOnExistence,
@@ -43,27 +42,21 @@
 
 
 def update_orm_from_pydantic(
-    pool_name: str,
+    pool: Pool,
     patch_body: PoolBody | PoolPatchBody,
     update_mask: list[str] | None,
-    session: SessionDep,
 ) -> Pool:
     """
     Update an existing pool.
 
-    :param pool_name: The name of the existing Pool to be updated.
+    :param pool: The existing Pool ORM object to update.
     :param patch_body: Pydantic model containing the fields to update.
     :param update_mask: Specific fields to update. If None, all provided fields will be considered.
-    :param session: The database session dependency.
     :return: The updated Pool instance.
     :raises HTTPException: If attempting to update disallowed fields on ``default_pool``.
     """
     # Special restriction: default pool only allows limited fields to be patched
-    pool = session.scalar(select(Pool).where(Pool.pool == pool_name).limit(1))
-    if not pool:
-        raise HTTPException(
-            status.HTTP_404_NOT_FOUND, detail=f"The Pool with name: `{pool_name}` was not found"
-        )
+    pool_name = pool.pool
     if pool_name == Pool.DEFAULT_POOL_NAME:
         if update_mask and all(mask.strip() in {"slots", "include_deferred"} for mask in update_mask):
             # Validate only slots/include_deferred
@@ -170,7 +163,9 @@ def handle_bulk_create(self, action: BulkCreateAction[PoolBody], results: BulkAc
     def handle_bulk_update(self, action: BulkUpdateAction[PoolBody], results: BulkActionResponse) -> None:
         """Bulk Update pools."""
         to_update_pool_names = {pool.pool for pool in action.entities}
-        _, matched_pool_names, not_found_pool_names = self.categorize_pools(to_update_pool_names)
+        existing_pools_dict, matched_pool_names, not_found_pool_names = self.categorize_pools(
+            to_update_pool_names
+        )
         try:
             if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_pool_names:
                 raise HTTPException(
@@ -185,7 +180,9 @@ def handle_bulk_update(self, action: BulkUpdateAction[PoolBody], results: BulkAc
                 if pool.pool not in update_pool_names:
                     continue
 
-                updated_pool = update_orm_from_pydantic(pool.pool, pool, action.update_mask, self.session)
+                updated_pool = update_orm_from_pydantic(
+                    existing_pools_dict[pool.pool], pool, action.update_mask
+                )
 
                 results.success.append(str(updated_pool.pool))  # use request field, always consistent
 

airflow-core/src/airflow/api_fastapi/core_api/services/public/variables.py

@@ -24,7 +24,6 @@
 from pydantic import ValidationError
 from sqlalchemy import select
 
-from airflow.api_fastapi.common.db.common import SessionDep
 from airflow.api_fastapi.core_api.datamodels.common import (
     BulkActionNotOnExistence,
     BulkActionOnExistence,
@@ -42,30 +41,17 @@
 
 
 def update_orm_from_pydantic(
-    variable_key: str, patch_body: VariableBody, update_mask: list[str] | None, session: SessionDep
+    old_variable: Variable, patch_body: VariableBody, update_mask: list[str] | None
 ) -> Variable:
     """
     Update an existing Variable.
 
-    :param variable_key: The name of the existing Variable_key to update.
+    :param old_variable: The existing Variable ORM object to update.
     :param patch_body: The patch request body containing fields to update.
     :param update_mask: List of fields to update. If None, all provided fields will be updated.
-    :param session: The database session dependency.
     :return: The updated Variable object.
     :raises HTTPException: If attempting to update restricted fields (e.g., ``key``).
     """
-    # Key field is immutable → cannot be patched
-
-    if patch_body.key != variable_key:
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST, "Invalid body, key from request body doesn't match uri parameter"
-        )
-    old_variable = session.scalar(select(Variable).filter_by(key=variable_key).limit(1))
-    if not old_variable:
-        raise HTTPException(
-            status.HTTP_404_NOT_FOUND, f"The Variable with key: `{variable_key}` was not found"
-        )
-
     if update_mask:
         fields_to_update = patch_body.model_fields_set & set(update_mask)
         try:
@@ -79,12 +65,6 @@ def update_orm_from_pydantic(
             raise RequestValidationError(errors=e.errors())
     non_update_fields = {"key"}
 
-    if patch_body.key != old_variable.key:
-        raise HTTPException(
-            status.HTTP_400_BAD_REQUEST,
-            "Invalid body, key from request body doesn't match uri parameter",
-        )
-
     # Apply patch via utility
     return cast(
         "Variable",
@@ -141,7 +121,7 @@ def handle_bulk_create(self, action: BulkCreateAction, results: BulkActionRespon
     def handle_bulk_update(self, action: BulkUpdateAction, results: BulkActionResponse) -> None:
         """Bulk Update variables."""
         to_update_keys = {variable.key for variable in action.entities}
-        _, matched_keys, not_found_keys = self.categorize_keys(to_update_keys)
+        existing_variables_dict, matched_keys, not_found_keys = self.categorize_keys(to_update_keys)
         try:
             if action.action_on_non_existence == BulkActionNotOnExistence.FAIL and not_found_keys:
                 raise HTTPException(
@@ -157,7 +137,7 @@ def handle_bulk_update(self, action: BulkUpdateAction, results: BulkActionRespon
                 if variable.key not in update_keys:
                     continue
                 updated_variable = update_orm_from_pydantic(
-                    variable.key, variable, action.update_mask, self.session
+                    existing_variables_dict[variable.key], variable, action.update_mask
                 )
 
                 results.success.append(updated_variable.key)

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_pools.py

@@ -1204,6 +1204,48 @@ def test_bulk_delete_query_count_is_independent_of_pool_count(self, test_client,
             f"A regression that re-queries pools inside the loop would add one SELECT per pool."
         )
 
+    @pytest.mark.parametrize(
+        ("pool_count"),
+        [5, 10, 20],
+    )
+    def test_bulk_update_query_count_is_independent_of_pool_count(self, test_client, session, pool_count):
+        # Regression guard for the N+1 fix in BulkPoolService.handle_bulk_update:
+        # the query count for a bulk update must be the same regardless of how
+        # many pools are updated. A regression that re-queries each pool inside
+        # the loop would add one SELECT per pool, so the larger run would issue
+        # strictly more queries than the smaller one.
+
+        EXPECTED_QUERY_COUNT = 4
+
+        pool_names = [f"perf_update_pool_{pool_count}_{i}" for i in range(pool_count)]
+        session.add_all(Pool(pool=name, slots=1, include_deferred=False) for name in pool_names)
+        session.commit()
+
+        request_body = {
+            "actions": [
+                {
+                    "action": "update",
+                    "entities": [
+                        {"name": name, "slots": 99, "include_deferred": False} for name in pool_names
+                    ],
+                    "action_on_non_existence": "fail",
+                }
+            ]
+        }
+
+        with count_queries() as result:
+            response = test_client.patch("/pools", json=request_body)
+
+        assert response.status_code == 200
+        assert sorted(response.json()["update"]["success"]) == sorted(pool_names)
+
+        query_count = sum(result.values())
+
+        assert query_count == EXPECTED_QUERY_COUNT, (
+            f"Bulk-update query count {query_count} does not match expected {EXPECTED_QUERY_COUNT}. "
+            f"A regression that re-queries pools inside the loop would add one SELECT per pool."
+        )
+
     def test_should_respond_401(self, unauthenticated_test_client):
         response = unauthenticated_test_client.patch("/pools", json={})
         assert response.status_code == 401

airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py

@@ -1605,3 +1605,34 @@ def test_bulk_delete_resolves_existence_in_single_query(self, session, num_varia
 
         session.commit()
         assert session.scalars(select(Variable.key).where(Variable.key.in_(keys))).all() == []
+
+    @pytest.mark.parametrize("num_variables", [1, 25])
+    def test_bulk_update_resolves_existence_in_single_query(self, session, num_variables):
+        """Bulk update looks up all targeted variables in one query, not one per key (no N+1)."""
+        keys = [f"bulk_update_var_{i}" for i in range(num_variables)]
+        for key in keys:
+            Variable.set(key=key, value="old_value", session=session)
+        session.commit()
+
+        request = BulkBody[VariableBody].model_validate(
+            {
+                "actions": [
+                    {
+                        "action": "update",
+                        "entities": [{"key": key, "value": "new_value"} for key in keys],
+                        "action_on_non_existence": "skip",
+                    }
+                ]
+            }
+        )
+        service = BulkVariableService(session=session, request=request)
+
+        with assert_queries_count(1):
+            response = service.handle_request()
+
+        assert response.update is not None
+        assert sorted(response.update.success) == sorted(keys)
+
+        session.commit()
+        updated_variables = session.scalars(select(Variable).where(Variable.key.in_(keys))).all()
+        assert {variable.val for variable in updated_variables} == {"new_value"}