Skip to content

Commit 6465b92

Browse files
authored
fix: improve cardinality many relationship fetch (#476)
* fix: improve cardinality many relationship fetch Signed-off-by: Fatih Acar <[email protected]>
1 parent 8d89883 commit 6465b92

File tree

2 files changed

+61
-3
lines changed

2 files changed

+61
-3
lines changed

infrahub_sdk/node/relationship.py

Lines changed: 47 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,14 @@
11
from __future__ import annotations
22

3+
from collections import defaultdict
34
from collections.abc import Iterable
45
from typing import TYPE_CHECKING, Any
56

67
from ..exceptions import (
8+
Error,
79
UninitializedError,
810
)
11+
from ..types import Order
912
from .constants import PROPERTIES_FLAG, PROPERTIES_OBJECT
1013
from .related_node import RelatedNode, RelatedNodeSync
1114

@@ -156,8 +159,29 @@ async def fetch(self) -> None:
156159
self.peers = rm.peers
157160
self.initialized = True
158161

162+
ids_per_kind_map = defaultdict(list)
159163
for peer in self.peers:
160-
await peer.fetch() # type: ignore[misc]
164+
if not peer.id or not peer.typename:
165+
raise Error("Unable to fetch the peer, id and/or typename are not defined")
166+
if peer.typename not in ids_per_kind_map:
167+
ids_per_kind_map[peer.typename] = [peer.id]
168+
else:
169+
ids_per_kind_map[peer.typename].append(peer.id)
170+
171+
batch = await self.client.create_batch()
172+
for kind, ids in ids_per_kind_map.items():
173+
batch.add(
174+
task=self.client.filters,
175+
kind=kind,
176+
ids=ids,
177+
populate_store=True,
178+
branch=self.branch,
179+
parallel=True,
180+
order=Order(disable=True),
181+
)
182+
183+
async for _ in batch.execute():
184+
pass
161185

162186
def add(self, data: str | RelatedNode | dict) -> None:
163187
"""Add a new peer to this relationship."""
@@ -261,8 +285,29 @@ def fetch(self) -> None:
261285
self.peers = rm.peers
262286
self.initialized = True
263287

288+
ids_per_kind_map = defaultdict(list)
264289
for peer in self.peers:
265-
peer.fetch()
290+
if not peer.id or not peer.typename:
291+
raise Error("Unable to fetch the peer, id and/or typename are not defined")
292+
if peer.typename not in ids_per_kind_map:
293+
ids_per_kind_map[peer.typename] = [peer.id]
294+
else:
295+
ids_per_kind_map[peer.typename].append(peer.id)
296+
297+
batch = self.client.create_batch()
298+
for kind, ids in ids_per_kind_map.items():
299+
batch.add(
300+
task=self.client.filters,
301+
kind=kind,
302+
ids=ids,
303+
populate_store=True,
304+
branch=self.branch,
305+
parallel=True,
306+
order=Order(disable=True),
307+
)
308+
309+
for _ in batch.execute():
310+
pass
266311

267312
def add(self, data: str | RelatedNodeSync | dict) -> None:
268313
"""Add a new peer to this relationship."""

tests/unit/sdk/test_node.py

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1883,6 +1883,19 @@ async def test_node_fetch_relationship(
18831883
)
18841884

18851885
response2 = {
1886+
"data": {
1887+
"BuiltinTag": {
1888+
"count": 1,
1889+
}
1890+
}
1891+
}
1892+
1893+
httpx_mock.add_response(
1894+
method="POST",
1895+
json=response2,
1896+
)
1897+
1898+
response3 = {
18861899
"data": {
18871900
"BuiltinTag": {
18881901
"count": 1,
@@ -1895,7 +1908,7 @@ async def test_node_fetch_relationship(
18951908

18961909
httpx_mock.add_response(
18971910
method="POST",
1898-
json=response2,
1911+
json=response3,
18991912
match_headers={"X-Infrahub-Tracker": "query-builtintag-page1"},
19001913
)
19011914

0 commit comments

Comments
 (0)