Skip to content

Commit ee7c390

Browse files
fatih-acarfatih-acar
committed
fix(backend): fetch only required attributes for group members
Signed-off-by: Fatih Acar <[email protected]>
1 parent 87c7fe2 commit ee7c390

File tree

4 files changed

+117
-105
lines changed

4 files changed

+117
-105
lines changed

‎backend/infrahub/generators/tasks.py‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
)
2222
from infrahub.git.base import extract_repo_file_information
2323
from infrahub.git.repository import get_initialized_repo
24+
from infrahub.git.utils import fetch_definition_targets
2425
from infrahub.workers.dependencies import get_client, get_workflow
2526
from infrahub.workflows.catalogue import REQUEST_GENERATOR_DEFINITION_RUN, REQUEST_GENERATOR_RUN
2627
from infrahub.workflows.utils import add_tags
@@ -177,14 +178,9 @@ async def request_generator_definition_run(
177178
branch=model.branch,
178179
)
179180

180-
group = await client.get(
181-
kind=InfrahubKind.GENERICGROUP,
182-
prefetch_relationships=True,
183-
populate_store=True,
184-
id=model.generator_definition.group_id,
185-
branch=model.branch,
181+
group = await fetch_definition_targets(
182+
client=client, branch=model.model.branch, definition=model.generator_definition
186183
)
187-
await group.members.fetch()
188184

189185
instance_by_member = {}
190186
for instance in existing_instances:

‎backend/infrahub/git/tasks.py‎

Lines changed: 3 additions & 87 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,3 @@
1-
from collections import defaultdict
2-
from typing import Any
3-
41
from infrahub_sdk import InfrahubClient
52
from infrahub_sdk.protocols import (
63
CoreArtifact,
@@ -10,7 +7,6 @@
107
CoreRepositoryValidator,
118
CoreUserValidator,
129
)
13-
from infrahub_sdk.types import Order
1410
from infrahub_sdk.uuidt import UUIDT
1511
from prefect import flow, task
1612
from prefect.cache_policies import NONE
@@ -57,72 +53,7 @@
5753
UserCheckDefinitionData,
5854
)
5955
from .repository import InfrahubReadOnlyRepository, InfrahubRepository, get_initialized_repo
60-
61-
62-
def _collect_parameter_first_segments(params: Any) -> set[str]:
63-
segments: set[str] = set()
64-
65-
def _walk(value: Any) -> None:
66-
if isinstance(value, str):
67-
segment = value.split("__", 1)[0]
68-
if segment:
69-
segments.add(segment)
70-
elif isinstance(value, dict):
71-
for nested in value.values():
72-
_walk(nested)
73-
elif isinstance(value, (list, tuple, set)):
74-
for nested in value:
75-
_walk(nested)
76-
77-
_walk(params)
78-
return segments
79-
80-
81-
async def _prefetch_group_member_nodes(
82-
*,
83-
client: InfrahubClient,
84-
members: Any,
85-
branch: str,
86-
required_fields: set[str],
87-
) -> None:
88-
ids_per_kind: dict[str, list[str]] = defaultdict(list)
89-
for related in getattr(members, "peers", []):
90-
related_id = getattr(related, "id", None)
91-
related_type = getattr(related, "typename", None)
92-
if related_id and related_type:
93-
ids_per_kind[related_type].append(related_id)
94-
95-
if not ids_per_kind:
96-
return
97-
98-
batch = await client.create_batch()
99-
100-
for kind, ids in ids_per_kind.items():
101-
schema = await client.schema.get(kind=kind, branch=branch)
102-
103-
keep_attributes = {field for field in required_fields if field in schema.attribute_names}
104-
keep_relationships = {field for field in required_fields if field in schema.relationship_names}
105-
106-
exclude_attributes = [attr for attr in schema.attribute_names if attr not in keep_attributes]
107-
exclude_relationships = [rel for rel in schema.relationship_names if rel not in keep_relationships]
108-
109-
unique_ids = list(dict.fromkeys(ids))
110-
kwargs: dict[str, Any] = {
111-
"kind": kind,
112-
"ids": unique_ids,
113-
"branch": branch,
114-
"exclude": exclude_attributes + exclude_relationships,
115-
"populate_store": True,
116-
"order": Order(disable=True),
117-
}
118-
119-
if keep_relationships:
120-
kwargs["include"] = list(keep_relationships)
121-
122-
batch.add(task=client.filters, **kwargs)
123-
124-
async for _ in batch.execute():
125-
pass
56+
from .utils import fetch_definition_targets
12657

12758

12859
@flow(
@@ -393,20 +324,7 @@ async def generate_request_artifact_definition(
393324
kind=CoreArtifactDefinition, id=model.artifact_definition_id, branch=model.branch
394325
)
395326

396-
group = await client.get(
397-
kind=artifact_definition.targets.typename,
398-
id=artifact_definition.targets.id,
399-
branch=model.branch,
400-
include=["members"],
401-
)
402-
403-
parameter_fields = _collect_parameter_first_segments(artifact_definition.parameters.value)
404-
await _prefetch_group_member_nodes(
405-
client=client,
406-
members=group.members,
407-
branch=model.branch,
408-
required_fields=parameter_fields,
409-
)
327+
group = await fetch_definition_targets(client=client, branch=model.branch, artifact_definition=artifact_definition)
410328

411329
current_members = [member.id for member in group.members.peers]
412330

@@ -659,9 +577,7 @@ async def trigger_repository_user_checks_definitions(model: UserCheckDefinitionD
659577

660578
if definition.targets.id:
661579
# Check against a group of targets
662-
await definition.targets.fetch()
663-
group = definition.targets.peer
664-
await group.members.fetch()
580+
group = await fetch_definition_targets(client=client, branch=model.branch_name, definition=definition)
665581
check_models = []
666582
for relationship in group.members.peers:
667583
member = relationship.peer

‎backend/infrahub/git/utils.py‎

Lines changed: 106 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,15 @@
1-
from typing import TYPE_CHECKING
1+
from collections import defaultdict
2+
from typing import TYPE_CHECKING, Any
3+
4+
from infrahub_sdk import InfrahubClient
5+
from infrahub_sdk.protocols import CoreArtifactDefinition, CoreCheckDefinition, CoreGroup
6+
from infrahub_sdk.types import Order
27

38
from infrahub.core import registry
49
from infrahub.core.constants import InfrahubKind
510
from infrahub.core.manager import NodeManager
611
from infrahub.database import InfrahubDatabase
12+
from infrahub.generators.models import ProposedChangeGeneratorDefinition
713

814
from .models import RepositoryBranchInfo, RepositoryData
915

@@ -46,3 +52,102 @@ async def get_repositories_commit_per_branch(
4652
)
4753

4854
return repositories
55+
56+
57+
def _collect_parameter_first_segments(params: Any) -> set[str]:
58+
segments: set[str] = set()
59+
60+
def _walk(value: Any) -> None:
61+
if isinstance(value, str):
62+
segment = value.split("__", 1)[0]
63+
if segment:
64+
segments.add(segment)
65+
elif isinstance(value, dict):
66+
for nested in value.values():
67+
_walk(nested)
68+
elif isinstance(value, (list, tuple, set)):
69+
for nested in value:
70+
_walk(nested)
71+
72+
_walk(params)
73+
return segments
74+
75+
76+
async def _prefetch_group_member_nodes(
77+
*,
78+
client: InfrahubClient,
79+
members: Any,
80+
branch: str,
81+
required_fields: set[str],
82+
) -> None:
83+
ids_per_kind: dict[str, list[str]] = defaultdict(list)
84+
for related in getattr(members, "peers", []):
85+
related_id = getattr(related, "id", None)
86+
related_type = getattr(related, "typename", None)
87+
if related_id and related_type:
88+
ids_per_kind[related_type].append(related_id)
89+
90+
if not ids_per_kind:
91+
return
92+
93+
batch = await client.create_batch()
94+
95+
for kind, ids in ids_per_kind.items():
96+
schema = await client.schema.get(kind=kind, branch=branch)
97+
98+
keep_attributes = {field for field in required_fields if field in schema.attribute_names}
99+
keep_relationships = {field for field in required_fields if field in schema.relationship_names}
100+
101+
exclude_attributes = [attr for attr in schema.attribute_names if attr not in keep_attributes]
102+
exclude_relationships = [rel for rel in schema.relationship_names if rel not in keep_relationships]
103+
104+
unique_ids = list(dict.fromkeys(ids))
105+
kwargs: dict[str, Any] = {
106+
"kind": kind,
107+
"ids": unique_ids,
108+
"branch": branch,
109+
"exclude": exclude_attributes + exclude_relationships,
110+
"populate_store": True,
111+
"order": Order(disable=True),
112+
}
113+
114+
if keep_relationships:
115+
kwargs["include"] = list(keep_relationships)
116+
117+
batch.add(task=client.filters, **kwargs)
118+
119+
async for _ in batch.execute():
120+
pass
121+
122+
123+
async def fetch_definition_targets(
124+
client: InfrahubClient,
125+
branch: str,
126+
definition: CoreArtifactDefinition | CoreCheckDefinition | ProposedChangeGeneratorDefinition,
127+
) -> CoreGroup:
128+
group_id: str
129+
parameters: Any
130+
131+
if isinstance(definition, CoreArtifactDefinition, CoreCheckDefinition):
132+
group_id = definition.targets.id
133+
parameters = definition.parameters.value
134+
if isinstance(definition, ProposedChangeGeneratorDefinition):
135+
group_id = definition.group_id
136+
parameters = definition.parameters
137+
138+
group = await client.get(
139+
kind=CoreGroup,
140+
id=group_id,
141+
branch=branch,
142+
include=["members"],
143+
)
144+
145+
parameter_fields = _collect_parameter_first_segments(parameters)
146+
await _prefetch_group_member_nodes(
147+
client=client,
148+
members=group.members,
149+
branch=branch,
150+
required_fields=parameter_fields,
151+
)
152+
153+
return group

‎backend/infrahub/proposed_change/tasks.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828

2929
from infrahub import config, lock
3030
from infrahub.artifacts.models import CheckArtifactCreate
31+
from infrahub.artifacts.utils import get_artifact_definition_targets
3132
from infrahub.context import InfrahubContext # noqa: TC001 needed for prefect flow
3233
from infrahub.core import registry
3334
from infrahub.core.branch import Branch
@@ -58,6 +59,7 @@
5859
from infrahub.git.base import extract_repo_file_information
5960
from infrahub.git.models import TriggerRepositoryInternalChecks, TriggerRepositoryUserChecks
6061
from infrahub.git.repository import InfrahubRepository, get_initialized_repo
62+
from infrahub.git.utils import fetch_definition_targets
6163
from infrahub.log import get_logger
6264
from infrahub.message_bus.types import (
6365
ProposedChangeArtifactDefinition,
@@ -652,9 +654,7 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, c
652654
branch=model.source_branch,
653655
)
654656

655-
await artifact_definition.targets.fetch()
656-
group = artifact_definition.targets.peer
657-
await group.members.fetch()
657+
group = await get_artifact_definition_targets(client=client, artifact_definition=artifact_definition)
658658

659659
artifacts_by_member = {}
660660
for artifact in existing_artifacts:
@@ -918,14 +918,9 @@ async def request_generator_definition_check(model: RequestGeneratorDefinitionCh
918918
branch=model.source_branch,
919919
)
920920

921-
group = await client.get(
922-
kind=InfrahubKind.GENERICGROUP,
923-
prefetch_relationships=True,
924-
populate_store=True,
925-
id=model.generator_definition.group_id,
926-
branch=model.source_branch,
921+
group = await fetch_definition_targets(
922+
client=client, branch=model.source_branch, definition=model.generator_definition
927923
)
928-
await group.members.fetch()
929924

930925
instance_by_member = {}
931926
for instance in existing_instances:

0 commit comments

Comments
 (0)