Skip to content

Commit 33f0dd3

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 c5ec5a5 commit 33f0dd3

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,
@@ -36,7 +39,12 @@
3639
from ai.backend.manager.repositories.app_config_fragment.types import (
3740
AppConfigFragmentSearchScope,
3841
)
39-
from ai.backend.manager.repositories.base import BatchQuerier, QueryCondition, QueryOrder
42+
from ai.backend.manager.repositories.base import (
43+
BatchQuerier,
44+
OffsetPagination,
45+
QueryCondition,
46+
QueryOrder,
47+
)
4048
from ai.backend.manager.services.app_config_fragment.actions.admin_search import (
4149
AdminSearchAppConfigFragmentsAction,
4250
)
@@ -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
@@ -14,6 +14,9 @@
1414
from ai.backend.common.dto.manager.v2.rbac.response import EntityNode # pants: no-infer-dep
1515
from ai.backend.manager.api.adapters.registry import Adapters # pants: no-infer-dep
1616
from ai.backend.manager.api.gql.agent.types import AgentV2GQL # pants: no-infer-dep
17+
from ai.backend.manager.api.gql.app_config_fragment.types import ( # pants: no-infer-dep
18+
AppConfigFragmentGQL,
19+
)
1720
from ai.backend.manager.api.gql.artifact.types import ( # pants: no-infer-dep
1821
ArtifactRevision,
1922
)
@@ -113,6 +116,22 @@ class DataLoaders:
113116
def __init__(self, adapters: Adapters) -> None:
114117
self._adapters = adapters
115118

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

0 commit comments

Comments
 (0)