This is an automated email from the ASF dual-hosted git repository.
pierrejeambrun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 562e7307c2c Fix N+1 query in bulk update for Variables and Pools
(#71918)
562e7307c2c is described below
commit 562e7307c2cffa1af1dc74813a215221d4c7ad9b
Author: Jyun-An Chen <[email protected]>
AuthorDate: Thu Aug 27 23:21:22 2026 +0800
Fix N+1 query in bulk update for Variables and Pools (#71918)
BulkVariableService.handle_bulk_update and
BulkPoolService.handle_bulk_update
each ran one batched existence-check query up front, then discarded the
result and re-queried every entity individually inside the loop. The
sibling handle_bulk_delete methods in the same two files, and
connections.py's handle_bulk_update, already reuse the batched-lookup
dict correctly -- this was the one spot left over from the earlier
bulk-delete N+1 fixes.
update_orm_from_pydantic in both files now takes the already-fetched
ORM object instead of a key/name string, matching the design
connections.py already uses, so the update loop no longer needs its
own per-item query.
---
.../api_fastapi/core_api/routes/public/pools.py | 10 ++++--
.../core_api/routes/public/variables.py | 11 +++++-
.../api_fastapi/core_api/services/public/pools.py | 21 +++++------
.../core_api/services/public/variables.py | 28 +++------------
.../core_api/routes/public/test_pools.py | 42 ++++++++++++++++++++++
.../core_api/routes/public/test_variables.py | 31 ++++++++++++++++
6 files changed, 104 insertions(+), 39 deletions(-)
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/pools.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/pools.py
index f2c0330c815..07f81b3a0d6 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/pools.py
+++ b/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(
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/variables.py
b/airflow-core/src/airflow/api_fastapi/core_api/routes/public/variables.py
index 705880c4898..e48031cd68d 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/routes/public/variables.py
+++ b/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
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py
b/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py
index 476c45b8f5c..1a058ec8f6a 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/pools.py
@@ -24,7 +24,6 @@ from fastapi.exceptions import RequestValidationError
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 @@ from airflow.models.pool import Pool
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 @@ class BulkPoolService(BulkService[PoolBody]):
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 @@ class BulkPoolService(BulkService[PoolBody]):
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
diff --git
a/airflow-core/src/airflow/api_fastapi/core_api/services/public/variables.py
b/airflow-core/src/airflow/api_fastapi/core_api/services/public/variables.py
index 37941fed550..b5e63fb2d84 100644
--- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/variables.py
+++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/variables.py
@@ -24,7 +24,6 @@ from fastapi.exceptions import RequestValidationError
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 @@ from airflow.models.variable import Variable
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 @@ class BulkVariableService(BulkService[VariableBody]):
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 @@ class BulkVariableService(BulkService[VariableBody]):
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)
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 d56aceecfa0..92c197e7423 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
@@ -1204,6 +1204,48 @@ class TestBulkPools(TestPoolsEndpoint):
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
diff --git
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py
index 1198f748481..ec5f915dc43 100644
---
a/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py
+++
b/airflow-core/tests/unit/api_fastapi/core_api/routes/public/test_variables.py
@@ -1605,3 +1605,34 @@ class TestBulkVariables(TestVariableEndpoint):
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"}
