infrahub-server 1.1.8__py3-none-any.whl → 1.1.9__py3-none-any.whl

infrahub/config.py

@@ -249,6 +249,12 @@ class DatabaseSettings(BaseSettings):
     retry_limit: int = Field(
         default=3, description="Maximum number of times a transient issue in a transaction should be retried."
     )
+    max_concurrent_queries: int = Field(
+        default=0, ge=0, description="Maximum number of concurrent queries that can run (0 means unlimited)."
+    )
+    max_concurrent_queries_delay: float = Field(
+        default=0.01, ge=0, description="Delay to add when max_concurrent_queries is reached."
+    )
 
     @property
     def database_name(self) -> str:

infrahub/core/diff/query/merge.py

@@ -36,6 +36,7 @@ class DiffMergeQuery(Query):
             "target_branch": self.target_branch.name,
             "source_branch": self.source_branch_name,
         }
+        # ruff: noqa: E501
         query = """
 UNWIND $node_diff_dicts AS node_diff_map
 CALL {
@@ -242,9 +243,11 @@ CALL {
                 CASE
                     WHEN startNode(source_r_rel_2).uuid = r.uuid THEN "r"
                     ELSE "l"
-                END AS r2_dir
+                END AS r2_dir,
+                source_r_rel_1.hierarchy AS r1_hierarchy,
+                source_r_rel_2.hierarchy AS r2_hierarchy
             }
-            WITH n, r, r1_dir, r2_dir, rel_name, rel_peer_id, related_rel_status
+            WITH n, r, r1_dir, r2_dir, r1_hierarchy, r2_hierarchy, rel_name, rel_peer_id, related_rel_status
             CALL {
                 WITH n, rel_name, rel_peer_id, related_rel_status
                 OPTIONAL MATCH (n)
@@ -258,12 +261,12 @@ CALL {
                 SET target_r_rel_1.to = $at
                 SET target_r_rel_2.to = $at
             }
-            WITH n, r, r1_dir, r2_dir, rel_name, rel_peer_id, related_rel_status
+            WITH n, r, r1_dir, r2_dir, r1_hierarchy, r2_hierarchy, rel_name, rel_peer_id, related_rel_status
             // ------------------------------
             // conditionally create new IS_RELATED relationships on target_branch, if necessary
             // ------------------------------
             CALL {
-                WITH n, r, r1_dir, r2_dir, rel_name, rel_peer_id, related_rel_status
+                WITH n, r, r1_dir, r2_dir, r1_hierarchy, r2_hierarchy, rel_name, rel_peer_id, related_rel_status
                 MATCH (p:Node {uuid: rel_peer_id})
                 OPTIONAL MATCH (n)
                     -[r_rel_1:IS_RELATED {branch: $target_branch, status: related_rel_status}]
@@ -274,42 +277,42 @@ CALL {
                 AND (r_rel_1.to >= $at OR r_rel_1.to IS NULL)
                 AND r_rel_2.from <= $at
                 AND (r_rel_2.to >= $at OR r_rel_2.to IS NULL)
-                WITH n, r, r1_dir, r2_dir, p, related_rel_status, r_rel_1, r_rel_2
+                WITH n, r, r1_dir, r2_dir, r1_hierarchy, r2_hierarchy, p, related_rel_status, r_rel_1, r_rel_2
                 WHERE r_rel_1 IS NULL
                 AND r_rel_2 IS NULL
                 // ------------------------------
                 // create IS_RELATED relationships with directions maintained from source
                 // ------------------------------
                 CALL {
-                    WITH n, r, r1_dir, related_rel_status
-                    WITH n, r, r1_dir, related_rel_status
+                    WITH n, r, r1_dir, r1_hierarchy, related_rel_status
+                    WITH n, r, r1_dir, r1_hierarchy, related_rel_status
                     WHERE r1_dir = "r"
                     CREATE (n)
-                        -[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status}]
+                        -[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status, hierarchy: r1_hierarchy}]
                         ->(r)
                 }
                 CALL {
-                    WITH n, r, r1_dir, related_rel_status
-                    WITH n, r, r1_dir, related_rel_status
+                    WITH n, r, r1_dir, r1_hierarchy, related_rel_status
+                    WITH n, r, r1_dir, r1_hierarchy, related_rel_status
                     WHERE r1_dir = "l"
                     CREATE (n)
-                        <-[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status}]
+                        <-[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status, hierarchy: r1_hierarchy}]
                         -(r)
                 }
                 CALL {
-                    WITH r, p, r2_dir, related_rel_status
-                    WITH r, p, r2_dir, related_rel_status
+                    WITH r, p, r2_dir, r2_hierarchy, related_rel_status
+                    WITH r, p, r2_dir, r2_hierarchy, related_rel_status
                     WHERE r2_dir = "r"
                     CREATE (r)
-                        -[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status}]
+                        -[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status, hierarchy: r2_hierarchy}]
                         ->(p)
                 }
                 CALL {
-                    WITH r, p, r2_dir, related_rel_status
-                    WITH r, p, r2_dir, related_rel_status
+                    WITH r, p, r2_dir, r2_hierarchy, related_rel_status
+                    WITH r, p, r2_dir, r2_hierarchy, related_rel_status
                     WHERE r2_dir = "l"
                     CREATE (r)
-                        <-[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status}]
+                        <-[:IS_RELATED {branch: $target_branch, branch_level: $branch_level, from: $at, status: related_rel_status, hierarchy: r2_hierarchy}]
                         -(p)
                 }
             }

infrahub/core/graph/__init__.py

@@ -1 +1 @@
-GRAPH_VERSION = 20
+GRAPH_VERSION = 21

infrahub/core/migrations/graph/__init__.py

@@ -22,6 +22,7 @@ from .m017_add_core_profile import Migration017
 from .m018_uniqueness_nulls import Migration018
 from .m019_restore_rels_to_time import Migration019
 from .m020_duplicate_edges import Migration020
+from .m021_missing_hierarchy_merge import Migration021
 
 if TYPE_CHECKING:
     from infrahub.core.root import Root
@@ -49,6 +50,7 @@ MIGRATIONS: list[type[Union[GraphMigration, InternalSchemaMigration, ArbitraryMi
     Migration018,
     Migration019,
     Migration020,
+    Migration021,
 ]
 
 

infrahub/core/migrations/graph/m021_missing_hierarchy_merge.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Sequence
+
+from infrahub.core.migrations.shared import GraphMigration, MigrationResult
+from infrahub.log import get_logger
+
+from ...query import Query, QueryType
+
+if TYPE_CHECKING:
+    from infrahub.database import InfrahubDatabase
+
+log = get_logger()
+
+
+class SetMissingHierarchyQuery(Query):
+    name = "set_missing_hierarchy"
+    type = QueryType.WRITE
+    insert_return = False
+
+    async def query_init(self, db: InfrahubDatabase, **kwargs: dict[str, Any]) -> None:
+        query = """
+        MATCH (r:Root)
+        WITH r.default_branch AS default_branch
+        MATCH (n:Node)-[main_e:IS_RELATED {branch: default_branch}]-(rel:Relationship)
+        WHERE main_e.hierarchy IS NULL
+        CALL {
+            WITH n, main_e, rel
+            MATCH (n)-[branch_e:IS_RELATED]-(rel)
+            WHERE branch_e.hierarchy IS NOT NULL
+            AND branch_e.branch <> main_e.branch
+            AND branch_e.from < main_e.from
+            SET main_e.hierarchy = branch_e.hierarchy
+        }
+        """
+        self.add_to_query(query)
+
+
+class Migration021(GraphMigration):
+    """
+    A bug in diff merge logic caused the hierarchy information on IS_RELATED edges to be lost when merged into
+    main. This migration sets the missing hierarchy data.
+    """
+
+    name: str = "021_replace_hierarchy"
+    minimum_version: int = 20
+    queries: Sequence[type[Query]] = [SetMissingHierarchyQuery]
+
+    async def validate_migration(self, db: InfrahubDatabase) -> MigrationResult:
+        result = MigrationResult()
+        return result

infrahub/core/query/node.py

@@ -666,18 +666,21 @@ class NodeListGetRelationshipsQuery(Query):
         MATCH paths_in = ((n)<-[r1:IS_RELATED]-(rel:Relationship)<-[r2:IS_RELATED]-(peer))
         WHERE ($relationship_identifiers IS NULL OR rel.name in $relationship_identifiers)
         AND all(r IN relationships(paths_in) WHERE (%(filters)s))
+        AND n.uuid <> peer.uuid
         RETURN n, rel, peer, r1, r2, "inbound" as direction
         UNION
         MATCH (n:Node) WHERE n.uuid IN $ids
         MATCH paths_out = ((n)-[r1:IS_RELATED]->(rel:Relationship)-[r2:IS_RELATED]->(peer))
         WHERE ($relationship_identifiers IS NULL OR rel.name in $relationship_identifiers)
         AND all(r IN relationships(paths_out) WHERE (%(filters)s))
+        AND n.uuid <> peer.uuid
         RETURN n, rel, peer, r1, r2, "outbound" as direction
         UNION
         MATCH (n:Node) WHERE n.uuid IN $ids
         MATCH paths_bidir = ((n)-[r1:IS_RELATED]->(rel:Relationship)<-[r2:IS_RELATED]-(peer))
         WHERE ($relationship_identifiers IS NULL OR rel.name in $relationship_identifiers)
         AND all(r IN relationships(paths_bidir) WHERE (%(filters)s))
+        AND n.uuid <> peer.uuid
         RETURN n, rel, peer, r1, r2, "bidirectional" as direction
         """ % {"filters": rels_filter}
 

infrahub/database/__init__.py

@@ -34,7 +34,7 @@ from infrahub.utils import InfrahubStringEnum
 
 from .constants import DatabaseType, Neo4jRuntime
 from .memgraph import DatabaseManagerMemgraph
-from .metrics import QUERY_EXECUTION_METRICS, TRANSACTION_RETRIES
+from .metrics import CONNECTION_POOL_USAGE, QUERY_EXECUTION_METRICS, TRANSACTION_RETRIES
 from .neo4j import DatabaseManagerNeo4j
 
 if TYPE_CHECKING:
@@ -335,6 +335,14 @@ class InfrahubDatabase:
         context: dict[str, str] | None = None,
         type: QueryType | None = None,  # pylint: disable=redefined-builtin
     ) -> tuple[list[Record], dict[str, Any]]:
+        connpool_usage = self._driver._pool.in_use_connection_count(self._driver._pool.address)
+        CONNECTION_POOL_USAGE.labels(self._driver._pool.address).set(float(connpool_usage))
+
+        if config.SETTINGS.database.max_concurrent_queries:
+            while connpool_usage > config.SETTINGS.database.max_concurrent_queries:  # noqa: ASYNC110
+                await asyncio.sleep(config.SETTINGS.database.max_concurrent_queries_delay)
+                connpool_usage = self._driver._pool.in_use_connection_count(self._driver._pool.address)
+
         with trace.get_tracer(__name__).start_as_current_span("execute_db_query_with_metadata") as span:
             span.set_attribute("query", query)
             if name:

infrahub/database/metrics.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from prometheus_client import Counter, Histogram
+from prometheus_client import Counter, Gauge, Histogram
 
 METRIC_PREFIX = "infrahub_db"
 
@@ -16,3 +16,9 @@ TRANSACTION_RETRIES = Counter(
     "Number of transaction that have been retried due to transcient error",
     labelnames=["name"],
 )
+
+CONNECTION_POOL_USAGE = Gauge(
+    f"{METRIC_PREFIX}_last_connection_pool_usage",
+    "Number of last known active connections in the pool",
+    labelnames=["address"],
+)

infrahub/graphql/initialization.py

@@ -8,6 +8,7 @@ from starlette.background import BackgroundTasks
 from infrahub.core import registry
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import InitializationError
+from infrahub.graphql.resolvers.many_relationship import ManyRelationshipResolver
 from infrahub.graphql.resolvers.single_relationship import SingleRelationshipResolver
 from infrahub.permissions import PermissionManager
 
@@ -35,6 +36,7 @@ class GraphqlContext:
     branch: Branch
     types: dict
     single_relationship_resolver: SingleRelationshipResolver
+    many_relationship_resolver: ManyRelationshipResolver
     at: Timestamp | None = None
     related_node_ids: set | None = None
     service: InfrahubServices | None = None
@@ -107,6 +109,7 @@ async def prepare_graphql_params(
             db=db,
             branch=branch,
             single_relationship_resolver=SingleRelationshipResolver(),
+            many_relationship_resolver=ManyRelationshipResolver(),
             at=Timestamp(at),
             types=gqlm.get_graphql_types(),
             related_node_ids=set(),

infrahub/graphql/loaders/node.py

@@ -10,6 +10,8 @@ from infrahub.core.node import Node
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
 
+from .shared import to_frozen_set
+
 
 @dataclass
 class GetManyParams:
@@ -68,15 +70,3 @@ class NodeDataLoader(DataLoader[str, Node | None]):
         for node_id in keys:
             results.append(nodes_by_id.get(node_id, None))
         return results
-
-
-def to_frozen_set(to_freeze: dict[str, Any]) -> frozenset:
-    freezing_dict = {}
-    for k, v in to_freeze.items():
-        if isinstance(v, dict):
-            freezing_dict[k] = to_frozen_set(v)
-        elif isinstance(v, (list, set)):
-            freezing_dict[k] = frozenset(v)
-        else:
-            freezing_dict[k] = v
-    return frozenset(freezing_dict)

infrahub/graphql/loaders/peers.py

@@ -0,0 +1,77 @@
+from dataclasses import dataclass
+from typing import Any
+
+from aiodataloader import DataLoader
+
+from infrahub.core.branch.models import Branch
+from infrahub.core.manager import NodeManager
+from infrahub.core.relationship.model import Relationship
+from infrahub.core.schema.relationship_schema import RelationshipSchema
+from infrahub.core.timestamp import Timestamp
+from infrahub.database import InfrahubDatabase
+
+from .shared import to_frozen_set
+
+
+@dataclass
+class QueryPeerParams:
+    branch: Branch | str
+    source_kind: str
+    schema: RelationshipSchema
+    filters: dict[str, Any]
+    fields: dict | None = None
+    at: Timestamp | str | None = None
+    branch_agnostic: bool = False
+
+    def __hash__(self) -> int:
+        frozen_fields: frozenset | None = None
+        if self.fields:
+            frozen_fields = to_frozen_set(self.fields)
+        frozen_filters = to_frozen_set(self.filters)
+        timestamp = Timestamp(self.at)
+        branch = self.branch.name if isinstance(self.branch, Branch) else self.branch
+        hash_str = "|".join(
+            [
+                str(hash(frozen_fields)),
+                str(hash(frozen_filters)),
+                timestamp.to_string(),
+                branch,
+                self.schema.name,
+                str(self.source_kind),
+                str(self.branch_agnostic),
+            ]
+        )
+        return hash(hash_str)
+
+
+class PeerRelationshipsDataLoader(DataLoader[str, list[Relationship]]):
+    def __init__(self, db: InfrahubDatabase, query_params: QueryPeerParams, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self.query_params = query_params
+        self.db = db
+
+    async def batch_load_fn(self, keys: list[Any]) -> list[list[Relationship]]:  # pylint: disable=method-hidden
+        async with self.db.start_session() as db:
+            peer_rels = await NodeManager.query_peers(
+                db=db,
+                ids=keys,
+                source_kind=self.query_params.source_kind,
+                schema=self.query_params.schema,
+                filters=self.query_params.filters,
+                fields=self.query_params.fields,
+                at=self.query_params.at,
+                branch=self.query_params.branch,
+                branch_agnostic=self.query_params.branch_agnostic,
+                fetch_peers=True,
+            )
+        peer_rels_by_node_id: dict[str, list[Relationship]] = {}
+        for rel in peer_rels:
+            node_id = rel.node_id
+            if node_id not in peer_rels_by_node_id:
+                peer_rels_by_node_id[node_id] = []
+            peer_rels_by_node_id[node_id].append(rel)
+
+        results = []
+        for node_id in keys:
+            results.append(peer_rels_by_node_id.get(node_id, []))
+        return results

infrahub/graphql/loaders/shared.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+
+def to_frozen_set(to_freeze: dict[str, Any]) -> frozenset:
+    freezing_dict = {}
+    for k, v in to_freeze.items():
+        if isinstance(v, dict):
+            freezing_dict[k] = to_frozen_set(v)
+        elif isinstance(v, (list, set)):
+            freezing_dict[k] = frozenset(v)
+        else:
+            freezing_dict[k] = v
+    return frozenset(freezing_dict)

infrahub/graphql/resolvers/many_relationship.py

@@ -0,0 +1,264 @@
+from typing import TYPE_CHECKING, Any
+
+from graphql import GraphQLResolveInfo
+from infrahub_sdk.utils import deep_merge_dict, extract_fields
+
+from infrahub.core.branch.models import Branch
+from infrahub.core.constants import BranchSupportType, RelationshipHierarchyDirection
+from infrahub.core.manager import NodeManager
+from infrahub.core.query.node import NodeGetHierarchyQuery
+from infrahub.core.schema.node_schema import NodeSchema
+from infrahub.core.schema.relationship_schema import RelationshipSchema
+from infrahub.core.timestamp import Timestamp
+from infrahub.database import InfrahubDatabase
+
+from ..loaders.peers import PeerRelationshipsDataLoader, QueryPeerParams
+from ..types import RELATIONS_PROPERTY_MAP, RELATIONS_PROPERTY_MAP_REVERSED
+
+if TYPE_CHECKING:
+    from infrahub.core.schema import MainSchemaTypes
+
+    from ..initialization import GraphqlContext
+
+
+class ManyRelationshipResolver:
+    def __init__(self) -> None:
+        self._data_loader_instances: dict[QueryPeerParams, PeerRelationshipsDataLoader] = {}
+
+    async def get_descendant_ids(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        at: Timestamp | None,
+        parent_id: str,
+        node_schema: NodeSchema,
+    ) -> list[str]:
+        async with db.start_session() as dbs:
+            query = await NodeGetHierarchyQuery.init(
+                db=dbs,
+                direction=RelationshipHierarchyDirection.DESCENDANTS,
+                node_id=parent_id,
+                node_schema=node_schema,
+                at=at,
+                branch=branch,
+            )
+            await query.execute(db=dbs)
+        return list(query.get_peer_ids())
+
+    async def get_peer_count(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        at: Timestamp | None,
+        ids: list[str],
+        source_kind: str,
+        rel_schema: RelationshipSchema,
+        filters: dict[str, Any],
+    ) -> int:
+        async with db.start_session() as dbs:
+            return await NodeManager.count_peers(
+                db=dbs,
+                ids=ids,
+                source_kind=source_kind,
+                schema=rel_schema,
+                filters=filters,
+                at=at,
+                branch=branch,
+                branch_agnostic=rel_schema.branch is BranchSupportType.AGNOSTIC,
+            )
+
+    async def resolve(
+        self,
+        parent: dict,
+        info: GraphQLResolveInfo,
+        include_descendants: bool = False,
+        offset: int | None = None,
+        limit: int | None = None,
+        **kwargs: Any,
+    ) -> dict[str, Any]:
+        """Resolver for relationships of cardinality=one for Edged responses
+
+        This resolver is used for paginated responses and as such we redefined the requested
+        fields by only reusing information below the 'node' key.
+        """
+        # Extract the InfraHub schema by inspecting the GQL Schema
+
+        node_schema: MainSchemaTypes = info.parent_type.graphene_type._meta.schema  # type: ignore[attr-defined]
+
+        context: GraphqlContext = info.context
+
+        # Extract the name of the fields in the GQL query
+        fields = await extract_fields(info.field_nodes[0].selection_set)
+        edges = fields.get("edges", {})
+        node_fields = edges.get("node", {})
+        property_fields = edges.get("properties", {})
+        for key, value in property_fields.items():
+            mapped_name = RELATIONS_PROPERTY_MAP[key]
+            node_fields[mapped_name] = value
+
+        filters = {
+            f"{info.field_name}__{key}": value
+            for key, value in kwargs.items()
+            if "__" in key and value or key in ["id", "ids"]
+        }
+
+        response: dict[str, Any] = {"edges": [], "count": None}
+
+        # Extract the schema of the node on the other end of the relationship from the GQL Schema
+        node_rel = node_schema.get_relationship(info.field_name)
+        source_kind = node_schema.kind
+        ids = [parent["id"]]
+        if isinstance(node_schema, NodeSchema):
+            if node_schema.hierarchy:
+                source_kind = node_schema.hierarchy
+
+            if include_descendants:
+                descendant_ids = await self.get_descendant_ids(
+                    db=context.db,
+                    branch=context.branch,
+                    at=context.at,
+                    parent_id=ids[0],
+                    node_schema=node_schema,
+                )
+                ids.extend(descendant_ids)
+
+        if "count" in fields:
+            peer_count = await self.get_peer_count(
+                db=context.db,
+                branch=context.branch,
+                at=context.at,
+                ids=ids,
+                source_kind=source_kind,
+                rel_schema=node_rel,
+                filters=filters,
+            )
+            response["count"] = peer_count
+
+        if not node_fields:
+            return response
+
+        if offset or limit:
+            node_graph = await self._get_entities_simple(
+                db=context.db,
+                branch=context.branch,
+                ids=ids,
+                at=context.at,
+                related_node_ids=context.related_node_ids,
+                source_kind=source_kind,
+                rel_schema=node_rel,
+                filters=filters,
+                node_fields=node_fields,
+                offset=offset,
+                limit=limit,
+            )
+        else:
+            node_graph = await self._get_entities_with_data_loader(
+                db=context.db,
+                branch=context.branch,
+                ids=ids,
+                at=context.at,
+                related_node_ids=context.related_node_ids,
+                source_kind=source_kind,
+                rel_schema=node_rel,
+                filters=filters,
+                node_fields=node_fields,
+            )
+
+        if not node_graph:
+            return response
+
+        entries = []
+        for node in node_graph:
+            entry: dict[str, dict[str, Any]] = {"node": {}, "properties": {}}
+            for key, mapped in RELATIONS_PROPERTY_MAP_REVERSED.items():
+                value = node.pop(key, None)
+                if value:
+                    entry["properties"][mapped] = value
+            entry["node"] = node
+            entries.append(entry)
+
+        response["edges"] = entries
+        return response
+
+    async def _get_entities_simple(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        ids: list[str],
+        at: Timestamp | None,
+        related_node_ids: set[str] | None,
+        source_kind: str,
+        rel_schema: RelationshipSchema,
+        filters: dict[str, Any],
+        node_fields: dict[str, Any],
+        offset: int | None = None,
+        limit: int | None = None,
+    ) -> list[dict[str, Any]] | None:
+        async with db.start_session() as dbs:
+            objs = await NodeManager.query_peers(
+                db=dbs,
+                ids=ids,
+                source_kind=source_kind,
+                schema=rel_schema,
+                filters=filters,
+                fields=node_fields,
+                offset=offset,
+                limit=limit,
+                at=at,
+                branch=branch,
+                branch_agnostic=rel_schema.branch is BranchSupportType.AGNOSTIC,
+                fetch_peers=True,
+            )
+            if not objs:
+                return None
+            return [await obj.to_graphql(db=dbs, fields=node_fields, related_node_ids=related_node_ids) for obj in objs]
+
+    async def _get_entities_with_data_loader(
+        self,
+        db: InfrahubDatabase,
+        branch: Branch,
+        ids: list[str],
+        at: Timestamp | None,
+        related_node_ids: set[str] | None,
+        source_kind: str,
+        rel_schema: RelationshipSchema,
+        filters: dict[str, Any],
+        node_fields: dict[str, Any],
+    ) -> list[dict[str, Any]] | None:
+        if node_fields and "display_label" in node_fields:
+            schema_branch = db.schema.get_schema_branch(name=branch.name)
+            display_label_fields = schema_branch.generate_fields_for_display_label(name=rel_schema.peer)
+            if display_label_fields:
+                node_fields = deep_merge_dict(dicta=node_fields, dictb=display_label_fields)
+
+        if node_fields and "hfid" in node_fields:
+            peer_schema = db.schema.get(name=rel_schema.peer, branch=branch, duplicate=False)
+            hfid_fields = peer_schema.generate_fields_for_hfid()
+            if hfid_fields:
+                node_fields = deep_merge_dict(dicta=node_fields, dictb=hfid_fields)
+
+        query_params = QueryPeerParams(
+            branch=branch,
+            source_kind=source_kind,
+            schema=rel_schema,
+            filters=filters,
+            fields=node_fields,
+            at=at,
+            branch_agnostic=rel_schema.branch is BranchSupportType.AGNOSTIC,
+        )
+        if query_params in self._data_loader_instances:
+            loader = self._data_loader_instances[query_params]
+        else:
+            loader = PeerRelationshipsDataLoader(db=db, query_params=query_params)
+            self._data_loader_instances[query_params] = loader
+        all_peer_rels = []
+        for node_id in ids:
+            node_peer_rels = await loader.load(key=node_id)
+            all_peer_rels.extend(node_peer_rels)
+        if not all_peer_rels:
+            return None
+        async with db.start_session() as dbs:
+            return [
+                await obj.to_graphql(db=dbs, fields=node_fields, related_node_ids=related_node_ids)
+                for obj in all_peer_rels
+            ]

infrahub/graphql/resolvers/resolver.py

@@ -6,13 +6,11 @@ from infrahub_sdk.utils import extract_fields
 
 from infrahub.core.constants import BranchSupportType, InfrahubKind, RelationshipHierarchyDirection
 from infrahub.core.manager import NodeManager
-from infrahub.core.query.node import NodeGetHierarchyQuery
 from infrahub.exceptions import NodeNotFoundError
 
 from ..models import OrderModel
 from ..parser import extract_selection
 from ..permissions import get_permissions
-from ..types import RELATIONS_PROPERTY_MAP, RELATIONS_PROPERTY_MAP_REVERSED
 
 if TYPE_CHECKING:
     from graphql import GraphQLResolveInfo
@@ -217,109 +215,11 @@ async def single_relationship_resolver(parent: dict, info: GraphQLResolveInfo, *
 
 
 async def many_relationship_resolver(
-    parent: dict, info: GraphQLResolveInfo, include_descendants: Optional[bool] = False, **kwargs: Any
+    parent: dict, info: GraphQLResolveInfo, include_descendants: bool | None = False, **kwargs: Any
 ) -> dict[str, Any]:
-    """Resolver for relationships of cardinality=many for Edged responses
-
-    This resolver is used for paginated responses and as such we redefined the requested
-    fields by only reusing information below the 'node' key.
-    """
-    # Extract the InfraHub schema by inspecting the GQL Schema
-    node_schema: NodeSchema = info.parent_type.graphene_type._meta.schema
-
     context: GraphqlContext = info.context
-
-    # Extract the name of the fields in the GQL query
-    fields = await extract_fields(info.field_nodes[0].selection_set)
-    edges = fields.get("edges", {})
-    node_fields = edges.get("node", {})
-    property_fields = edges.get("properties", {})
-    for key, value in property_fields.items():
-        mapped_name = RELATIONS_PROPERTY_MAP[key]
-        node_fields[mapped_name] = value
-
-    # Extract the schema of the node on the other end of the relationship from the GQL Schema
-    node_rel = node_schema.get_relationship(info.field_name)
-
-    # Extract only the filters from the kwargs and prepend the name of the field to the filters
-    offset = kwargs.pop("offset", None)
-    limit = kwargs.pop("limit", None)
-
-    filters = {
-        f"{info.field_name}__{key}": value
-        for key, value in kwargs.items()
-        if "__" in key and value or key in ["id", "ids"]
-    }
-
-    response: dict[str, Any] = {"edges": [], "count": None}
-
-    source_kind = node_schema.kind
-
-    async with context.db.start_session() as db:
-        ids = [parent["id"]]
-        if include_descendants:
-            query = await NodeGetHierarchyQuery.init(
-                db=db,
-                direction=RelationshipHierarchyDirection.DESCENDANTS,
-                node_id=parent["id"],
-                node_schema=node_schema,
-                at=context.at,
-                branch=context.branch,
-            )
-            if node_schema.hierarchy:
-                source_kind = node_schema.hierarchy
-            await query.execute(db=db)
-            descendants_ids = list(query.get_peer_ids())
-            ids.extend(descendants_ids)
-
-        if "count" in fields:
-            response["count"] = await NodeManager.count_peers(
-                db=db,
-                ids=ids,
-                source_kind=source_kind,
-                schema=node_rel,
-                filters=filters,
-                at=context.at,
-                branch=context.branch,
-                branch_agnostic=node_rel.branch is BranchSupportType.AGNOSTIC,
-            )
-
-        if not node_fields:
-            return response
-
-        objs = await NodeManager.query_peers(
-            db=db,
-            ids=ids,
-            source_kind=source_kind,
-            schema=node_rel,
-            filters=filters,
-            fields=node_fields,
-            offset=offset,
-            limit=limit,
-            at=context.at,
-            branch=context.branch,
-            branch_agnostic=node_rel.branch is BranchSupportType.AGNOSTIC,
-            fetch_peers=True,
-        )
-
-        if not objs:
-            return response
-        node_graph = [
-            await obj.to_graphql(db=db, fields=node_fields, related_node_ids=context.related_node_ids) for obj in objs
-        ]
-
-        entries = []
-        for node in node_graph:
-            entry = {"node": {}, "properties": {}}
-            for key, mapped in RELATIONS_PROPERTY_MAP_REVERSED.items():
-                value = node.pop(key, None)
-                if value:
-                    entry["properties"][mapped] = value
-            entry["node"] = node
-            entries.append(entry)
-        response["edges"] = entries
-
-        return response
+    resolver = context.many_relationship_resolver
+    return await resolver.resolve(parent=parent, info=info, include_descendants=include_descendants, **kwargs)
 
 
 async def ancestors_resolver(parent: dict, info: GraphQLResolveInfo, **kwargs) -> dict[str, Any]:

infrahub/graphql/subscription/graphql_query.py

@@ -8,6 +8,7 @@ from infrahub.core.constants import InfrahubKind
 from infrahub.core.manager import NodeManager
 from infrahub.core.protocols import CoreGraphQLQuery
 from infrahub.core.timestamp import Timestamp
+from infrahub.graphql.resolvers.many_relationship import ManyRelationshipResolver
 from infrahub.graphql.resolvers.single_relationship import SingleRelationshipResolver
 from infrahub.log import get_logger
 
@@ -48,6 +49,7 @@ async def resolver_graphql_query(
                     related_node_ids=set(),
                     types=context.types,
                     single_relationship_resolver=SingleRelationshipResolver(),
+                    many_relationship_resolver=ManyRelationshipResolver(),
                 ),
                 root_value=None,
                 variable_values=params or {},

{infrahub_server-1.1.8.dist-info → infrahub_server-1.1.9.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: infrahub-server
-Version: 1.1.8
+Version: 1.1.9
 Summary: Infrahub is taking a new approach to Infrastructure Management by providing a new generation of datastore to organize and control all the data that defines how an infrastructure should run.
 Home-page: https://opsmill.com
 License: AGPL-3.0-only

{infrahub_server-1.1.8.dist-info → infrahub_server-1.1.9.dist-info}/RECORD

@@ -33,7 +33,7 @@ infrahub/computed_attribute/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMp
 infrahub/computed_attribute/constants.py,sha256=oTMPEfRuf2mcfCkBpRLWRALO6nsLHpFm9jJGu0lowS4,446
 infrahub/computed_attribute/models.py,sha256=icUzsu0DrGoxMkBVXpNiv17rMo0OwpSE-QBJyWblMM0,2637
 infrahub/computed_attribute/tasks.py,sha256=rzzWBKrmCQ_ZXJ4Uv5XE7IDhdx39I0_HwyDTLYwkLKo,33238
-infrahub/config.py,sha256=DQFXa6e6nbW6McdLVViuCcKYlGjw9ytmYDsKU5l-Jdo,33605
+infrahub/config.py,sha256=lAf4zimE5IBU7FiQGDJAKRdBnbWawpGzcELwv6yMs8Y,33912
 infrahub/core/__init__.py,sha256=z6EJBZyCYCBqinoBtX9li6BTBbbGV8WCkE_4CrEsmDA,104
 infrahub/core/account.py,sha256=sggpuO_QpwYH7wXG_lZDnrB5Izmej486o_CYiYjYin8,26497
 infrahub/core/attribute.py,sha256=H0vHLoNlp1KHI2Jw-Ti5FLmlwJqpYDcv4-9o3K1AsR8,42150
@@ -91,7 +91,7 @@ infrahub/core/diff/query/field_summary.py,sha256=ZXYxX5f-nJ0rG_ZzfSxmF01OmkBYdSj
 infrahub/core/diff/query/filters.py,sha256=HnbeTo3W2n0gQcKYM4my7sQ0Syoj7aPy5WFkVp38qLc,3646
 infrahub/core/diff/query/get_conflict_query.py,sha256=fXJNx5hiZW3NLAbGITvC-iuRDhRkXFhLF3oRTOzgEqM,876
 infrahub/core/diff/query/has_conflicts_query.py,sha256=nftOmnAtDsNxxhpEoLcbQdalhgqx5ROckabou1CSv-Y,2531
-infrahub/core/diff/query/merge.py,sha256=HZgsf9tMIEZ_O-ztWBmF-L75nAO9LOe1lKo53Ehub9I,22902
+infrahub/core/diff/query/merge.py,sha256=mk80O6Rsk7-P27C6YiOI96pS4smLqMuNfVFyo-vxXIY,23369
 infrahub/core/diff/query/merge_tracking_id.py,sha256=puDyqQv-g0jHA-ZuNNwHyJ2spbWOUoSZr131d0NOu0M,817
 infrahub/core/diff/query/roots_metadata.py,sha256=mwS8elSVxH-xHYdMd0maAGJ94sGY9oXk85jp59RsCQ4,2166
 infrahub/core/diff/query/save.py,sha256=tfeWCy5FdBgNcPWIO86NZb5DyemX18fLSCDqJ86zsCo,21617
@@ -104,7 +104,7 @@ infrahub/core/diff/repository/deserializer.py,sha256=gkFPByBY1nnmA6sjTJer76RWvwH
 infrahub/core/diff/repository/repository.py,sha256=xgTzmd_fdc-n7iX8E83sd3fOz25O4P3CEDQFpRMZjpI,24946
 infrahub/core/diff/tasks.py,sha256=0EuasZKVk76ECpv5yKuqJ7_kGoplqx-_x7w-Di9IAsY,3018
 infrahub/core/enums.py,sha256=5wMcX9x6acU9CTa4B4b6rFwgRZ31N9c9TR3n2EO0BuI,490
-infrahub/core/graph/__init__.py,sha256=Exfs4Nq05DCpF2AkOkBelW_MeAF7Te0NRo0FEhOQHl8,19
+infrahub/core/graph/__init__.py,sha256=MvHNUs4wiJa0gTdCMF6WObMfqStA5V391r31Y9RIVtE,19
 infrahub/core/graph/constraints.py,sha256=lmuzrKDFoeSKRiLtycB9PXi6zhMYghczKrPYvfWyy90,10396
 infrahub/core/graph/index.py,sha256=oR6wyYpJbq2IVVzUdiuGyWA511hw2AvgklFoBmQk-bM,1619
 infrahub/core/graph/schema.py,sha256=FmEPPb1XOFv3nnS_XJCuUqlp8HsStX5A2frHjlhoqvE,10105
@@ -123,7 +123,7 @@ infrahub/core/ipam/utilization.py,sha256=Urv0thyR6xYgwyQaZDnx170Wcw8nKKZkBymwNTM
 infrahub/core/manager.py,sha256=e2rawz77UG3tLOz3J-JL6-UpJDyJui3uSDg_APnFnAE,46544
 infrahub/core/merge.py,sha256=ibXM0Rb8qVoBuGiW8q6JYdFsLQ9DS-PPTBoK4R2mPhg,10354
 infrahub/core/migrations/__init__.py,sha256=PBewY3fZkqVMABRo_oTZkDtdD7HfCC9nCn-DXtTca1g,1150
-infrahub/core/migrations/graph/__init__.py,sha256=f1OGl0qaYxr9Rotk3bl08fZ8CrBW6NO2TfT_axD6u7Y,2178
+infrahub/core/migrations/graph/__init__.py,sha256=jKTE6FK5RRuudsnal6cL1wXoy7pAHUR3k6GZGBMFlqQ,2251
 infrahub/core/migrations/graph/m001_add_version_to_graph.py,sha256=x_dYBnrZtNQiB6TSl4xwXR-Phn7-4EgrJce76ZfPe_c,1499
 infrahub/core/migrations/graph/m002_attribute_is_default.py,sha256=m0uW3rkDjcYEmuHcTH8ngWxenJ5K0_TkVV00L-Ec6Mw,1004
 infrahub/core/migrations/graph/m003_relationship_parent_optional.py,sha256=69RIuLB6YoNOOJ9og1DuUxuWSdTbxbSv6uk25-KsHUI,2313
@@ -144,6 +144,7 @@ infrahub/core/migrations/graph/m017_add_core_profile.py,sha256=T-IrK3DW7m_xS4xp3
 infrahub/core/migrations/graph/m018_uniqueness_nulls.py,sha256=QPQT9ID6idIXizDG0xtzfob_XHG_5LNdv90qD7Lwjng,4785
 infrahub/core/migrations/graph/m019_restore_rels_to_time.py,sha256=H0pQLkn-iup_-dl7S-jA6BHAeAmz5oBx9IqVLIfcMkw,11718
 infrahub/core/migrations/graph/m020_duplicate_edges.py,sha256=ec5Z_HF_5MKEQ6gAEvzRM8sxdAsD3lg3eR96fMQ-2lI,6722
+infrahub/core/migrations/graph/m021_missing_hierarchy_merge.py,sha256=JSfkkuqvmSwlXDgyAN__Ypn4J5qhCUPdzAKiaZgqWYU,1606
 infrahub/core/migrations/query/__init__.py,sha256=JoWOUWlV6IzwxWxObsfCnAAKUOHJkE7dZlOsfB64ZEo,876
 infrahub/core/migrations/query/attribute_add.py,sha256=gLibqL1TKtt8ia1UQBxL8vyVDibUPlP3w_vp5bp4XsQ,3507
 infrahub/core/migrations/query/attribute_rename.py,sha256=lotPE_XRqyJQisnrgXtH_cg-0qHMqhtUlT_xcRP84TU,6955
@@ -186,7 +187,7 @@ infrahub/core/query/branch.py,sha256=OmYS1n1U4NEbDXBrPAM9lvh_qb_prbRdNnC07N1k-Mw
 infrahub/core/query/delete.py,sha256=BsCeUb11Le3NU2hLMcu05acvWrRVrGncxGxpkWTIzJE,1902
 infrahub/core/query/diff.py,sha256=73Jy-D1wd6GYpUoRb7MGDMkgSuX6KeQ_8swaYLhgJQU,31942
 infrahub/core/query/ipam.py,sha256=xWCemBaiOA00eU3sB8XmOdNOJ1L69Lls3L2GIw5P_CU,28123
-infrahub/core/query/node.py,sha256=mPEChMQe5ivOxPBzLpWeHW3tJT5-oWn0PIgsm5TW05o,61092
+infrahub/core/query/node.py,sha256=nA7o6ZvPMY4Jx1C-LBzqmF_BTNPs_FI7gFuB11lyWe4,61188
 infrahub/core/query/relationship.py,sha256=Fzg9FvbFj9KsvlzWSwmOj2Ae9qPoaFJa4mzYe77D1qE,38628
 infrahub/core/query/resource_manager.py,sha256=3B-4-gTNomAyvkw88KHwHbjClIp4SlSJL0cGb-dZVyw,12568
 infrahub/core/query/standard_node.py,sha256=Wh9ekHhfamxizri_1M5EQFcgJWyIvtAMQRz9nEKOL3w,4437
@@ -266,12 +267,12 @@ infrahub/core/validators/uniqueness/checker.py,sha256=IqreVcMPZj7nNhC5cc5SlBSZ-P
 infrahub/core/validators/uniqueness/index.py,sha256=yu-clITQF4MrgK36hsyuXllvR4QkVTqy4ugi_Y_C_Sg,5081
 infrahub/core/validators/uniqueness/model.py,sha256=EPl8X91BSGXGU7GWbUSue6laNGhAtIiXj7rFaz56Kvk,5197
 infrahub/core/validators/uniqueness/query.py,sha256=DigQCR5278wBiXNXhy1FFTTVQewlf2Q8rCcuSmyilkQ,10171
-infrahub/database/__init__.py,sha256=NJfVx7yXiSBiolqU4JcSH3jDYrjjL9uDQVvHa-6cROo,20196
+infrahub/database/__init__.py,sha256=lF6RbdP0huD_XG9Y168yo6lBZBVHL93JsYFE_JuygzI,20766
 infrahub/database/constants.py,sha256=WmV1iuOk4xulxZHOVvO3sS_VF1eTf7fKh0TPe_RnfV4,507
 infrahub/database/index.py,sha256=y0sWXO3tdIr1wL1XC9O6iNRV-Elu2KAXFOiYXRIIhN4,1659
 infrahub/database/manager.py,sha256=BDXNw1RNBeSFV-EZd0aGFbPNuoqlKwrkDqmYB7sy4tU,317
 infrahub/database/memgraph.py,sha256=Jj5f8Rq2eNx1OSUiLVf8eM0Egf_8HG3wIuZJOPt9hAA,2018
-infrahub/database/metrics.py,sha256=lpE81u2glhrqCNj9RRv07GT9jaFhsVKZh9U0Y9OIeNA,560
+infrahub/database/metrics.py,sha256=xU4OSKFbsxcw_yZlt_39PmGtF7S7yPbPuOIlSCu5sI0,739
 infrahub/database/neo4j.py,sha256=IE5lSevKqu-tJa3KF_bj_gOUx-SpZNdmGOl6oheNhiY,2624
 infrahub/dependencies/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 infrahub/dependencies/builder/constraint/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -370,9 +371,11 @@ infrahub/graphql/auth/query_permission_checker/super_admin_checker.py,sha256=8gC
 infrahub/graphql/constants.py,sha256=iVvo3HK-ch7YmHw1Eg2E_ja3I45cNAwjpYahsnu85CI,37
 infrahub/graphql/directives.py,sha256=wyIkJFp7l0J4JqNl1Lqu7YfKXP7glrewlQFMDTUAPcE,645
 infrahub/graphql/enums.py,sha256=CjGgywkA44MrmTmHxoOoieRPtdmk1bEYzVGxAXUT5as,819
-infrahub/graphql/initialization.py,sha256=LMoVOcAeOtnFN6mZV_rLcnjSCBDTY564fLS8TivZ1ak,4055
+infrahub/graphql/initialization.py,sha256=pnXsXHHtmeJghm_yR3ddzk4NNR_nQ379p72Q5uqnddU,4261
 infrahub/graphql/loaders/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-infrahub/graphql/loaders/node.py,sha256=gPV8tmbmmiusRbLk-wHlOyDIoGBXnpiDYiuKNsXapaM,2969
+infrahub/graphql/loaders/node.py,sha256=sY6mMEG7ApYi3zGEzyIMTkBxMelX3dPGSgDroTSbW8Y,2637
+infrahub/graphql/loaders/peers.py,sha256=fOdgaaPBtBiiLKg8Rws4zXvpyCzUhszgxhDAeXkzETA,2785
+infrahub/graphql/loaders/shared.py,sha256=7O0Jaeb5-urt75XfRQ97a_pgJgs6ZY0TnKAZPUI9eKs,390
 infrahub/graphql/manager.py,sha256=qF23ffrdbzy4p2crY1RATX8w-AklZkENolM-bNTKJJY,44303
 infrahub/graphql/metrics.py,sha256=viq_M57mDYd4DDK7suUttf1FJTgzQ3U50yOuSw_Nd-s,2267
 infrahub/graphql/models.py,sha256=7kr3DSO_rujPocMIfPyZ5Hwy3Mpnu4ySDMAIE9G5Y7Y,147
@@ -416,12 +419,13 @@ infrahub/graphql/queries/status.py,sha256=F3UpRHocUrJRQeOyI1PYa9W-2YDwTBTdGU-CyC
 infrahub/graphql/queries/task.py,sha256=qBL54BBKRW3vRxg1bIIM88c6Yfrq6IiH2-zc-m-NdTk,3491
 infrahub/graphql/query.py,sha256=FbYKa9Y5h33jKR8JA4mD3W7FXkCsTd8qWbJ2CdiUXKk,1587
 infrahub/graphql/resolvers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-infrahub/graphql/resolvers/resolver.py,sha256=GXtO1ik4-msL98MShk-aSzMucLvuKNM5jj936AcOCho,14099
+infrahub/graphql/resolvers/many_relationship.py,sha256=jthzwVi8IZ-ZVpzr0OkFuknd9MD0rWmdRIzSVmigkd0,9587
+infrahub/graphql/resolvers/resolver.py,sha256=1UMP36G467BCkIKjzweGxNfiMzNqNL_QZcY9WuqqElY,10579
 infrahub/graphql/resolvers/single_relationship.py,sha256=3MtfoR0w-OUcuOkZXVje3BbTTCSPzabh74jSm3tyAPs,6660
 infrahub/graphql/schema.py,sha256=4M1YYW3yUO99TKsCUDBG2DAyHsvHVoy6eBbRmg7s_Qk,3498
 infrahub/graphql/subscription/__init__.py,sha256=BaGpWIjP_O8kSHKvWtRQseJtSpcfRlePNPhZf6HrTa0,1189
 infrahub/graphql/subscription/events.py,sha256=dlawTJljVc7SN5kKag5jOyW1lQuxEBq3PIMSHXQ-09M,1873
-infrahub/graphql/subscription/graphql_query.py,sha256=j-W9GzfTH6vIfI0GSMmCiSEYu9Kg6bfjjfGAFIxm7_Q,2044
+infrahub/graphql/subscription/graphql_query.py,sha256=4XimqnKfEF7E318ZzA2rLazeu6_U11sSyq16Itw8c4I,2201
 infrahub/graphql/types/__init__.py,sha256=oP4DhYAteHkc8LiQsIRHE1q2jaksfc-VvAO1urkmI4k,1895
 infrahub/graphql/types/attribute.py,sha256=bc2q44q8j5DTNdwBMe0SgG7Rbk8si1h-8SyqGIa-Rn0,6594
 infrahub/graphql/types/branch.py,sha256=8C07E3ne1Rk0MQtzLchaJwiQqPNDHIVT59t_y1d3wNs,1393
@@ -673,12 +677,12 @@ infrahub_sdk/utils.py,sha256=S2uQvg2WROdOK2gtbQ6HSVWStTVqVkzmEUXcksuMZGI,11513
 infrahub_sdk/uuidt.py,sha256=Tz-4nHkJwbi39UT3gaIe2wJeZNAoBqf6tm3sw7LZbXc,2155
 infrahub_sdk/yaml.py,sha256=L_sj5ds-0_uKe3aIfZu86kDLq8tffKzle9dcyDUTaEc,2937
 infrahub_testcontainers/__init__.py,sha256=oPpmesGgYBSdKTg1L37FGwYBeao1EHury5SJGul-CT8,216
-infrahub_testcontainers/container.py,sha256=dtweuZol1LAZiJBItRnB-Yrd-k6ohVuOdcgWcd0F-RM,4855
-infrahub_testcontainers/docker-compose.test.yml,sha256=l-ac0Bi8MxXCNM0w7QPakmeMsYzMnJXPhENN1yGRHAk,5993
+infrahub_testcontainers/container.py,sha256=AuHmQv4fqB5Til40-GUsZhop9hzHWLWkqsRV44UZPUY,5870
+infrahub_testcontainers/docker-compose.test.yml,sha256=gDmMNoIh5ROQxVAw02_twaAn_wi3Fw63D1Ek2k8YlWc,5966
 infrahub_testcontainers/haproxy.cfg,sha256=QF_DJSll10uGWICLnwFpj0-0VWIeB3CteCDthVoWpCU,1316
 infrahub_testcontainers/helpers.py,sha256=NlzyZN1QpByUQ6t7NwIDQxSSM7xFvY2jCM6-P8RYE6g,5253
-infrahub_server-1.1.8.dist-info/LICENSE.txt,sha256=TfPDBt3ar0uv_f9cqCDMZ5rIzW3CY8anRRd4PkL6ejs,34522
-infrahub_server-1.1.8.dist-info/METADATA,sha256=BJcVX7r_XMQiEkLZjsOAZo-mVuL7k1prexbpw7eAmrs,8081
-infrahub_server-1.1.8.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
-infrahub_server-1.1.8.dist-info/entry_points.txt,sha256=JNQoBcLpUyfeOMhls_-uX1CdJ8Vl-AFSh9UhzTcKdjA,329
-infrahub_server-1.1.8.dist-info/RECORD,,
+infrahub_server-1.1.9.dist-info/LICENSE.txt,sha256=TfPDBt3ar0uv_f9cqCDMZ5rIzW3CY8anRRd4PkL6ejs,34522
+infrahub_server-1.1.9.dist-info/METADATA,sha256=N42-nCLMaBJBOiUs29mwCq-Sr_CHvZ3e4WT-RA8R7IU,8081
+infrahub_server-1.1.9.dist-info/WHEEL,sha256=Nq82e9rUAnEjt98J6MlVmMCZb-t9cYE2Ir1kpBmnWfs,88
+infrahub_server-1.1.9.dist-info/entry_points.txt,sha256=JNQoBcLpUyfeOMhls_-uX1CdJ8Vl-AFSh9UhzTcKdjA,329
+infrahub_server-1.1.9.dist-info/RECORD,,

infrahub_testcontainers/container.py

@@ -24,6 +24,7 @@ INFRAHUB_SERVICES: dict[str, ContainerService] = {
 
 PROJECT_ENV_VARIABLES: dict[str, str] = {
     "INFRAHUB_TESTING_DOCKER_IMAGE": "registry.opsmill.io/opsmill/infrahub",
+    "INFRAHUB_TESTING_DOCKER_ENTRYPOINT": f"gunicorn --config backend/infrahub/serve/gunicorn_config.py -w {os.environ.get("INFRAHUB_TESTING_WEB_CONCURRENCY", 4)} --logger-class infrahub.serve.log.GunicornLogger infrahub.server:app",  # noqa: E501
     "INFRAHUB_TESTING_IMAGE_VERSION": infrahub_version,
     "INFRAHUB_TESTING_PRODUCTION": "false",
     "INFRAHUB_TESTING_DB_ADDRESS": "database",
@@ -38,6 +39,7 @@ PROJECT_ENV_VARIABLES: dict[str, str] = {
     "INFRAHUB_TESTING_BROKER_ADDRESS": "message-queue",
     "INFRAHUB_TESTING_CACHE_ADDRESS": "cache",
     "INFRAHUB_TESTING_WORKFLOW_ADDRESS": "task-manager",
+    "INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE": "infrahubasync",
     "INFRAHUB_TESTING_TIMEOUT": "60",
     "INFRAHUB_TESTING_PREFECT_API": "http://task-manager:4200/api",
     "INFRAHUB_TESTING_LOCAL_REMOTE_GIT_DIRECTORY": "repos",
@@ -96,6 +98,15 @@ class InfrahubDockerCompose(DockerCompose):
         env_file = directory / ".env"
 
         PROJECT_ENV_VARIABLES.update({"INFRAHUB_TESTING_IMAGE_VERSION": version})
+        if os.environ.get("INFRAHUB_TESTING_ENTERPRISE"):
+            PROJECT_ENV_VARIABLES.update(
+                {
+                    "INFRAHUB_TESTING_DOCKER_IMAGE": "registry.opsmill.io/opsmill/infrahub-enterprise",
+                    "INFRAHUB_TESTING_DOCKER_ENTRYPOINT": f"gunicorn --config community/backend/infrahub/serve/gunicorn_config.py -w {os.environ.get("INFRAHUB_TESTING_WEB_CONCURRENCY", 4)} --logger-class infrahub.serve.log.GunicornLogger infrahub_enterprise.server:app",  # noqa: E501
+                    "INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE": "infrahubentasync",
+                    "NEO4J_DOCKER_IMAGE": "neo4j:5.20.0-enterprise",
+                }
+            )
 
         with env_file.open(mode="w", encoding="utf-8") as file:
             for key, value in PROJECT_ENV_VARIABLES.items():

infrahub_testcontainers/docker-compose.test.yml

@@ -102,11 +102,7 @@ services:
       mode: replicated
       replicas: ${INFRAHUB_TESTING_API_SERVER_COUNT}
     image: "${INFRAHUB_TESTING_DOCKER_IMAGE}:${INFRAHUB_TESTING_IMAGE_VERSION}"
-    command: >
-      gunicorn --config backend/infrahub/serve/gunicorn_config.py
-      -w ${INFRAHUB_TESTING_WEB_CONCURRENCY}
-      --logger-class infrahub.serve.log.GunicornLogger
-      infrahub.server:app
+    command: ${INFRAHUB_TESTING_DOCKER_ENTRYPOINT}
     environment:
       INFRAHUB_PRODUCTION: ${INFRAHUB_TESTING_PRODUCTION}
       INFRAHUB_LOG_LEVEL: ${INFRAHUB_TESTING_LOG_LEVEL:-INFO}
@@ -114,6 +110,7 @@ services:
       INFRAHUB_CACHE_ADDRESS: ${INFRAHUB_TESTING_CACHE_ADDRESS}
       INFRAHUB_DB_ADDRESS: ${INFRAHUB_TESTING_DB_ADDRESS}
       INFRAHUB_WORKFLOW_ADDRESS: ${INFRAHUB_TESTING_WORKFLOW_ADDRESS}
+      INFRAHUB_WORKFLOW_DEFAULT_WORKER_TYPE: ${INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE}
       INFRAHUB_INITIAL_ADMIN_TOKEN: ${INFRAHUB_TESTING_INITIAL_ADMIN_TOKEN}
       INFRAHUB_INITIAL_AGENT_TOKEN: ${INFRAHUB_TESTING_INITIAL_AGENT_TOKEN}
       INFRAHUB_SECURITY_SECRET_KEY: ${INFRAHUB_TESTING_SECURITY_SECRET_KEY}
@@ -142,7 +139,7 @@ services:
       mode: replicated
       replicas: ${INFRAHUB_TESTING_TASK_WORKER_COUNT}
     image: "${INFRAHUB_TESTING_DOCKER_IMAGE}:${INFRAHUB_TESTING_IMAGE_VERSION}"
-    command: prefect worker start --type infrahubasync --pool infrahub-worker --with-healthcheck
+    command: prefect worker start --type ${INFRAHUB_TESTING_WORKFLOW_DEFAULT_WORKER_TYPE} --pool infrahub-worker --with-healthcheck
     environment:
       INFRAHUB_PRODUCTION: ${INFRAHUB_TESTING_PRODUCTION}
       INFRAHUB_LOG_LEVEL: ${INFRAHUB_TESTING_LOG_LEVEL}