feat(BA-4325): Apply DataLoaders to fair share entity back-reference resolvers#8702
feat(BA-4325): Apply DataLoaders to fair share entity back-reference resolvers#8702HyeockJinKim merged 3 commits intomainfrom
Conversation
…resolvers Replace individual fetch_*() calls with core entity DataLoaders (domain_loader, project_loader, user_loader) in fair share GQL types to prevent N+1 queries. Also document cross-entity reference resolver pattern with strawberry.lazy() in api-guide skill. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: octodog <mu001@lablup.com>
…resolvers (#8702) Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: octodog <mu001@lablup.com>
There was a problem hiding this comment.
Pull request overview
This PR updates the fair share GraphQL node back-reference resolvers (domain/project/user) to use core entity DataLoaders to prevent N+1 query patterns, and documents the recommended cross-entity resolver pattern using strawberry.lazy().
Changes:
- Replace
fetch_*()calls withinfo.context.data_loaders.{domain,project,user}_loader.load(...)in fair share type resolvers. - Convert loaded data into
*V2GQLnodes viafrom_data(...)constructors. - Add an API guide section documenting cross-entity reference resolver patterns (lazy types + DataLoaders).
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated 9 comments.
| File | Description |
|---|---|
src/ai/backend/manager/api/gql/fair_share/types/user.py |
Switch fair-share → user back-reference resolver to DataLoader-based loading. |
src/ai/backend/manager/api/gql/fair_share/types/project.py |
Switch fair-share → project back-reference resolver to DataLoader-based loading. |
src/ai/backend/manager/api/gql/fair_share/types/domain.py |
Switch fair-share → domain back-reference resolver to DataLoader-based loading. |
.claude/skills/api-guide/SKILL.md |
Document cross-entity resolver pattern (strawberry.lazy() + DataLoader). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ) -> ( | ||
| Annotated[ | ||
| DomainV2GQL, | ||
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | ||
| ] | ||
| | None | ||
| ): | ||
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | ||
|
|
||
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | ||
| if domain_data is None: | ||
| return None |
There was a problem hiding this comment.
The DataLoader returns DomainData | None; calling DomainV2GQL.from_data() is only safe when data is present. Instead of returning None here, raise a not-found error (matching the previous fetch_domain() behavior) so missing domains don’t get silently swallowed.
| ) -> ( | |
| Annotated[ | |
| DomainV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | |
| ] | |
| | None | |
| ): | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| if domain_data is None: | |
| return None | |
| ) -> Annotated[ | |
| DomainV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | |
| ]: | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| if domain_data is None: | |
| raise KeyError(f"Domain not found: {self.domain_name}") |
| When a GQL node references another entity node, use `strawberry.lazy()` to avoid circular imports. Strawberry requires runtime type resolution, so `TYPE_CHECKING` imports alone are insufficient. | ||
|
|
There was a problem hiding this comment.
This section says "TYPE_CHECKING imports alone are insufficient" but earlier in the document it states "NEVER use TYPE_CHECKING for Strawberry types"; these statements conflict, and the codebase uses TYPE_CHECKING for cross-entity node types alongside strawberry.lazy(). Consider rewording to clarify that TYPE_CHECKING is fine for static typing, but runtime schema generation must rely on strawberry.lazy() / forward references (and adjust the earlier rule accordingly).
| | None | ||
| ): | ||
| from ai.backend.manager.api.gql.user_v2.types.node import UserV2GQL | ||
|
|
||
| user_data = await info.context.data_loaders.user_loader.load(self.user_uuid) | ||
| if user_data is None: | ||
| return None |
There was a problem hiding this comment.
The user field used to be non-nullable (and fetch_user() would raise UserNotFound on missing users). Making the return type ... | None changes the GraphQL schema/contract and can silently hide dangling references. Keep the field non-nullable and, if the DataLoader returns None, raise the same not-found error as the old fetcher.
| | None | |
| ): | |
| from ai.backend.manager.api.gql.user_v2.types.node import UserV2GQL | |
| user_data = await info.context.data_loaders.user_loader.load(self.user_uuid) | |
| if user_data is None: | |
| return None | |
| ): | |
| from ai.backend.manager.api.gql.user_v2.types.node import UserNotFound, UserV2GQL | |
| user_data = await info.context.data_loaders.user_loader.load(self.user_uuid) | |
| if user_data is None: | |
| raise UserNotFound(str(self.user_uuid)) |
|
|
||
| user_data = await info.context.data_loaders.user_loader.load(self.user_uuid) | ||
| if user_data is None: | ||
| return None |
There was a problem hiding this comment.
Returning None when the loader can’t find the user changes behavior compared to the previous fetch_user() call (which raises). If you keep the DataLoader approach, translate None into the appropriate not-found exception so clients get a consistent error instead of a null field.
| return None | |
| raise strawberry.exceptions.GraphQLError( | |
| f"User with UUID {self.user_uuid} not found." | |
| ) |
| ) -> ( | ||
| Annotated[ | ||
| ProjectV2GQL, | ||
| strawberry.lazy("ai.backend.manager.api.gql.project_v2.types.node"), | ||
| ] | ||
| | None | ||
| ): | ||
| from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL | ||
|
|
||
| project_data = await info.context.data_loaders.project_loader.load(self.project_id) | ||
| if project_data is None: | ||
| return None |
There was a problem hiding this comment.
The project field previously returned a non-null ProjectV2GQL (and fetch_project() raises on missing projects). Changing the return type to ... | None alters the GraphQL schema and can break clients. Prefer keeping it non-nullable and raising ProjectNotFound (or the same error raised by the old fetcher) when the DataLoader yields None.
| ) -> ( | |
| Annotated[ | |
| ProjectV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.project_v2.types.node"), | |
| ] | |
| | None | |
| ): | |
| from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL | |
| project_data = await info.context.data_loaders.project_loader.load(self.project_id) | |
| if project_data is None: | |
| return None | |
| ) -> Annotated[ | |
| ProjectV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.project_v2.types.node"), | |
| ]: | |
| from ai.backend.manager.api.gql.project_v2.types.node import ProjectV2GQL | |
| project_data = await info.context.data_loaders.project_loader.load(self.project_id) | |
| if project_data is None: | |
| raise RuntimeError( | |
| f"Project not found for fair share record (project_id={self.project_id})" | |
| ) |
|
|
||
| project_data = await info.context.data_loaders.project_loader.load(self.project_id) | ||
| if project_data is None: | ||
| return None |
There was a problem hiding this comment.
Returning None on missing project data will silently null out the back-reference instead of surfacing an error like the old fetch_project() path. Consider raising the appropriate not-found exception when project_loader.load() returns None to preserve prior semantics.
| return None | |
| raise KeyError(f"Project not found for id {self.project_id}") |
| ]: | ||
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | ||
| data = await info.context.data_loaders.domain_loader.load(self.domain_name) |
There was a problem hiding this comment.
The example resolver uses DomainV2GQL.from_data(data) but domain_loader.load() returns an optional type; as written, the snippet will fail if data is None. Update the guide to either return None (with a nullable field type) or raise DomainNotFound when the loader yields None, matching the recommended pattern.
| ]: | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| ] | None: | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| if data is None: | |
| return None |
| **DataLoaders** (`info.context.data_loaders`): Use DataLoaders instead of individual fetch functions to prevent N+1 queries. See `api/gql/data_loader/data_loaders.py` for available loaders. | ||
|
|
||
| **See examples:** | ||
| - `api/gql/fair_share/types/domain.py` - Cross-entity reference with DataLoader | ||
| - `api/gql/domain_v2/types/node.py` - DomainV2GQL with fair_shares/usage_buckets |
There was a problem hiding this comment.
The file paths in this new section are inconsistent with earlier “Key Files” paths (e.g., src/ai/backend/manager/...). api/gql/data_loader/data_loaders.py and the example paths look like they’re missing the src/ai/backend/manager/ prefix; please make the path style consistent so readers can locate the files reliably.
| ) -> ( | ||
| Annotated[ | ||
| DomainV2GQL, | ||
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | ||
| ] | ||
| | None | ||
| ): | ||
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | ||
|
|
||
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | ||
| if domain_data is None: | ||
| return None |
There was a problem hiding this comment.
The domain back-reference was previously non-null and would error if the domain didn’t exist via fetch_domain(). Making it nullable (... | None) changes the GraphQL schema/contract. Keep it non-nullable and map a None DataLoader result to DomainNotFound (or the equivalent error from the old fetcher).
| ) -> ( | |
| Annotated[ | |
| DomainV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | |
| ] | |
| | None | |
| ): | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| if domain_data is None: | |
| return None | |
| ) -> Annotated[ | |
| DomainV2GQL, | |
| strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"), | |
| ]: | |
| from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL | |
| domain_data = await info.context.data_loaders.domain_loader.load(self.domain_name) | |
| if domain_data is None: | |
| raise RuntimeError("Domain not found") |
Replace individual fetch_*() calls with core entity DataLoaders (domain_loader, project_loader, user_loader) in fair share GQL types to prevent N+1 queries. Also document cross-entity reference resolver pattern with strawberry.lazy() in api-guide skill.
Co-Authored-By: Claude Opus 4.6 noreply@anthropic.com
resolves #NNN (BA-MMM)
Checklist: (if applicable)
ai.backend.testdocsdirectory