Skip to content

feat(BA-4325): Apply DataLoaders to fair share entity back-reference resolvers #8702

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
HyeockJinKim merged 3 commits into from
Feb 9, 2026
Merged

feat(BA-4325): Apply DataLoaders to fair share entity back-reference resolvers #8702

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
3380aa6
feat(BA-4325): Apply DataLoaders to fair share entity back-reference …
HyeockJinKim Feb 9, 2026
89be4c1
Add new fragment
HyeockJinKim Feb 9, 2026
922f32f
chore: update api schema dump
HyeockJinKim Feb 9, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions .claude/skills/api-guide/SKILL.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,37 @@ GraphQL Resolver → check_admin_only (if admin) → Processor → Service → R
- NEVER put optional fields in Scope - use Filter instead
- Scope fields must all be required (no default values, no Optional types)

### Cross-Entity Reference Resolvers

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.

Comment on lines +244 to +245
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Copilot uses AI. Check for mistakes.
**Pattern:**
```python
# 1. TYPE_CHECKING: for static analysis (mypy)
if TYPE_CHECKING:
from ai.backend.manager.api.gql.domain_v2.types.node import DomainV2GQL

# 2. Return type: Annotated with strawberry.lazy() for runtime resolution
# 3. Function body: runtime import + DataLoader
async def domain(self, info: Info[StrawberryGQLContext]) -> 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
data = await info.context.data_loaders.domain_loader.load(self.domain_name)
Comment on lines +257 to +259
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
]:
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

Copilot uses AI. Check for mistakes.
return DomainV2GQL.from_data(data)
```

**Optional return type:** `| None` must be **outside** `Annotated[]`:
```python
) -> Annotated[DomainV2GQL, strawberry.lazy("...")] | None: # ✅
) -> Annotated[DomainV2GQL | None, strawberry.lazy("...")]: # ❌ lazy cannot resolve union
```

**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
Comment on lines +269 to 273
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copilot uses AI. Check for mistakes.
- `api/gql/fair_share/types/*.py` - Scope and Filter patterns

Expand Down
1 change: 1 addition & 0 deletions changes/8702.feature.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Apply DataLoaders to fair share entity back-reference resolvers
6 changes: 3 additions & 3 deletions docs/manager/graphql-reference/supergraph.graphql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3341,7 +3341,7 @@ type DomainFairShare implements Node
"""
Added in 26.2.0. The domain entity associated with this fair share record.
"""
domain: DomainV2!
domain: DomainV2
}

"""
Expand Down Expand Up @@ -7594,7 +7594,7 @@ type ProjectFairShare implements Node
"""
Added in 26.2.0. The project entity associated with this fair share record.
"""
project: ProjectV2!
project: ProjectV2
}

"""
Expand Down Expand Up @@ -11412,7 +11412,7 @@ type UserFairShare implements Node
"""
Added in 26.2.0. The user entity associated with this fair share record.
"""
user: UserV2!
user: UserV2
}

"""
Expand Down
6 changes: 3 additions & 3 deletions docs/manager/graphql-reference/v2-schema.graphql
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1731,7 +1731,7 @@ type DomainFairShare implements Node {
"""
Added in 26.2.0. The domain entity associated with this fair share record.
"""
domain: DomainV2!
domain: DomainV2
}

"""
Expand Down Expand Up @@ -3975,7 +3975,7 @@ type ProjectFairShare implements Node {
"""
Added in 26.2.0. The project entity associated with this fair share record.
"""
project: ProjectV2!
project: ProjectV2
}

"""
Expand Down Expand Up @@ -6526,7 +6526,7 @@ type UserFairShare implements Node {
"""
Added in 26.2.0. The user entity associated with this fair share record.
"""
user: UserV2!
user: UserV2
}

"""
Expand Down
24 changes: 15 additions & 9 deletions src/ai/backend/manager/api/gql/fair_share/types/domain.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,7 +12,7 @@
from strawberry.relay import Connection, Edge, Node, NodeID

from ai.backend.manager.api.gql.base import OrderDirection, StringFilter
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy, StrawberryGQLContext
from ai.backend.manager.data.fair_share.types import DomainFairShareData
from ai.backend.manager.repositories.base import (
QueryCondition,
Expand Down Expand Up @@ -65,14 +65,20 @@ class DomainFairShareGQL(Node):
)
async def domain(
self,
info: Info,
) -> Annotated[
DomainV2GQL,
strawberry.lazy("ai.backend.manager.api.gql.domain_v2.types.node"),
]:
from ai.backend.manager.api.gql.domain_v2.fetcher.domain import fetch_domain

return await fetch_domain(info=info, domain_name=self.domain_name)
info: Info[StrawberryGQLContext],
) -> (
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
Comment on lines +69 to +80
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
) -> (
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}")

Copilot uses AI. Check for mistakes.
Comment on lines +69 to +80
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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).

Suggested change
) -> (
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")

Copilot uses AI. Check for mistakes.
return DomainV2GQL.from_data(domain_data)

@classmethod
def from_dataclass(cls, data: DomainFairShareData) -> DomainFairShareGQL:
Expand Down
24 changes: 15 additions & 9 deletions src/ai/backend/manager/api/gql/fair_share/types/project.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from strawberry.relay import Connection, Edge, Node, NodeID

from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, UUIDFilter
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy, StrawberryGQLContext
from ai.backend.manager.data.fair_share.types import ProjectFairShareData
from ai.backend.manager.data.group.types import ProjectType
from ai.backend.manager.repositories.base import (
Expand Down Expand Up @@ -68,14 +68,20 @@ class ProjectFairShareGQL(Node):
)
async def project(
self,
info: Info,
) -> Annotated[
ProjectV2GQL,
strawberry.lazy("ai.backend.manager.api.gql.project_v2.types.node"),
]:
from ai.backend.manager.api.gql.project_v2.fetcher.project import fetch_project

return await fetch_project(info=info, project_id=self.project_id)
info: Info[StrawberryGQLContext],
) -> (
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
Comment on lines +72 to +83
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
) -> (
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})"
)

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
return None
raise KeyError(f"Project not found for id {self.project_id}")

Copilot uses AI. Check for mistakes.
return ProjectV2GQL.from_data(project_data)

@classmethod
def from_dataclass(cls, data: ProjectFairShareData) -> ProjectFairShareGQL:
Expand Down
24 changes: 15 additions & 9 deletions src/ai/backend/manager/api/gql/fair_share/types/user.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
from strawberry.relay import Connection, Edge, Node, NodeID

from ai.backend.manager.api.gql.base import OrderDirection, StringFilter, UUIDFilter
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy
from ai.backend.manager.api.gql.types import GQLFilter, GQLOrderBy, StrawberryGQLContext
from ai.backend.manager.data.fair_share.types import UserFairShareData
from ai.backend.manager.repositories.base import (
QueryCondition,
Expand Down Expand Up @@ -68,14 +68,20 @@ class UserFairShareGQL(Node):
)
async def user(
self,
info: Info,
) -> Annotated[
UserV2GQL,
strawberry.lazy("ai.backend.manager.api.gql.user_v2.types.node"),
]:
from ai.backend.manager.api.gql.user_v2.fetcher.user import fetch_user

return await fetch_user(info=info, user_uuid=self.user_uuid)
info: Info[StrawberryGQLContext],
) -> (
Annotated[
UserV2GQL,
strawberry.lazy("ai.backend.manager.api.gql.user_v2.types.node"),
]
| 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
Comment on lines +77 to +83
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
| 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))

Copilot uses AI. Check for mistakes.
Copy link

Copilot AI Feb 9, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Suggested change
return None
raise strawberry.exceptions.GraphQLError(
f"User with UUID {self.user_uuid} not found."
)

Copilot uses AI. Check for mistakes.
return UserV2GQL.from_data(user_data)

@classmethod
def from_dataclass(cls, data: UserFairShareData) -> UserFairShareGQL:
Expand Down
Loading