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

infrahub/graphql/mutations/diff.py

@@ -1,6 +1,6 @@
 from typing import TYPE_CHECKING
 
-from graphene import Boolean, DateTime, InputObjectType, Mutation, String
+from graphene import Boolean, DateTime, Field, InputObjectType, Mutation, String
 from graphql import GraphQLResolveInfo
 
 from infrahub.core import registry
@@ -9,11 +9,12 @@ from infrahub.core.diff.model.path import NameTrackingId
 from infrahub.core.diff.models import RequestDiffUpdate
 from infrahub.core.diff.repository.repository import DiffRepository
 from infrahub.core.timestamp import Timestamp
-from infrahub.database import retry_db_transaction
 from infrahub.dependencies.registry import get_component_registry
 from infrahub.exceptions import ValidationError
 from infrahub.workflows.catalogue import DIFF_UPDATE
 
+from ..types.task import TaskInfo
+
 if TYPE_CHECKING:
     from ..initialization import GraphqlContext
 
@@ -23,25 +24,30 @@ class DiffUpdateInput(InputObjectType):
     name = String(required=False)
     from_time = DateTime(required=False)
     to_time = DateTime(required=False)
-    wait_for_completion = Boolean(required=False)
+    wait_for_completion = Boolean(required=False, deprecation_reason="Please use `wait_until_completion` instead")
 
 
 class DiffUpdateMutation(Mutation):
     class Arguments:
         data = DiffUpdateInput(required=True)
+        wait_until_completion = Boolean(required=False)
 
     ok = Boolean()
+    task = Field(TaskInfo, required=False)
 
     @classmethod
-    @retry_db_transaction(name="diff_update")
     async def mutate(
         cls,
         root: dict,  # pylint: disable=unused-argument
         info: GraphQLResolveInfo,
         data: DiffUpdateInput,
-    ) -> dict[str, bool]:
+        wait_until_completion: bool = False,
+    ) -> dict[str, bool | dict[str, str]]:
         context: GraphqlContext = info.context
 
+        if data.wait_for_completion is True:
+            wait_until_completion = True
+
         from_timestamp_str = DateTime.serialize(data.from_time) if data.from_time else None
         to_timestamp_str = DateTime.serialize(data.to_time) if data.to_time else None
         if (data.from_time or data.to_time) and not data.name:
@@ -53,11 +59,11 @@ class DiffUpdateMutation(Mutation):
         diff_repository = await component_registry.get_component(DiffRepository, db=context.db, branch=diff_branch)
 
         tracking_id = NameTrackingId(name=data.name)
-        existing_diffs_metatdatas = await diff_repository.get_roots_metadata(
+        existing_diffs_metadatas = await diff_repository.get_roots_metadata(
             diff_branch_names=[diff_branch.name], base_branch_names=[base_branch.name], tracking_id=tracking_id
         )
-        if existing_diffs_metatdatas:
-            metadata = existing_diffs_metatdatas[0]
+        if existing_diffs_metadatas:
+            metadata = existing_diffs_metadatas[0]
             from_time = Timestamp(from_timestamp_str) if from_timestamp_str else None
             to_time = Timestamp(to_timestamp_str) if to_timestamp_str else None
             branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
@@ -68,7 +74,7 @@ class DiffUpdateMutation(Mutation):
             if to_time and to_time < metadata.to_time:
                 raise ValidationError(f"to_time must be null or greater than or equal to {metadata.to_time}")
 
-        if data.wait_for_completion is True:
+        if wait_until_completion is True:
             diff_coordinator = await component_registry.get_component(
                 DiffCoordinator, db=context.db, branch=diff_branch
             )
@@ -89,6 +95,7 @@ class DiffUpdateMutation(Mutation):
             to_time=to_timestamp_str,
         )
         if context.service:
-            await context.service.workflow.submit_workflow(workflow=DIFF_UPDATE, parameters={"model": model})
+            workflow = await context.service.workflow.submit_workflow(workflow=DIFF_UPDATE, parameters={"model": model})
+            return {"ok": True, "task": {"id": str(workflow.id)}}
 
         return {"ok": True}

infrahub/graphql/mutations/resource_manager.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any
 
-from graphene import Boolean, Field, InputField, InputObjectType, Int, Mutation, String
+from graphene import Boolean, Field, InputField, InputObjectType, Int, List, Mutation, String
 from graphene.types.generic import GenericScalar
 from typing_extensions import Self
 
@@ -30,7 +30,7 @@ if TYPE_CHECKING:
 
 class IPPrefixPoolGetResourceInput(InputObjectType):
     id = InputField(String(required=False), description="ID of the pool to allocate from")
-    hfid = InputField(String(required=False), description="HFID of the pool to allocate from")
+    hfid = InputField(List(of_type=String, required=False), description="HFID of the pool to allocate from")
     identifier = InputField(String(required=False), description="Identifier for the allocated resource")
     prefix_length = InputField(Int(required=False), description="Size of the prefix to allocate")
     member_type = InputField(String(required=False), description="Type of members for the newly created prefix")
@@ -40,7 +40,7 @@ class IPPrefixPoolGetResourceInput(InputObjectType):
 
 class IPAddressPoolGetResourceInput(InputObjectType):
     id = InputField(String(required=False), description="ID of the pool to allocate from")
-    hfid = InputField(String(required=False), description="HFID of the pool to allocate from")
+    hfid = InputField(List(of_type=String, required=False), description="HFID of the pool to allocate from")
     identifier = InputField(String(required=False), description="Identifier for the allocated resource")
     prefix_length = InputField(
         Int(required=False), description="Size of the prefix mask to allocate on the new IP address"

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_sdk/batch.py

@@ -30,7 +30,7 @@ class BatchTaskSync:
         result = None
         try:
             result = self.task(*self.args, **self.kwargs)
-        except Exception as exc:  # pylint: disable=broad-exception-caught
+        except Exception as exc:
             if return_exceptions:
                 return self.node, exc
             raise exc
@@ -44,7 +44,7 @@ async def execute_batch_task_in_pool(
     async with semaphore:
         try:
             result = await task.task(*task.args, **task.kwargs)
-        except Exception as exc:  # pylint: disable=broad-exception-caught
+        except Exception as exc:
             if return_exceptions:
                 return (task.node, exc)
             raise exc

infrahub_sdk/client.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import copy
 import logging
+import time
 from collections.abc import Coroutine, MutableMapping
 from functools import wraps
 from time import sleep
@@ -38,6 +39,7 @@ from .exceptions import (
     NodeNotFoundError,
     ServerNotReachableError,
     ServerNotResponsiveError,
+    URLNotFoundError,
 )
 from .graphql import Mutation, Query
 from .node import (
@@ -878,7 +880,8 @@ class InfrahubClient(BaseClient):
 
         retry = True
         resp = None
-        while retry:
+        start_time = time.time()
+        while retry and time.time() - start_time < self.config.max_retry_duration:
             retry = self.retry_on_failure
             try:
                 resp = await self._post(url=url, payload=payload, headers=headers, timeout=timeout)
@@ -902,6 +905,8 @@ class InfrahubClient(BaseClient):
                     errors = response.get("errors", [])
                     messages = [error.get("message") for error in errors]
                     raise AuthenticationError(" | ".join(messages)) from exc
+                if exc.response.status_code == 404:
+                    raise URLNotFoundError(url=url)
 
         if not resp:
             raise Error("Unexpected situation, resp hasn't been initialized.")
@@ -1613,7 +1618,8 @@ class InfrahubClientSync(BaseClient):
 
         retry = True
         resp = None
-        while retry:
+        start_time = time.time()
+        while retry and time.time() - start_time < self.config.max_retry_duration:
             retry = self.retry_on_failure
             try:
                 resp = self._post(url=url, payload=payload, headers=headers, timeout=timeout)
@@ -1637,6 +1643,8 @@ class InfrahubClientSync(BaseClient):
                     errors = response.get("errors", [])
                     messages = [error.get("message") for error in errors]
                     raise AuthenticationError(" | ".join(messages)) from exc
+                if exc.response.status_code == 404:
+                    raise URLNotFoundError(url=url)
 
         if not resp:
             raise Error("Unexpected situation, resp hasn't been initialized.")

infrahub_sdk/config.py

@@ -56,6 +56,9 @@ class ConfigBase(BaseSettings):
     pagination_size: int = Field(default=50, description="Page size for queries to the server")
     retry_delay: int = Field(default=5, description="Number of seconds to wait until attempting a retry.")
     retry_on_failure: bool = Field(default=False, description="Retry operation in case of failure")
+    max_retry_duration: int = Field(
+        default=300, description="Maximum duration until we stop attempting to retry if enabled."
+    )
     schema_converge_timeout: int = Field(
         default=60, description="Number of seconds to wait for schema to have converged"
     )
@@ -113,7 +116,7 @@ class ConfigBase(BaseSettings):
 
     @model_validator(mode="after")
     def validate_proxy_config(self) -> Self:
-        if self.proxy and self.proxy_mounts.is_set:  # pylint: disable=no-member
+        if self.proxy and self.proxy_mounts.is_set:
             raise ValueError("'proxy' and 'proxy_mounts' are mutually exclusive")
         return self
 

infrahub_sdk/ctl/check.py

@@ -5,7 +5,7 @@ import sys
 from asyncio import run as aiorun
 from dataclasses import dataclass
 from pathlib import Path
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Optional
 
 import typer
 from rich.console import Console
@@ -50,8 +50,8 @@ def run(
     format_json: bool,
     list_available: bool,
     variables: dict[str, str],
-    name: str | None = None,
-    branch: str | None = None,
+    name: Optional[str] = None,
+    branch: Optional[str] = None,
 ) -> None:
     """Locate and execute all checks under the defined path."""
 
@@ -121,7 +121,7 @@ async def run_check(
     except QueryNotFoundError as exc:
         log.warning(f"{module_name}::{check}: unable to find query ({exc!s})")
         passed = False
-    except Exception as exc:  # pylint: disable=broad-exception-caught
+    except Exception as exc:
         log.warning(f"{module_name}::{check}: An error occurred during execution ({exc})")
         passed = False