infrahub-server 1.3.5__py3-none-any.whl → 1.4.0b0__py3-none-any.whl

infrahub/graphql/field_extractor.py

@@ -0,0 +1,69 @@
+from typing import Any
+
+from graphql import (
+    FieldNode,
+    FragmentSpreadNode,
+    GraphQLResolveInfo,
+    InlineFragmentNode,
+    SelectionSetNode,
+)
+
+
+class GraphQLFieldExtractor:
+    """Class to extract fields from a GraphQL selection set."""
+
+    def __init__(self, info: GraphQLResolveInfo):
+        self.info = info
+        self.fragments = info.fragments
+
+    def get_fields(self) -> dict[str, Any]:
+        """Extract fields from the GraphQL selection set."""
+        fields = self._extract_fields(selection_set=self.info.field_nodes[0].selection_set)
+        return fields or {}
+
+    def _extract_fields(self, selection_set: SelectionSetNode | None) -> dict[str, dict] | None:
+        """This function extract all the requested fields in a tree of Dict from a SelectionSetNode
+
+        The goal of this function is to limit the fields that we need to query from the backend.
+
+        Currently the function support Fields and InlineFragments but in a combined tree where the fragments are merged together
+        This implementation may seam counter intuitive but in the current implementation
+        it's better to have slightly more information at time passed to the query manager.
+
+        In the future we'll probably need to redesign how we read GraphQL queries to generate better Database query.
+        """
+
+        if not selection_set:
+            return None
+
+        fields: dict[str, dict | Any] = {}
+        for node in selection_set.selections:
+            sub_selection_set = getattr(node, "selection_set", None)
+            if isinstance(node, FieldNode):
+                value = self._extract_fields(sub_selection_set)
+                if node.name.value not in fields:
+                    fields[node.name.value] = value
+                elif isinstance(fields[node.name.value], dict) and isinstance(value, dict):
+                    fields[node.name.value].update(value)
+
+            elif isinstance(node, InlineFragmentNode):
+                for sub_node in node.selection_set.selections:
+                    sub_sub_selection_set = getattr(sub_node, "selection_set", None)
+                    value = self._extract_fields(sub_sub_selection_set)
+                    sub_node_name = getattr(sub_node, "name", "")
+                    sub_node_name_value = getattr(sub_node_name, "value", "")
+                    if sub_node_name_value not in fields:
+                        fields[sub_node_name_value] = self._extract_fields(sub_sub_selection_set)
+                    elif isinstance(fields[sub_node_name_value], dict) and isinstance(value, dict):
+                        fields[sub_node_name_value].update(value)
+            elif isinstance(node, FragmentSpreadNode):
+                if node.name.value in self.info.fragments:
+                    if fragment_fields := self._extract_fields(self.info.fragments[node.name.value].selection_set):
+                        fields.update(fragment_fields)
+
+        return fields
+
+
+def extract_graphql_fields(info: GraphQLResolveInfo) -> dict[str, Any]:
+    graphql_extractor = GraphQLFieldExtractor(info=info)
+    return graphql_extractor.get_fields()

infrahub/graphql/manager.py

@@ -40,6 +40,7 @@ from .mutations.resource_manager import (
     InfrahubNumberPoolMutation,
 )
 from .mutations.webhook import InfrahubWebhookMutation
+from .resolvers.ipam import ipam_paginated_list_resolver
 from .resolvers.resolver import (
     account_resolver,
     ancestors_resolver,
@@ -57,8 +58,8 @@ from .types import (
     InfrahubObject,
     PaginatedObjectPermission,
     RelatedIPAddressNodeInput,
+    RelatedIPPrefixNodeInput,
     RelatedNodeInput,
-    RelatedPrefixNodeInput,
 )
 from .types.attribute import BaseAttribute as BaseAttributeType
 from .types.attribute import TextAttributeType
@@ -343,14 +344,10 @@ class GraphQLSchemaManager:
 
     def _get_related_input_type(self, relationship: RelationshipSchema) -> type[RelatedNodeInput]:
         peer_schema = self.schema.get(name=relationship.peer, duplicate=False)
-        if (isinstance(peer_schema, NodeSchema) and peer_schema.is_ip_prefix()) or (
-            isinstance(peer_schema, GenericSchema) and relationship.peer == InfrahubKind.IPPREFIX
-        ):
-            return RelatedPrefixNodeInput
+        if peer_schema.is_ip_prefix:
+            return RelatedIPPrefixNodeInput
 
-        if (isinstance(peer_schema, NodeSchema) and peer_schema.is_ip_address()) or (
-            isinstance(peer_schema, GenericSchema) and relationship.peer == InfrahubKind.IPADDRESS
-        ):
+        if peer_schema.is_ip_address:
             return RelatedIPAddressNodeInput
 
         return RelatedNodeInput
@@ -495,7 +492,9 @@ class GraphQLSchemaManager:
 
             class_attrs[node_schema.kind] = graphene.Field(
                 node_type,
-                resolver=default_paginated_list_resolver,
+                resolver=default_paginated_list_resolver
+                if node_name not in [InfrahubKind.IPADDRESS, InfrahubKind.IPPREFIX]
+                else ipam_paginated_list_resolver,
                 required=True,
                 **node_filters,
             )
@@ -530,9 +529,9 @@ class GraphQLSchemaManager:
                 InfrahubKind.NODETRIGGERRELATIONSHIPMATCH: InfrahubTriggerRuleMatchMutation,
             }
 
-            if isinstance(node_schema, NodeSchema) and node_schema.is_ip_prefix():
+            if isinstance(node_schema, NodeSchema) and node_schema.is_ip_prefix:
                 base_class = InfrahubIPPrefixMutation
-            elif isinstance(node_schema, NodeSchema) and node_schema.is_ip_address():
+            elif isinstance(node_schema, NodeSchema) and node_schema.is_ip_address:
                 base_class = InfrahubIPAddressMutation
             else:
                 base_class = mutation_map.get(node_schema.kind, InfrahubMutation)
@@ -912,6 +911,11 @@ class GraphQLSchemaManager:
             filters.update(get_attribute_type().get_graphql_filters(name="any"))
             filters["partial_match"] = graphene.Boolean()
 
+            if schema.kind in [InfrahubKind.IPADDRESS, InfrahubKind.IPPREFIX]:
+                # This is only available for IPAM generics
+                filters["include_available"] = graphene.Boolean()
+                filters["kinds"] = graphene.List(graphene.NonNull(graphene.String))
+
         if not top_level:
             return filters
 

infrahub/graphql/mutations/account.py

@@ -2,7 +2,6 @@ from typing import TYPE_CHECKING, Any
 
 from graphene import Boolean, Field, InputField, InputObjectType, Mutation, String
 from graphql import GraphQLResolveInfo
-from infrahub_sdk.utils import extract_fields
 from infrahub_sdk.uuidt import UUIDT
 from typing_extensions import Self
 
@@ -14,6 +13,7 @@ from infrahub.core.protocols import CoreAccount, CoreNode, InternalAccountToken
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase, retry_db_transaction
 from infrahub.exceptions import NodeNotFoundError, PermissionDeniedError
+from infrahub.graphql.field_extractor import extract_graphql_fields
 
 from ..models import OrderModel
 from ..types import InfrahubObjectType
@@ -101,7 +101,7 @@ class AccountMixin:
         async with db.start_transaction() as dbt:
             await obj.save(db=dbt)
 
-        fields = await extract_fields(info.field_nodes[0].selection_set)
+        fields = extract_graphql_fields(info=info)
         return cls(object=await obj.to_graphql(db=db, fields=fields.get("object", {})), ok=True)  # type: ignore[call-arg]
 
     @classmethod

infrahub/graphql/mutations/action.py

@@ -46,11 +46,14 @@ class InfrahubTriggerRuleMutation(InfrahubMutationMixin, Mutation):
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
         db = database or graphql_context.db
         _validate_node_kind(data=data, db=db)
-        trigger_rule_definition, result = await super().mutate_create(info=info, data=data, branch=branch, database=db)
+        trigger_rule_definition, result = await super().mutate_create(
+            info=info, data=data, branch=branch, database=db, override_data=override_data
+        )
 
         return trigger_rule_definition, result
 
@@ -93,11 +96,14 @@ class InfrahubTriggerRuleMatchMutation(InfrahubMutationMixin, Mutation):
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,  # noqa: ARG003
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
 
         async with graphql_context.db.start_transaction() as dbt:
-            trigger_match, result = await super().mutate_create(info=info, data=data, branch=branch, database=dbt)
+            trigger_match, result = await super().mutate_create(
+                info=info, data=data, branch=branch, database=dbt, override_data=override_data
+            )
             trigger_match_model = cast(CoreNodeTriggerAttributeMatch | CoreNodeTriggerRelationshipMatch, trigger_match)
             node_trigger_rule = await trigger_match_model.trigger.get_peer(db=dbt, raise_on_error=True)
             node_trigger_rule_model = cast(CoreNodeTriggerRule, node_trigger_rule)

infrahub/graphql/mutations/artifact_definition.py

@@ -49,10 +49,13 @@ class InfrahubArtifactDefinitionMutation(InfrahubMutationMixin, Mutation):
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,  # noqa: ARG003
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
 
-        artifact_definition, result = await super().mutate_create(info=info, data=data, branch=branch)
+        artifact_definition, result = await super().mutate_create(
+            info=info, data=data, branch=branch, override_data=override_data
+        )
 
         if graphql_context.service:
             model = RequestArtifactDefinitionGenerate(

infrahub/graphql/mutations/branch.py

@@ -3,13 +3,14 @@ from __future__ import annotations
 from typing import TYPE_CHECKING
 
 from graphene import Boolean, Field, InputField, InputObjectType, Mutation, String
-from infrahub_sdk.utils import extract_fields, extract_fields_first_node
 from opentelemetry import trace
 from typing_extensions import Self
 
+from infrahub.branch.merge_mutation_checker import verify_branch_merge_mutation_allowed
 from infrahub.core.branch import Branch
 from infrahub.database import retry_db_transaction
 from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.field_extractor import extract_graphql_fields
 from infrahub.graphql.types.context import ContextInput
 from infrahub.log import get_logger
 from infrahub.workflows.catalogue import (
@@ -84,7 +85,7 @@ class BranchCreate(Mutation):
 
         # Retrieve created branch
         obj = await Branch.get_by_name(db=graphql_context.db, name=model.name)
-        fields = await extract_fields(info.field_nodes[0].selection_set)
+        fields = extract_graphql_fields(info=info)
         return cls(object=await obj.to_graphql(fields=fields.get("object", {})), ok=True, task=task)
 
 
@@ -202,7 +203,7 @@ class BranchRebase(Mutation):
             )
             task = {"id": workflow.id}
 
-        fields = await extract_fields_first_node(info=info)
+        fields = extract_graphql_fields(info=info)
         ok = True
 
         return cls(object=await obj.to_graphql(fields=fields.get("object", {})), ok=ok, task=task)
@@ -245,7 +246,7 @@ class BranchValidate(Mutation):
             )
             task = {"id": workflow.id}
 
-        fields = await extract_fields_first_node(info=info)
+        fields = extract_graphql_fields(info=info)
 
         return cls(object=await obj.to_graphql(fields=fields.get("object", {})), ok=ok, task=task)
 
@@ -274,6 +275,10 @@ class BranchMerge(Mutation):
         graphql_context: GraphqlContext = info.context
         await apply_external_context(graphql_context=graphql_context, context_input=context)
 
+        await verify_branch_merge_mutation_allowed(
+            db=graphql_context.db, account_session=graphql_context.active_account_session
+        )
+
         if wait_until_completion:
             await graphql_context.active_service.workflow.execute_workflow(
                 workflow=BRANCH_MERGE_MUTATION,
@@ -291,7 +296,7 @@ class BranchMerge(Mutation):
         # Pull the latest information about the branch from the database directly
         obj = await Branch.get_by_name(db=graphql_context.db, name=branch_name)
 
-        fields = await extract_fields(info.field_nodes[0].selection_set)
+        fields = extract_graphql_fields(info=info)
         ok = True
 
         return cls(object=await obj.to_graphql(fields=fields.get("object", {})), ok=ok, task=task)

infrahub/graphql/mutations/graphql_query.py

@@ -68,6 +68,7 @@ class InfrahubGraphQLQueryMutation(InfrahubMutationMixin, Mutation):
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,  # noqa: ARG003
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
 
@@ -75,7 +76,7 @@ class InfrahubGraphQLQueryMutation(InfrahubMutationMixin, Mutation):
             await cls.extract_query_info(info=info, data=data, branch=graphql_context.branch, db=graphql_context.db)
         )
 
-        obj, result = await super().mutate_create(info=info, data=data, branch=branch)
+        obj, result = await super().mutate_create(info=info, data=data, branch=branch, override_data=override_data)
 
         return obj, result
 

infrahub/graphql/mutations/main.py

@@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any
 
 from graphene import InputObjectType, Mutation
 from graphene.types.mutation import MutationOptions
-from infrahub_sdk.utils import extract_fields
 from typing_extensions import Self
 
 from infrahub import config, lock
@@ -28,6 +27,7 @@ from infrahub.dependencies.registry import get_component_registry
 from infrahub.events.generator import generate_node_mutation_events
 from infrahub.exceptions import HFIDViolatedError, InitializationError, NodeNotFoundError
 from infrahub.graphql.context import apply_external_context
+from infrahub.graphql.field_extractor import extract_graphql_fields
 from infrahub.lock import InfrahubMultiLock, build_object_lock_name
 from infrahub.log import get_log_data, get_logger
 
@@ -147,7 +147,9 @@ class InfrahubMutationMixin:
         return mutation
 
     @classmethod
-    async def _call_mutate_create_object(cls, data: InputObjectType, db: InfrahubDatabase, branch: Branch) -> Node:
+    async def _call_mutate_create_object(
+        cls, data: InputObjectType, db: InfrahubDatabase, branch: Branch, override_data: dict[str, Any] | None = None
+    ) -> Node:
         """
         Wrapper around mutate_create_object to potentially activate locking.
         """
@@ -157,9 +159,9 @@ class InfrahubMutationMixin:
         )
         if lock_names:
             async with InfrahubMultiLock(lock_registry=lock.registry, locks=lock_names):
-                return await cls.mutate_create_object(data=data, db=db, branch=branch)
+                return await cls.mutate_create_object(data=data, db=db, branch=branch, override_data=override_data)
 
-        return await cls.mutate_create_object(data=data, db=db, branch=branch)
+        return await cls.mutate_create_object(data=data, db=db, branch=branch, override_data=override_data)
 
     @classmethod
     async def mutate_create(
@@ -168,10 +170,11 @@ class InfrahubMutationMixin:
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         graphql_context: GraphqlContext = info.context
         db = database or graphql_context.db
-        obj = await cls._call_mutate_create_object(data=data, db=db, branch=branch)
+        obj = await cls._call_mutate_create_object(data=data, db=db, branch=branch, override_data=override_data)
         result = await cls.mutate_create_to_graphql(info=info, db=db, obj=obj)
         return obj, result
 
@@ -182,12 +185,15 @@ class InfrahubMutationMixin:
         data: InputObjectType,
         db: InfrahubDatabase,
         branch: Branch,
+        override_data: dict[str, Any] | None = None,
     ) -> Node:
         schema = cls._meta.active_schema
         if isinstance(schema, GenericSchema):
             raise ValueError(f"Node of generic schema `{schema.name=}` can not be instantiated.")
+        create_data = dict(data)
+        create_data.update(override_data or {})
         return await create_node(
-            data=dict(data),
+            data=create_data,
             db=db,
             branch=branch,
             schema=schema,
@@ -195,7 +201,7 @@ class InfrahubMutationMixin:
 
     @classmethod
     async def mutate_create_to_graphql(cls, info: GraphQLResolveInfo, db: InfrahubDatabase, obj: Node) -> Self:
-        fields = await extract_fields(info.field_nodes[0].selection_set)
+        fields = extract_graphql_fields(info=info)
         result: dict[str, Any] = {"ok": True}
         if "object" in fields:
             result["object"] = await obj.to_graphql(db=db, fields=fields.get("object", {}))
@@ -314,7 +320,7 @@ class InfrahubMutationMixin:
         info: GraphQLResolveInfo,
         obj: Node,
     ) -> Self:
-        fields_object = await extract_fields(info.field_nodes[0].selection_set)
+        fields_object = extract_graphql_fields(info=info)
         fields_object = fields_object.get("object", {})
         result: dict[str, Any] = {"ok": True}
         if fields_object:

infrahub/graphql/mutations/menu.py

@@ -54,10 +54,11 @@ class InfrahubCoreMenuMutation(InfrahubMutationMixin, Mutation):
         data: InputObjectType,
         branch: Branch,
         database: InfrahubDatabase | None = None,  # noqa: ARG003
+        override_data: dict[str, Any] | None = None,
     ) -> tuple[Node, Self]:
         validate_namespace(data=data)
 
-        obj, result = await super().mutate_create(info=info, data=data, branch=branch)
+        obj, result = await super().mutate_create(info=info, data=data, branch=branch, override_data=override_data)
 
         return obj, result