infrahub-server 1.4.12__py3-none-any.whl → 1.5.0__py3-none-any.whl

infrahub/events/branch_action.py

@@ -109,7 +109,7 @@ class BranchRebasedEvent(InfrahubEvent):
 
     event_name: ClassVar[str] = f"{EVENT_NAMESPACE}.branch.rebased"
 
-    branch_id: str = Field(..., description="The ID of the mutated node")
+    branch_id: str = Field(..., description="The ID of the branch")
     branch_name: str = Field(..., description="The name of the branch")
 
     def get_resource(self) -> dict[str, str]:
@@ -128,3 +128,29 @@ class BranchRebasedEvent(InfrahubEvent):
             RefreshRegistryRebasedBranch(branch=self.branch_name),
         ]
         return events
+
+
+class BranchMigratedEvent(InfrahubEvent):
+    """Event generated when a branch has been migrated"""
+
+    event_name: ClassVar[str] = f"{EVENT_NAMESPACE}.branch.migrated"
+
+    branch_id: str = Field(..., description="The ID of the branch")
+    branch_name: str = Field(..., description="The name of the branch")
+
+    def get_resource(self) -> dict[str, str]:
+        return {
+            "prefect.resource.id": f"infrahub.branch.{self.branch_name}",
+            "infrahub.branch.id": self.branch_id,
+            "infrahub.branch.name": self.branch_name,
+        }
+
+    def get_messages(self) -> list[InfrahubMessage]:
+        events: list[InfrahubMessage] = [
+            # EventBranchMigrated(
+            #     branch=self.branch,
+            #     meta=self.get_message_meta(),
+            # ),
+            RefreshRegistryRebasedBranch(branch=self.branch_name),
+        ]
+        return events

infrahub/events/group_action.py

@@ -22,7 +22,7 @@ class GroupMutatedEvent(InfrahubEvent):
     def get_related(self) -> list[dict[str, str]]:
         related = super().get_related()
 
-        if self.kind in [InfrahubKind.GENERATORGROUP, InfrahubKind.GRAPHQLQUERYGROUP]:
+        if self.kind in [InfrahubKind.GENERATORGROUP, InfrahubKind.GENERATORAWAREGROUP, InfrahubKind.GRAPHQLQUERYGROUP]:
             # Temporary workaround to avoid too large payloads for the related field
             return related
 

infrahub/events/node_action.py

@@ -24,7 +24,7 @@ class NodeMutatedEvent(InfrahubEvent):
 
     def get_related(self) -> list[dict[str, str]]:
         related = super().get_related()
-        if self.kind in [InfrahubKind.GENERATORGROUP, InfrahubKind.GRAPHQLQUERYGROUP]:
+        if self.kind in [InfrahubKind.GENERATORGROUP, InfrahubKind.GENERATORAWAREGROUP, InfrahubKind.GRAPHQLQUERYGROUP]:
             # Temporary workaround to avoid too large payloads for the related field
             return related
 

infrahub/generators/constants.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class GeneratorDefinitionRunSource(Enum):
+    PROPOSED_CHANGE = "proposed_change"
+    MERGE = "merge"
+    UNKNOWN = "unknown"

infrahub/generators/models.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, ConfigDict, Field
 class RequestGeneratorRun(BaseModel):
     """Runs a generator."""
 
-    generator_definition: ProposedChangeGeneratorDefinition = Field(..., description="The Generator definition")
+    generator_definition: GeneratorDefinitionModel = Field(..., description="The Generator definition")
     generator_instance: str | None = Field(
         default=None, description="The id of the generator instance if it previously existed"
     )
@@ -31,14 +31,40 @@ class RequestGeneratorDefinitionRun(BaseModel):
     target_members: list[str] = Field(default_factory=list, description="List of targets to run the generator for")
 
 
-class ProposedChangeGeneratorDefinition(BaseModel):
-    definition_id: str
-    definition_name: str
-    query_name: str
-    convert_query_response: bool
-    query_models: list[str]
-    repository_id: str
-    class_name: str
-    file_path: str
-    parameters: dict
-    group_id: str
+class GeneratorDefinitionModel(BaseModel):
+    definition_id: str = Field(..., description="The id of the generator definition.")
+    definition_name: str = Field(..., description="The name of the generator definition.")
+    query_name: str = Field(..., description="The name of the query to use when collecting data.")
+    convert_query_response: bool = Field(
+        ...,
+        description="Decide if the generator should convert the result of the GraphQL query to SDK InfrahubNode objects.",
+    )
+    class_name: str = Field(..., description="The name of the generator class to run.")
+    file_path: str = Field(..., description="The file path of the generator in the repository.")
+    group_id: str = Field(..., description="The group to target when running this generator")
+    parameters: dict = Field(..., description="The input parameters required to run this check")
+
+    execute_in_proposed_change: bool = Field(
+        ..., description="Indicates if the generator should execute in a proposed change."
+    )
+    execute_after_merge: bool = Field(..., description="Indicates if the generator should execute after a merge.")
+
+    @classmethod
+    def from_pc_generator_definition(cls, model: ProposedChangeGeneratorDefinition) -> GeneratorDefinitionModel:
+        return GeneratorDefinitionModel(
+            definition_id=model.definition_id,
+            definition_name=model.definition_name,
+            query_name=model.query_name,
+            convert_query_response=model.convert_query_response,
+            class_name=model.class_name,
+            file_path=model.file_path,
+            group_id=model.group_id,
+            parameters=model.parameters,
+            execute_in_proposed_change=model.execute_in_proposed_change,
+            execute_after_merge=model.execute_after_merge,
+        )
+
+
+class ProposedChangeGeneratorDefinition(GeneratorDefinitionModel):
+    query_models: list[str] = Field(..., description="The models to use when collecting data.")
+    repository_id: str = Field(..., description="The id of the repository.")

infrahub/generators/tasks.py

@@ -14,7 +14,9 @@ from prefect.states import Completed, Failed
 from infrahub import lock
 from infrahub.context import InfrahubContext  # noqa: TC001 needed for prefect flow
 from infrahub.core.constants import GeneratorInstanceStatus, InfrahubKind
+from infrahub.generators.constants import GeneratorDefinitionRunSource
 from infrahub.generators.models import (
+    GeneratorDefinitionModel,
     ProposedChangeGeneratorDefinition,
     RequestGeneratorDefinitionRun,
     RequestGeneratorRun,
@@ -53,9 +55,12 @@ async def run_generator(model: RequestGeneratorRun) -> None:
         name=model.generator_definition.definition_name,
         class_name=model.generator_definition.class_name,
         file_path=model.generator_definition.file_path,
+        parameters=model.generator_definition.parameters,
         query=model.generator_definition.query_name,
         targets=model.generator_definition.group_id,
         convert_query_response=model.generator_definition.convert_query_response,
+        execute_in_proposed_change=model.generator_definition.execute_in_proposed_change,
+        execute_after_merge=model.generator_definition.execute_after_merge,
     )
 
     commit_worktree = repository.get_commit_worktree(commit=model.commit)
@@ -79,6 +84,8 @@ async def run_generator(model: RequestGeneratorRun) -> None:
             params=model.variables,
             generator_instance=generator_instance.id,
             convert_query_response=generator_definition.convert_query_response,
+            execute_in_proposed_change=generator_definition.execute_in_proposed_change,
+            execute_after_merge=generator_definition.execute_after_merge,
             infrahub_node=InfrahubNode,
         )
         await generator.run(identifier=generator_definition.name)
@@ -128,28 +135,39 @@ async def _define_instance(model: RequestGeneratorRun, client: InfrahubClient) -
 
 
 @flow(name="generator-definition-run", flow_run_name="Run all generators")
-async def run_generator_definition(branch: str, context: InfrahubContext) -> None:
+async def run_generator_definition(
+    branch: str, context: InfrahubContext, source: GeneratorDefinitionRunSource = GeneratorDefinitionRunSource.UNKNOWN
+) -> None:
     await add_tags(branches=[branch])
 
     generators = await get_client().filters(
         kind=InfrahubKind.GENERATORDEFINITION, prefetch_relationships=True, populate_store=True, branch=branch
     )
 
-    generator_definitions = [
-        ProposedChangeGeneratorDefinition(
-            definition_id=generator.id,
-            definition_name=generator.name.value,
-            class_name=generator.class_name.value,
-            file_path=generator.file_path.value,
-            query_name=generator.query.peer.name.value,
-            query_models=generator.query.peer.models.value,
-            repository_id=generator.repository.peer.id,
-            parameters=generator.parameters.value,
-            group_id=generator.targets.peer.id,
-            convert_query_response=generator.convert_query_response.value,
+    generator_definitions: list[ProposedChangeGeneratorDefinition] = []
+
+    for generator in generators:
+        if (
+            source == GeneratorDefinitionRunSource.PROPOSED_CHANGE and not generator.execute_in_proposed_change.value
+        ) or (source == GeneratorDefinitionRunSource.MERGE and not generator.execute_after_merge.value):
+            continue
+
+        generator_definitions.append(
+            ProposedChangeGeneratorDefinition(
+                definition_id=generator.id,
+                definition_name=generator.name.value,
+                class_name=generator.class_name.value,
+                file_path=generator.file_path.value,
+                query_name=generator.query.peer.name.value,
+                query_models=generator.query.peer.models.value,
+                repository_id=generator.repository.peer.id,
+                parameters=generator.parameters.value,
+                group_id=generator.targets.peer.id,
+                convert_query_response=generator.convert_query_response.value,
+                execute_in_proposed_change=generator.execute_in_proposed_change.value,
+                execute_after_merge=generator.execute_after_merge.value,
+            )
         )
-        for generator in generators
-    ]
 
     for generator_definition in generator_definitions:
         model = RequestGeneratorDefinitionRun(branch=branch, generator_definition=generator_definition)
@@ -209,7 +227,7 @@ async def request_generator_definition_run(
 
         generator_instance = instance_by_member.get(member.id)
         request_generator_run_model = RequestGeneratorRun(
-            generator_definition=model.generator_definition,
+            generator_definition=GeneratorDefinitionModel.from_pc_generator_definition(model.generator_definition),
             commit=repository.commit.value,
             generator_instance=generator_instance,
             repository_id=repository.id,

infrahub/git/base.py

@@ -7,7 +7,7 @@ from typing import TYPE_CHECKING, NoReturn
 from uuid import UUID  # noqa: TC003
 
 import git
-from git import Blob, Repo
+from git import BadName, Blob, Repo
 from git.exc import GitCommandError, InvalidGitRepositoryError
 from git.refs.remote import RemoteReference
 from infrahub_sdk import InfrahubClient  # noqa: TC002
@@ -67,6 +67,15 @@ class RepoFileInformation(BaseModel):
     """Extension of the file Example: py """
 
 
+class RepoChangedFiles(BaseModel):
+    added: list[str] = Field(default_factory=list)
+    copied: list[tuple[str, str]] = Field(default_factory=list)
+    deleted: list[str] = Field(default_factory=list)
+    renamed: list[tuple[str, str]] = Field(default_factory=list)
+    modified: list[str] = Field(default_factory=list)
+    type_changed: list[tuple[str, str]] = Field(default_factory=list)
+
+
 def extract_repo_file_information(
     full_filename: Path, repo_directory: Path, worktree_directory: Path | None = None
 ) -> RepoFileInformation:
@@ -932,7 +941,10 @@ class InfrahubRepositoryBase(BaseModel, ABC):
     def _raise_enriched_error_static(
         error: GitCommandError, name: str, location: str, branch_name: str | None = None
     ) -> NoReturn:
-        if "Repository not found" in error.stderr or "does not appear to be a git" in error.stderr:
+        if any(
+            err in error.stderr
+            for err in ("Repository not found", "does not appear to be a git", "Failed to connect to")
+        ):
             raise RepositoryConnectionError(identifier=name) from error
 
         if "error: pathspec" in error.stderr:
@@ -970,3 +982,31 @@ class InfrahubRepositoryBase(BaseModel, ABC):
         if branch_name == self.default_branch and branch_name != registry.default_branch:
             return registry.default_branch
         return branch_name
+
+    def get_changed_files(self, first_commit: str, second_commit: str | None = None) -> RepoChangedFiles:
+        """Return the changes between two commits in this repo."""
+        changes = RepoChangedFiles()
+        repo = self.get_git_repo_main()
+
+        try:
+            commit_a = repo.commit(first_commit)
+            commit_b = repo.commit(second_commit) if second_commit else repo.head.commit
+        except BadName as exc:
+            raise CommitNotFoundError(identifier=str(self.id), commit=exc.args[0]) from exc
+
+        for diff in commit_a.diff(commit_b):
+            match diff.change_type:
+                case "A":
+                    changes.added.append(diff.b_path)
+                case "C":
+                    changes.copied.append((diff.a_path, diff.b_path))
+                case "D":
+                    changes.deleted.append(diff.a_path)
+                case "R":
+                    changes.renamed.append((diff.a_path, diff.b_path))
+                case "M":
+                    changes.modified.append(diff.b_path)
+                case "T":
+                    changes.type_changed.append((diff.a_path, diff.b_path))
+
+        return changes

infrahub/git/integrator.py

@@ -226,7 +226,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             )
         )
 
-    @task(name="import-jinja2-tansforms", task_run_name="Import Jinja2 transform", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-jinja2-transforms", task_run_name="Import Jinja2 transform", cache_policy=NONE)
     async def import_jinja2_transforms(
         self,
         branch_name: str,
@@ -331,7 +331,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
 
         await existing_transform.save()
 
-    @task(name="import-artifact-definitions", task_run_name="Import Artifact Definitions", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-artifact-definitions", task_run_name="Import Artifact Definitions", cache_policy=NONE)
     async def import_artifact_definitions(
         self,
         branch_name: str,
@@ -432,7 +432,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
 
         await existing_artifact_definition.save()
 
-    @task(name="repository-get-config", task_run_name="get repository config", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="repository-get-config", task_run_name="get repository config", cache_policy=NONE)
     async def get_repository_config(self, branch_name: str, commit: str) -> InfrahubRepositoryConfig | None:
         branch_wt = self.get_worktree(identifier=commit or branch_name)
         log = get_run_logger()
@@ -469,7 +469,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             log.error(f"Unable to load the configuration file {config_file_name}, the format is not valid  : {exc}")
             return None
 
-    @task(name="import-schema-files", task_run_name="Import schema files", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-schema-files", task_run_name="Import schema files", cache_policy=NONE)
     async def import_schema_files(self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig) -> None:
         log = get_run_logger()
         branch_wt = self.get_worktree(identifier=commit or branch_name)
@@ -541,7 +541,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
         for schema_file in schemas_data:
             log.info(f"schema '{schema_file.identifier}' loaded successfully!")
 
-    @task(name="import-graphql-queries", task_run_name="Import GraphQL Queries", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-graphql-queries", task_run_name="Import GraphQL Queries", cache_policy=NONE)
     async def import_all_graphql_query(
         self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
     ) -> None:
@@ -599,7 +599,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
         await obj.save()
         return obj
 
-    @task(name="import-python-check-definitions", task_run_name="Import Python Check Definitions", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-python-check-definitions", task_run_name="Import Python Check Definitions", cache_policy=NONE)
     async def import_python_check_definitions(
         self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
     ) -> None:
@@ -670,7 +670,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             log.info(f"CheckDefinition '{check_name!r}' not found locally, deleting")
             await check_definition_in_graph[check_name].delete()
 
-    @task(name="import-generator-definitions", task_run_name="Import Generator Definitions", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-generator-definitions", task_run_name="Import Generator Definitions", cache_policy=NONE)
     async def import_generator_definitions(
         self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
     ) -> None:
@@ -756,11 +756,13 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             or existing_generator.parameters.value != generator.parameters
             or existing_generator.convert_query_response.value != generator.convert_query_response
             or existing_generator.targets.id != generator.targets
+            or existing_generator.execute_in_proposed_change.value != generator.execute_in_proposed_change
+            or existing_generator.execute_after_merge.value != generator.execute_after_merge
         ):
             return True
         return False
 
-    @task(name="import-python-transforms", task_run_name="Import Python Transforms", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-python-transforms", task_run_name="Import Python Transforms", cache_policy=NONE)
     async def import_python_transforms(
         self, branch_name: str, commit: str, config_file: InfrahubRepositoryConfig
     ) -> None:
@@ -885,7 +887,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
                 file_type=file_type,
             )
 
-    @task(name="import-objects", task_run_name="Import Objects", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="import-objects", task_run_name="Import Objects", cache_policy=NONE)
     async def import_objects(
         self,
         branch_name: str,
@@ -905,7 +907,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             object_type=RepositoryObjects.MENU,
         )
 
-    @task(name="check-definition-get", task_run_name="Get Check Definition", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="check-definition-get", task_run_name="Get Check Definition", cache_policy=NONE)
     async def get_check_definition(
         self,
         branch_name: str,
@@ -945,7 +947,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
             raise
         return checks
 
-    @task(name="python-transform-get", task_run_name="Get Python Transform", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="python-transform-get", task_run_name="Get Python Transform", cache_policy=NONE)
     async def get_python_transforms(
         self, branch_name: str, module: types.ModuleType, file_path: str, transform: InfrahubPythonTransformConfig
     ) -> list[TransformPythonInformation]:
@@ -1023,6 +1025,12 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
         if existing_generator.targets.id != generator.targets:
             existing_generator.targets = {"id": generator.targets, "source": str(self.id), "is_protected": True}
 
+        if existing_generator.execute_in_proposed_change.value != generator.execute_in_proposed_change:
+            existing_generator.execute_in_proposed_change.value = generator.execute_in_proposed_change
+
+        if existing_generator.execute_after_merge.value != generator.execute_after_merge:
+            existing_generator.execute_after_merge.value = generator.execute_after_merge
+
         await existing_generator.save()
 
     async def create_python_check_definition(
@@ -1152,7 +1160,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
         await self.import_python_transforms(branch_name=branch_name, commit=commit, config_file=config_file)  # type: ignore[misc]
         await self.import_generator_definitions(branch_name=branch_name, commit=commit, config_file=config_file)  # type: ignore[misc]
 
-    @task(name="jinja2-template-render", task_run_name="Render Jinja2 template", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="jinja2-template-render", task_run_name="Render Jinja2 template", cache_policy=NONE)
     async def render_jinja2_template(self, commit: str, location: str, data: dict) -> str:
         log = get_run_logger()
         commit_worktree = self.get_commit_worktree(commit=commit)
@@ -1168,7 +1176,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
                 repository_name=self.name, commit=commit, location=location, message=exc.message
             ) from exc
 
-    @task(name="python-check-execute", task_run_name="Execute Python Check", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="python-check-execute", task_run_name="Execute Python Check", cache_policy=NONE)
     async def execute_python_check(
         self,
         branch_name: str,
@@ -1227,7 +1235,7 @@ class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
                 repository_name=self.name, class_name=class_name, commit=commit, location=location, message=str(exc)
             ) from exc
 
-    @task(name="python-transform-execute", task_run_name="Execute Python Transform", cache_policy=NONE)  # type: ignore[arg-type]
+    @task(name="python-transform-execute", task_run_name="Execute Python Transform", cache_policy=NONE)
     async def execute_python_transform(
         self,
         branch_name: str,

infrahub/git/tasks.py

@@ -1,3 +1,5 @@
+from typing import Any
+
 from infrahub_sdk import InfrahubClient
 from infrahub_sdk.protocols import (
     CoreArtifact,
@@ -14,7 +16,12 @@ from prefect.logging import get_run_logger
 
 from infrahub import lock
 from infrahub.context import InfrahubContext
-from infrahub.core.constants import InfrahubKind, RepositoryInternalStatus, ValidatorConclusion
+from infrahub.core.constants import (
+    InfrahubKind,
+    RepositoryInternalStatus,
+    RepositoryOperationalStatus,
+    ValidatorConclusion,
+)
 from infrahub.core.manager import NodeManager
 from infrahub.core.registry import registry
 from infrahub.exceptions import CheckError, RepositoryError
@@ -152,6 +159,39 @@ async def create_branch(branch: str, branch_id: str) -> None:
         pass
 
 
+@flow(name="sync-git-repo-with-origin", flow_run_name="Sync git repo with origin")
+async def sync_git_repo_with_origin_and_tag_on_failure(
+    client: InfrahubClient,
+    repository_id: str,
+    repository_name: str,
+    repository_location: str,
+    internal_status: str,
+    default_branch_name: str,
+    operational_status: str,
+    staging_branch: str | None = None,
+    infrahub_branch: str | None = None,
+) -> None:
+    repo = await InfrahubRepository.init(
+        id=repository_id,
+        name=repository_name,
+        location=repository_location,
+        client=client,
+        internal_status=internal_status,
+        default_branch_name=default_branch_name,
+    )
+
+    try:
+        await repo.sync(staging_branch=staging_branch)
+    except RepositoryError:
+        if operational_status == RepositoryOperationalStatus.ONLINE.value:
+            params: dict[str, Any] = {
+                "branches": [infrahub_branch] if infrahub_branch else [],
+                "nodes": [str(repository_id)],
+            }
+            await add_tags(**params)
+        raise
+
+
 @flow(name="git_repositories_sync", flow_run_name="Sync Git Repositories")
 async def sync_remote_repositories() -> None:
     log = get_run_logger()
@@ -204,7 +244,17 @@ async def sync_remote_repositories() -> None:
                     continue
 
             try:
-                await repo.sync(staging_branch=staging_branch)
+                await sync_git_repo_with_origin_and_tag_on_failure(
+                    client=client,
+                    repository_id=repository_data.repository.id,
+                    repository_name=repository_data.repository.name.value,
+                    repository_location=repository_data.repository.location.value,
+                    internal_status=active_internal_status,
+                    default_branch_name=repository_data.repository.default_branch.value,
+                    operational_status=repository_data.repository.operational_status.value,
+                    staging_branch=staging_branch,
+                    infrahub_branch=infrahub_branch,
+                )
                 # Tell workers to fetch to stay in sync
                 message = messages.RefreshGitFetch(
                     meta=Meta(initiator_id=WORKER_IDENTITY, request_id=get_log_data().get("request_id", "")),

infrahub/graphql/analyzer.py

@@ -312,6 +312,13 @@ class GraphQLQueryReport:
         return []
 
     def required_argument(self, argument: GraphQLArgument) -> bool:
+        if argument.name == "ids" and argument.kind == "list_value":
+            for variable in self.variables:
+                if f"['${variable.name}']" == argument.as_variable_name and variable.required:
+                    return True
+
+            return False
+
         if not argument.is_variable:
             # If the argument isn't a variable it would have been
             # statically defined in the input and as such required
@@ -364,6 +371,8 @@ class GraphQLQueryReport:
                     if [[argument.name]] == query.infrahub_model.uniqueness_constraints:
                         if self.required_argument(argument=argument):
                             targets_single_query = True
+                    elif argument.name == "ids" and self.required_argument(argument=argument):
+                        targets_single_query = True
 
             if not targets_single_query:
                 return False

infrahub/graphql/api/dependencies.py

@@ -1,5 +1,3 @@
-from typing import Any
-
 from infrahub import config
 
 from ..app import InfrahubGraphQLApp
@@ -36,5 +34,5 @@ def build_graphql_query_permission_checker() -> GraphQLQueryPermissionChecker:
     )
 
 
-def build_graphql_app(**kwargs: Any) -> InfrahubGraphQLApp:
-    return InfrahubGraphQLApp(build_graphql_query_permission_checker(), **kwargs)
+def build_graphql_app() -> InfrahubGraphQLApp:
+    return InfrahubGraphQLApp(build_graphql_query_permission_checker())

infrahub/graphql/api/endpoints.py

@@ -2,14 +2,15 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Query
 from fastapi.responses import PlainTextResponse
-from graphql import print_schema
+from graphql import parse, print_ast, print_schema
 from starlette.routing import Route, WebSocketRoute
 
 from infrahub.api.dependencies import get_branch_dep, get_current_user
 from infrahub.core import registry
-from infrahub.graphql.manager import GraphQLSchemaManager
+from infrahub.graphql.registry import registry as graphql_registry
+from infrahub.graphql.schema_sort import sort_schema_ast
 
 from .dependencies import build_graphql_app
 
@@ -27,11 +28,20 @@ router.routes.append(WebSocketRoute(path="/graphql", endpoint=graphql_app))
 router.routes.append(WebSocketRoute(path="/graphql/{branch_name:str}", endpoint=graphql_app))
 
 
-@router.get("/schema.graphql", include_in_schema=False)
+@router.get("/schema.graphql")
 async def get_graphql_schema(
-    branch: Branch = Depends(get_branch_dep), _: AccountSession = Depends(get_current_user)
+    branch: Branch = Depends(get_branch_dep),
+    _: AccountSession = Depends(get_current_user),
+    sort_schema: bool = Query(default=False, alias="sorted", description="Whether to sort the schema alphabetically."),
 ) -> PlainTextResponse:
     schema_branch = registry.schema.get_schema_branch(name=branch.name)
-    gqlm = GraphQLSchemaManager.get_manager_for_branch(branch=branch, schema_branch=schema_branch)
+    gqlm = graphql_registry.get_manager_for_branch(branch=branch, schema_branch=schema_branch)
     graphql_schema = gqlm.get_graphql_schema()
+
+    if sort_schema:
+        schema_str = print_schema(graphql_schema)
+        schema_ast = parse(schema_str)
+        sorted_schema_ast = sort_schema_ast(schema_ast)
+        return PlainTextResponse(content=print_ast(sorted_schema_ast))
+
     return PlainTextResponse(content=print_schema(graphql_schema))

infrahub/graphql/app.py

@@ -23,7 +23,6 @@ from graphql import (
     ExecutionResult,
     GraphQLError,
     GraphQLFormattedError,
-    Middleware,
     OperationType,
     graphql,
     parse,
@@ -59,6 +58,7 @@ from .metrics import (
     GRAPHQL_RESPONSE_SIZE_METRICS,
     GRAPHQL_TOP_LEVEL_QUERIES_METRICS,
 )
+from .middleware import raise_on_mutation_on_branch_needing_rebase
 
 if TYPE_CHECKING:
     import graphene
@@ -99,7 +99,6 @@ class InfrahubGraphQLApp:
         *,
         on_get: Callable[[Request], Response | Awaitable[Response]] | None = None,
         root_value: RootValue = None,
-        middleware: Middleware | None = None,
         error_formatter: Callable[[GraphQLError], GraphQLFormattedError] = format_error,
         execution_context_class: type[ExecutionContext] | None = None,
     ) -> None:
@@ -107,7 +106,6 @@ class InfrahubGraphQLApp:
         self.on_get = on_get
         self.root_value = root_value
         self.error_formatter = error_formatter
-        self.middleware = middleware
         self.execution_context_class = execution_context_class
         self.logger = get_logger(name="infrahub.graphql")
         self.permission_checker = permission_checker
@@ -259,7 +257,7 @@ class InfrahubGraphQLApp:
                     source=query,
                     context_value=graphql_params.context,
                     root_value=self.root_value,
-                    middleware=self.middleware,
+                    middleware=[raise_on_mutation_on_branch_needing_rebase],
                     variable_values=variable_values,
                     operation_name=operation_name,
                     execution_context_class=self.execution_context_class,

infrahub/graphql/initialization.py

@@ -9,12 +9,11 @@ from infrahub.context import InfrahubContext
 from infrahub.core import registry
 from infrahub.core.timestamp import Timestamp
 from infrahub.exceptions import InitializationError
+from infrahub.graphql.registry import registry as graphql_registry
 from infrahub.graphql.resolvers.many_relationship import ManyRelationshipResolver
 from infrahub.graphql.resolvers.single_relationship import SingleRelationshipResolver
 from infrahub.permissions import PermissionManager
 
-from .manager import GraphQLSchemaManager
-
 if TYPE_CHECKING:
     from graphql import GraphQLSchema
     from starlette.requests import HTTPConnection
@@ -91,7 +90,7 @@ async def prepare_graphql_params(
 ) -> GraphqlParams:
     branch = registry.get_branch_from_registry(branch=branch)
     schema_branch = registry.schema.get_schema_branch(name=branch.name)
-    gqlm = GraphQLSchemaManager.get_manager_for_branch(branch=branch, schema_branch=schema_branch)
+    gqlm = graphql_registry.get_manager_for_branch(branch=branch, schema_branch=schema_branch)
     gql_schema = gqlm.get_graphql_schema(
         include_query=include_query,
         include_mutation=include_mutation,