Skip to content

Commit 56336f7

Browse files
jopemachineclaude
andcommitted
feat(BA-5829): add AppConfigFragment DataLoader for N+1 batching
Follows the `audit_log_loader` / `prometheus_query_preset_category_loader` pattern — adapter exposes `batch_load_by_ids(ids)` which goes through the `SearchAppConfigFragments` action path with an id-filter condition, and `DataLoaders.app_config_fragment_loader` wraps it for use from GQL resolvers via `info.context.data_loaders.app_config_fragment_loader`. Cross-scope `admin_search` is used instead of the scope-bound `search` action because DataLoader batches span scopes; authorization is already enforced at the parent resolver before the loader is ever invoked (same stance as other loaders in the registry). Adds: - `AppConfigFragmentAdapter.batch_load_by_ids` (adapter-level batching via `AppConfigFragmentConditions.by_ids` + `OffsetPagination(len(ids))`). - `DataLoaders.app_config_fragment_loader` cached-property. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 26210d8 commit 56336f7

2 files changed

Lines changed: 53 additions & 1 deletion

File tree

src/ai/backend/manager/api/adapters/app_config_fragment.py

Lines changed: 34 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,9 @@
22

33
from __future__ import annotations
44

5+
import uuid
6+
from collections.abc import Sequence
7+
58
from ai.backend.common.dto.manager.v2.app_config_fragment.request import (
69
AppConfigFragmentFilter,
710
AppConfigFragmentKeyInput,
@@ -35,7 +38,12 @@
3538
from ai.backend.manager.repositories.app_config_fragment.types import (
3639
AppConfigFragmentSearchScope,
3740
)
38-
from ai.backend.manager.repositories.base import BatchQuerier, QueryCondition, QueryOrder
41+
from ai.backend.manager.repositories.base import (
42+
BatchQuerier,
43+
OffsetPagination,
44+
QueryCondition,
45+
QueryOrder,
46+
)
3947
from ai.backend.manager.services.app_config_fragment.actions.admin_search import (
4048
AdminSearchAppConfigFragmentsAction,
4149
)
@@ -60,6 +68,31 @@
6068
class AppConfigFragmentAdapter(BaseAdapter):
6169
"""Adapter for AppConfigFragment domain operations (BEP-1052 §2)."""
6270

71+
async def batch_load_by_ids(
72+
self, ids: Sequence[uuid.UUID]
73+
) -> list[AppConfigFragmentNode | None]:
74+
"""Batch-load fragments by id for DataLoader use.
75+
76+
Returns `AppConfigFragmentNode` DTOs in the same order as the
77+
input `ids`, with `None` placeholders for rows that do not
78+
exist (Relay batching contract). Goes through the cross-scope
79+
search action because the scope-bound variant requires a single
80+
`(scope_type, scope_id)` slice — DataLoader batches span
81+
scopes. Authorization for the overall query is enforced at the
82+
parent resolver before this loader is invoked.
83+
"""
84+
if not ids:
85+
return []
86+
querier = BatchQuerier(
87+
pagination=OffsetPagination(limit=len(ids)),
88+
conditions=[AppConfigFragmentConditions.by_ids(ids)],
89+
)
90+
result = await self._processors.app_config_fragment.admin_search.wait_for_complete(
91+
AdminSearchAppConfigFragmentsAction(querier=querier)
92+
)
93+
by_id = {item.id: self._data_to_dto(item) for item in result.items}
94+
return [by_id.get(fragment_id) for fragment_id in ids]
95+
6396
async def create(self, input: CreateAppConfigFragmentInput) -> CreateAppConfigFragmentPayload:
6497
key = self._input_to_key(input.key)
6598
result = await self._processors.app_config_fragment.create.wait_for_complete(

src/ai/backend/manager/api/gql/data_loader/data_loaders.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,9 @@
1313
from ai.backend.common.dto.manager.v2.rbac.response import EntityNode # pants: no-infer-dep
1414
from ai.backend.manager.api.adapters.registry import Adapters # pants: no-infer-dep
1515
from ai.backend.manager.api.gql.agent.types import AgentV2GQL # pants: no-infer-dep
16+
from ai.backend.manager.api.gql.app_config_fragment.types import ( # pants: no-infer-dep
17+
AppConfigFragmentGQL,
18+
)
1619
from ai.backend.manager.api.gql.artifact.types import ( # pants: no-infer-dep
1720
ArtifactRevision,
1821
)
@@ -112,6 +115,22 @@ class DataLoaders:
112115
def __init__(self, adapters: Adapters) -> None:
113116
self._adapters = adapters
114117

118+
@cached_property
119+
def app_config_fragment_loader(
120+
self,
121+
) -> DataLoader[uuid.UUID, AppConfigFragmentGQL | None]:
122+
adapter = self._adapters.app_config_fragment
123+
124+
async def load_fn(ids: list[uuid.UUID]) -> list[AppConfigFragmentGQL | None]:
125+
from ai.backend.manager.api.gql.app_config_fragment.types import ( # pants: no-infer-dep
126+
AppConfigFragmentGQL as F,
127+
)
128+
129+
dtos = await adapter.batch_load_by_ids(ids)
130+
return [F.from_pydantic(dto) if dto is not None else None for dto in dtos]
131+
132+
return DataLoader(load_fn=load_fn)
133+
115134
@cached_property
116135
def audit_log_loader(
117136
self,

0 commit comments

Comments
 (0)