infrahub-server 1.5.0b2__py3-none-any.whl → 1.5.1__py3-none-any.whl

infrahub/proposed_change/tasks.py

@@ -11,6 +11,7 @@ import pytest
 from infrahub_sdk.exceptions import ModuleImportError, NodeNotFoundError, URLNotFoundError
 from infrahub_sdk.node import InfrahubNode
 from infrahub_sdk.protocols import (
+    CoreArtifactDefinition,
     CoreArtifactValidator,
     CoreGeneratorDefinition,
     CoreGeneratorValidator,
@@ -44,7 +45,7 @@ from infrahub.core.diff.model.diff import DiffElementType, SchemaConflict
 from infrahub.core.diff.model.path import NodeDiffFieldSummary
 from infrahub.core.integrity.object_conflict.conflict_recorder import ObjectConflictValidatorRecorder
 from infrahub.core.manager import NodeManager
-from infrahub.core.protocols import CoreArtifactDefinition, CoreDataCheck, CoreValidator
+from infrahub.core.protocols import CoreDataCheck, CoreValidator
 from infrahub.core.protocols import CoreProposedChange as InternalCoreProposedChange
 from infrahub.core.timestamp import Timestamp
 from infrahub.core.validators.checks_runner import run_checks_and_update_validator
@@ -59,6 +60,8 @@ from infrahub.git.base import extract_repo_file_information
 from infrahub.git.models import TriggerRepositoryInternalChecks, TriggerRepositoryUserChecks
 from infrahub.git.repository import InfrahubRepository, get_initialized_repo
 from infrahub.git.utils import fetch_artifact_definition_targets, fetch_proposed_change_generator_definition_targets
+from infrahub.graphql.analyzer import InfrahubGraphQLQueryAnalyzer
+from infrahub.graphql.initialization import prepare_graphql_params
 from infrahub.log import get_logger
 from infrahub.message_bus.types import (
     ProposedChangeArtifactDefinition,
@@ -672,6 +675,27 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, c
     repository = model.branch_diff.get_repository(repository_id=model.artifact_definition.repository_id)
     impacted_artifacts = model.branch_diff.get_subscribers_ids(kind=InfrahubKind.ARTIFACT)
 
+    source_schema_branch = registry.schema.get_schema_branch(name=model.source_branch)
+    source_branch = registry.get_branch_from_registry(branch=model.source_branch)
+
+    graphql_params = await prepare_graphql_params(db=await get_database(), branch=model.source_branch)
+    query_analyzer = InfrahubGraphQLQueryAnalyzer(
+        query=model.artifact_definition.query_payload,
+        branch=source_branch,
+        schema_branch=source_schema_branch,
+        schema=graphql_params.schema,
+    )
+
+    only_has_unique_targets = query_analyzer.query_report.only_has_unique_targets
+    if not only_has_unique_targets:
+        log.warning(
+            f"Artifact definition {artifact_definition.name.value} query does not guarantee unique targets. All targets will be processed."
+        )
+
+    managed_branch = model.source_branch_sync_with_git and model.branch_diff.has_file_modifications
+    if managed_branch:
+        log.info("Source branch is synced with Git repositories with updates, all artifacts will be processed")
+
     checks = []
 
     for relationship in group.members.peers:
@@ -679,8 +703,9 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, c
         artifact_id = artifacts_by_member.get(member.id)
         if _should_render_artifact(
             artifact_id=artifact_id,
-            managed_branch=model.source_branch_sync_with_git,
+            managed_branch=managed_branch,
             impacted_artifacts=impacted_artifacts,
+            only_has_unique_targets=only_has_unique_targets,
         ):
             log.info(f"Trigger Artifact processing for {member.display_label}")
 
@@ -726,21 +751,26 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, c
     )
 
 
-def _should_render_artifact(artifact_id: str | None, managed_branch: bool, impacted_artifacts: list[str]) -> bool:  # noqa: ARG001
+def _should_render_artifact(
+    artifact_id: str | None,
+    managed_branch: bool,
+    impacted_artifacts: list[str],
+    only_has_unique_targets: bool,
+) -> bool:
     """Returns a boolean to indicate if an artifact should be generated or not.
     Will return true if:
         * The artifact_id wasn't set which could be that it's a new object that doesn't have a previous artifact
-        * The source brance is not data only which would indicate that it could contain updates in git to the transform
+        * The source branch is not data only which would indicate that it could contain updates in git to the transform
         * The artifact_id exists in the impacted_artifacts list
+        * The query failes the only_has_unique_targets check
     Will return false if:
         * The source branch is a data only branch and the artifact_id exists and is not in the impacted list
     """
 
-    # if not artifact_id or managed_branch:
-    #    return True
-    # return artifact_id in impacted_artifacts
-    # Temporary workaround tracked in https://github.com/opsmill/infrahub/issues/4991
-    return True
+    if not only_has_unique_targets or not artifact_id or managed_branch:
+        return True
+
+    return artifact_id in impacted_artifacts
 
 
 @flow(
@@ -1262,6 +1292,9 @@ query GatherArtifactDefinitions {
                 name {
                   value
                 }
+                query {
+                  value
+                }
               }
             }
             ... on CoreTransformJinja2 {
@@ -1479,6 +1512,7 @@ def _parse_artifact_definitions(definitions: list[dict]) -> list[ProposedChangeA
             query_name=definition["node"]["transformation"]["node"]["query"]["node"]["name"]["value"],
             query_id=definition["node"]["transformation"]["node"]["query"]["node"]["id"],
             query_models=definition["node"]["transformation"]["node"]["query"]["node"]["models"]["value"] or [],
+            query_payload=definition["node"]["transformation"]["node"]["query"]["node"]["query"]["value"],
             repository_id=definition["node"]["transformation"]["node"]["repository"]["node"]["id"],
             transform_kind=definition["node"]["transformation"]["node"]["__typename"],
         )

infrahub_sdk/analyzer.py

@@ -30,7 +30,7 @@ class GraphQLOperation(BaseModel):
 
 
 class GraphQLQueryAnalyzer:
-    def __init__(self, query: str, schema: GraphQLSchema | None = None):
+    def __init__(self, query: str, schema: GraphQLSchema | None = None) -> None:
         self.query: str = query
         self.schema: GraphQLSchema | None = schema
         self.document: DocumentNode = parse(self.query)

infrahub_sdk/batch.py

@@ -58,7 +58,7 @@ class InfrahubBatch:
         semaphore: asyncio.Semaphore | None = None,
         max_concurrent_execution: int = 5,
         return_exceptions: bool = False,
-    ):
+    ) -> None:
         self._tasks: list[BatchTask] = []
         self.semaphore = semaphore or asyncio.Semaphore(value=max_concurrent_execution)
         self.return_exceptions = return_exceptions
@@ -90,7 +90,7 @@ class InfrahubBatch:
 
 
 class InfrahubBatchSync:
-    def __init__(self, max_concurrent_execution: int = 5, return_exceptions: bool = False):
+    def __init__(self, max_concurrent_execution: int = 5, return_exceptions: bool = False) -> None:
         self._tasks: list[BatchTaskSync] = []
         self.max_concurrent_execution = max_concurrent_execution
         self.return_exceptions = return_exceptions

infrahub_sdk/branch.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import warnings
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Literal, overload
 from urllib.parse import urlencode
 
@@ -14,6 +15,13 @@ if TYPE_CHECKING:
     from .client import InfrahubClient, InfrahubClientSync
 
 
+class BranchStatus(str, Enum):
+    OPEN = "OPEN"
+    NEED_REBASE = "NEED_REBASE"
+    NEED_UPGRADE_REBASE = "NEED_UPGRADE_REBASE"
+    DELETING = "DELETING"
+
+
 class BranchData(BaseModel):
     id: str
     name: str
@@ -21,6 +29,8 @@ class BranchData(BaseModel):
     sync_with_git: bool
     is_default: bool
     has_schema_changes: bool
+    graph_version: int | None = None
+    status: BranchStatus = BranchStatus.OPEN
     origin_branch: str | None = None
     branched_from: str
 
@@ -34,6 +44,8 @@ BRANCH_DATA = {
     "is_default": None,
     "sync_with_git": None,
     "has_schema_changes": None,
+    "graph_version": None,
+    "status": None,
 }
 
 BRANCH_DATA_FILTER = {"@filters": {"name": "$branch_name"}}
@@ -71,7 +83,7 @@ class InfraHubBranchManagerBase:
 
 
 class InfrahubBranchManager(InfraHubBranchManagerBase):
-    def __init__(self, client: InfrahubClient):
+    def __init__(self, client: InfrahubClient) -> None:
         self.client = client
 
     @overload
@@ -221,7 +233,7 @@ class InfrahubBranchManager(InfraHubBranchManagerBase):
 
 
 class InfrahubBranchManagerSync(InfraHubBranchManagerBase):
-    def __init__(self, client: InfrahubClientSync):
+    def __init__(self, client: InfrahubClientSync) -> None:
         self.client = client
 
     def all(self) -> dict[str, BranchData]:

infrahub_sdk/checks.py

@@ -43,7 +43,7 @@ class InfrahubCheck:
         initializer: InfrahubCheckInitializer | None = None,
         params: dict | None = None,
         client: InfrahubClient | None = None,
-    ):
+    ) -> None:
         self.git: GitRepoManager | None = None
         self.initializer = initializer or InfrahubCheckInitializer()
 

infrahub_sdk/client.py

@@ -140,7 +140,7 @@ class BaseClient:
         self,
         address: str = "",
         config: Config | dict[str, Any] | None = None,
-    ):
+    ) -> None:
         self.client = None
         self.headers = {"content-type": "application/json"}
         self.access_token: str = ""
@@ -327,7 +327,6 @@ class InfrahubClient(BaseClient):
         self.object_store = ObjectStore(self)
         self.store = NodeStore(default_branch=self.default_branch)
         self.task = InfrahubTaskManager(self)
-        self.concurrent_execution_limit = asyncio.Semaphore(self.max_concurrent_execution)
         self._request_method: AsyncRequester = self.config.requester or self._default_request_method
         self.group_context = InfrahubGroupContext(self)
 
@@ -1583,8 +1582,7 @@ class InfrahubClient(BaseClient):
 
     async def create_batch(self, return_exceptions: bool = False) -> InfrahubBatch:
         return InfrahubBatch(
-            semaphore=self.concurrent_execution_limit,
-            return_exceptions=return_exceptions,
+            max_concurrent_execution=self.max_concurrent_execution, return_exceptions=return_exceptions
         )
 
     async def get_list_repositories(

infrahub_sdk/ctl/branch.py

@@ -46,6 +46,7 @@ async def list_branch(_: str = CONFIG_PARAM) -> None:
     table.add_column("Sync with Git")
     table.add_column("Has Schema Changes")
     table.add_column("Is Default")
+    table.add_column("Status")
 
     # identify the default branch and always print it first
     default_branch = [branch for branch in branches.values() if branch.is_default][0]
@@ -57,6 +58,7 @@ async def list_branch(_: str = CONFIG_PARAM) -> None:
         "[green]True" if default_branch.sync_with_git else "[#FF7F50]False",
         "[green]True" if default_branch.has_schema_changes else "[#FF7F50]False",
         "[green]True" if default_branch.is_default else "[#FF7F50]False",
+        default_branch.status,
     )
 
     for branch in branches.values():
@@ -71,6 +73,7 @@ async def list_branch(_: str = CONFIG_PARAM) -> None:
             "[green]True" if branch.sync_with_git else "[#FF7F50]False",
             "[green]True" if default_branch.has_schema_changes else "[#FF7F50]False",
             "[green]True" if branch.is_default else "[#FF7F50]False",
+            branch.status,
         )
 
     console.print(table)

infrahub_sdk/ctl/cli_commands.py

@@ -32,6 +32,7 @@ from ..ctl.render import list_jinja2_transforms, print_template_errors
 from ..ctl.repository import app as repository_app
 from ..ctl.repository import find_repository_config_file, get_repository_config
 from ..ctl.schema import app as schema_app
+from ..ctl.task import app as task_app
 from ..ctl.transform import list_transforms
 from ..ctl.utils import (
     catch_exception,
@@ -64,6 +65,7 @@ app.add_typer(repository_app, name="repository")
 app.add_typer(menu_app, name="menu")
 app.add_typer(object_app, name="object")
 app.add_typer(graphql_app, name="graphql")
+app.add_typer(task_app, name="task")
 
 app.command(name="dump")(dump)
 app.command(name="load")(load)

infrahub_sdk/ctl/exceptions.py

@@ -3,6 +3,6 @@ class Error(Exception):
 
 
 class QueryNotFoundError(Error):
-    def __init__(self, name: str, message: str = ""):
+    def __init__(self, name: str, message: str = "") -> None:
         self.message = message or f"The requested query '{name}' was not found."
         super().__init__(self.message)

infrahub_sdk/ctl/task.py

@@ -0,0 +1,110 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+from ..async_typer import AsyncTyper
+from ..task.manager import TaskFilter
+from ..task.models import Task, TaskState
+from .client import initialize_client
+from .parameters import CONFIG_PARAM
+from .utils import catch_exception, init_logging
+
+app = AsyncTyper()
+console = Console()
+
+
+@app.callback()
+def callback() -> None:
+    """Manage Infrahub tasks."""
+
+
+def _parse_states(states: list[str] | None) -> list[TaskState] | None:
+    if not states:
+        return None
+
+    parsed_states: list[TaskState] = []
+    for state in states:
+        normalized_state = state.strip().upper()
+        try:
+            parsed_states.append(TaskState(normalized_state))
+        except ValueError as exc:  # pragma: no cover - typer will surface this as CLI error
+            raise typer.BadParameter(
+                f"Unsupported state '{state}'. Available states: {', '.join(item.value.lower() for item in TaskState)}"
+            ) from exc
+
+    return parsed_states
+
+
+def _render_table(tasks: list[Task]) -> None:
+    table = Table(title="Infrahub Tasks", box=None)
+    table.add_column("ID", style="cyan", overflow="fold")
+    table.add_column("Title", style="magenta", overflow="fold")
+    table.add_column("State", style="green")
+    table.add_column("Progress", justify="right")
+    table.add_column("Workflow", overflow="fold")
+    table.add_column("Branch", overflow="fold")
+    table.add_column("Updated")
+
+    if not tasks:
+        table.add_row("-", "No tasks found", "-", "-", "-", "-", "-")
+        console.print(table)
+        return
+
+    for task in tasks:
+        progress = f"{task.progress:.0%}" if task.progress is not None else "-"
+        table.add_row(
+            task.id,
+            task.title,
+            task.state.value,
+            progress,
+            task.workflow or "-",
+            task.branch or "-",
+            task.updated_at.isoformat(),
+        )
+
+    console.print(table)
+
+
+@app.command(name="list")
+@catch_exception(console=console)
+async def list_tasks(
+    state: list[str] = typer.Option(
+        None, "--state", "-s", help="Filter by task state. Can be provided multiple times."
+    ),
+    limit: Optional[int] = typer.Option(None, help="Maximum number of tasks to retrieve."),
+    offset: Optional[int] = typer.Option(None, help="Offset for pagination."),
+    include_related_nodes: bool = typer.Option(False, help="Include related nodes in the output."),
+    include_logs: bool = typer.Option(False, help="Include task logs in the output."),
+    json_output: bool = typer.Option(False, "--json", help="Output the result as JSON."),
+    debug: bool = False,
+    _: str = CONFIG_PARAM,
+) -> None:
+    """List Infrahub tasks."""
+
+    init_logging(debug=debug)
+
+    client = initialize_client()
+    filters = TaskFilter()
+    parsed_states = _parse_states(state)
+    if parsed_states:
+        filters.state = parsed_states
+
+    tasks = await client.task.filter(
+        filter=filters,
+        limit=limit,
+        offset=offset,
+        include_related_nodes=include_related_nodes,
+        include_logs=include_logs,
+    )
+
+    if json_output:
+        console.print_json(
+            data=[task.model_dump(mode="json") for task in tasks], indent=2, sort_keys=True, highlight=False
+        )
+        return
+
+    _render_table(tasks)

infrahub_sdk/exceptions.py

@@ -5,13 +5,13 @@ from typing import Any
 
 
 class Error(Exception):
-    def __init__(self, message: str | None = None):
+    def __init__(self, message: str | None = None) -> None:
         self.message = message
         super().__init__(self.message)
 
 
 class JsonDecodeError(Error):
-    def __init__(self, message: str | None = None, content: str | None = None, url: str | None = None):
+    def __init__(self, message: str | None = None, content: str | None = None, url: str | None = None) -> None:
         self.message = message
         self.content = content
         self.url = url
@@ -23,14 +23,14 @@ class JsonDecodeError(Error):
 
 
 class ServerNotReachableError(Error):
-    def __init__(self, address: str, message: str | None = None):
+    def __init__(self, address: str, message: str | None = None) -> None:
         self.address = address
         self.message = message or f"Unable to connect to '{address}'."
         super().__init__(self.message)
 
 
 class ServerNotResponsiveError(Error):
-    def __init__(self, url: str, timeout: int | None = None, message: str | None = None):
+    def __init__(self, url: str, timeout: int | None = None, message: str | None = None) -> None:
         self.url = url
         self.timeout = timeout
         self.message = message or f"Unable to read from '{url}'."
@@ -40,7 +40,7 @@ class ServerNotResponsiveError(Error):
 
 
 class GraphQLError(Error):
-    def __init__(self, errors: list[dict[str, Any]], query: str | None = None, variables: dict | None = None):
+    def __init__(self, errors: list[dict[str, Any]], query: str | None = None, variables: dict | None = None) -> None:
         self.query = query
         self.variables = variables
         self.errors = errors
@@ -49,21 +49,21 @@ class GraphQLError(Error):
 
 
 class BranchNotFoundError(Error):
-    def __init__(self, identifier: str, message: str | None = None):
+    def __init__(self, identifier: str, message: str | None = None) -> None:
         self.identifier = identifier
         self.message = message or f"Unable to find the branch '{identifier}' in the Database."
         super().__init__(self.message)
 
 
 class SchemaNotFoundError(Error):
-    def __init__(self, identifier: str, message: str | None = None):
+    def __init__(self, identifier: str, message: str | None = None) -> None:
         self.identifier = identifier
         self.message = message or f"Unable to find the schema '{identifier}'."
         super().__init__(self.message)
 
 
 class ModuleImportError(Error):
-    def __init__(self, message: str | None = None):
+    def __init__(self, message: str | None = None) -> None:
         self.message = message or "Unable to import the module"
         super().__init__(self.message)
 
@@ -75,7 +75,7 @@ class NodeNotFoundError(Error):
         message: str = "Unable to find the node in the database.",
         branch_name: str | None = None,
         node_type: str | None = None,
-    ):
+    ) -> None:
         self.node_type = node_type or "unknown"
         self.identifier = identifier
         self.branch_name = branch_name
@@ -97,25 +97,25 @@ class NodeInvalidError(NodeNotFoundError):
 class ResourceNotDefinedError(Error):
     """Raised when trying to access a resource that hasn't been defined."""
 
-    def __init__(self, message: str | None = None):
+    def __init__(self, message: str | None = None) -> None:
         self.message = message or "The requested resource was not found"
         super().__init__(self.message)
 
 
 class InfrahubCheckNotFoundError(Error):
-    def __init__(self, name: str, message: str | None = None):
+    def __init__(self, name: str, message: str | None = None) -> None:
         self.message = message or f"The requested InfrahubCheck '{name}' was not found."
         super().__init__(self.message)
 
 
 class InfrahubTransformNotFoundError(Error):
-    def __init__(self, name: str, message: str | None = None):
+    def __init__(self, name: str, message: str | None = None) -> None:
         self.message = message or f"The requested InfrahubTransform '{name}' was not found."
         super().__init__(self.message)
 
 
 class ValidationError(Error):
-    def __init__(self, identifier: str, message: str | None = None, messages: list[str] | None = None):
+    def __init__(self, identifier: str, message: str | None = None, messages: list[str] | None = None) -> None:
         self.identifier = identifier
         self.message = message
         self.messages = messages
@@ -130,7 +130,7 @@ class ValidationError(Error):
 
 
 class ObjectValidationError(Error):
-    def __init__(self, position: list[int | str], message: str):
+    def __init__(self, position: list[int | str], message: str) -> None:
         self.position = position
         self.message = message
         super().__init__(self.message)
@@ -140,13 +140,13 @@ class ObjectValidationError(Error):
 
 
 class AuthenticationError(Error):
-    def __init__(self, message: str | None = None):
+    def __init__(self, message: str | None = None) -> None:
         self.message = message or "Authentication Error, unable to execute the query."
         super().__init__(self.message)
 
 
 class URLNotFoundError(Error):
-    def __init__(self, url: str):
+    def __init__(self, url: str) -> None:
         self.message = f"`{url}` not found."
         super().__init__(self.message)
 
@@ -164,12 +164,12 @@ class InvalidResponseError(Error):
 
 
 class FileNotValidError(Error):
-    def __init__(self, name: str, message: str = ""):
+    def __init__(self, name: str, message: str = "") -> None:
         self.message = message or f"Cannot parse '{name}' content."
         super().__init__(self.message)
 
 
 class TimestampFormatError(Error):
-    def __init__(self, message: str | None = None):
+    def __init__(self, message: str | None = None) -> None:
         self.message = message or "Invalid timestamp format"
         super().__init__(self.message)

infrahub_sdk/graphql/query.py

@@ -9,7 +9,7 @@ class BaseGraphQLQuery:
     query_type: str = "not-defined"
     indentation: int = 4
 
-    def __init__(self, query: dict, variables: dict | None = None, name: str | None = None):
+    def __init__(self, query: dict, variables: dict | None = None, name: str | None = None) -> None:
         self.query = query
         self.variables = variables
         self.name = name or ""
@@ -46,7 +46,7 @@ class Query(BaseGraphQLQuery):
 class Mutation(BaseGraphQLQuery):
     query_type = "mutation"
 
-    def __init__(self, *args: Any, mutation: str, input_data: dict, **kwargs: Any):
+    def __init__(self, *args: Any, mutation: str, input_data: dict, **kwargs: Any) -> None:
         self.input_data = input_data
         self.mutation = mutation
         super().__init__(*args, **kwargs)

infrahub_sdk/node/attribute.py

@@ -15,7 +15,7 @@ if TYPE_CHECKING:
 class Attribute:
     """Represents an attribute of a Node, including its schema, value, and properties."""
 
-    def __init__(self, name: str, schema: AttributeSchemaAPI, data: Any | dict):
+    def __init__(self, name: str, schema: AttributeSchemaAPI, data: Any | dict) -> None:
         """
         Args:
             name (str): The name of the attribute.

infrahub_sdk/node/property.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 class NodeProperty:
     """Represents a property of a node, typically used for metadata like display labels."""
 
-    def __init__(self, data: dict | str):
+    def __init__(self, data: dict | str) -> None:
         """
         Args:
             data (Union[dict, str]): Data representing the node property.

infrahub_sdk/node/related_node.py

@@ -17,7 +17,7 @@ if TYPE_CHECKING:
 class RelatedNodeBase:
     """Base class for representing a related node in a relationship."""
 
-    def __init__(self, branch: str, schema: RelationshipSchemaAPI, data: Any | dict, name: str | None = None):
+    def __init__(self, branch: str, schema: RelationshipSchemaAPI, data: Any | dict, name: str | None = None) -> None:
         """
         Args:
             branch (str): The branch where the related node resides.
@@ -189,7 +189,7 @@ class RelatedNode(RelatedNodeBase):
         schema: RelationshipSchemaAPI,
         data: Any | dict,
         name: str | None = None,
-    ):
+    ) -> None:
         """
         Args:
             client (InfrahubClient): The client used to interact with the backend asynchronously.
@@ -236,7 +236,7 @@ class RelatedNodeSync(RelatedNodeBase):
         schema: RelationshipSchemaAPI,
         data: Any | dict,
         name: str | None = None,
-    ):
+    ) -> None:
         """
         Args:
             client (InfrahubClientSync): The client used to interact with the backend synchronously.

infrahub_sdk/node/relationship.py

@@ -4,7 +4,6 @@ from collections import defaultdict
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
-from ..batch import InfrahubBatch
 from ..exceptions import (
     Error,
     UninitializedError,
@@ -22,7 +21,7 @@ if TYPE_CHECKING:
 class RelationshipManagerBase:
     """Base class for RelationshipManager and RelationshipManagerSync"""
 
-    def __init__(self, name: str, branch: str, schema: RelationshipSchemaAPI):
+    def __init__(self, name: str, branch: str, schema: RelationshipSchemaAPI) -> None:
         """
         Args:
             name (str): The name of the relationship.
@@ -108,7 +107,7 @@ class RelationshipManager(RelationshipManagerBase):
         branch: str,
         schema: RelationshipSchemaAPI,
         data: Any | dict,
-    ):
+    ) -> None:
         """
         Args:
             name (str): The name of the relationship.
@@ -166,7 +165,7 @@ class RelationshipManager(RelationshipManagerBase):
                 raise Error("Unable to fetch the peer, id and/or typename are not defined")
             ids_per_kind_map[peer.typename].append(peer.id)
 
-        batch = InfrahubBatch(max_concurrent_execution=self.client.max_concurrent_execution)
+        batch = await self.client.create_batch()
         for kind, ids in ids_per_kind_map.items():
             batch.add(
                 task=self.client.filters,
@@ -231,7 +230,7 @@ class RelationshipManagerSync(RelationshipManagerBase):
         branch: str,
         schema: RelationshipSchemaAPI,
         data: Any | dict,
-    ):
+    ) -> None:
         """
         Args:
             name (str): The name of the relationship.
@@ -289,7 +288,6 @@ class RelationshipManagerSync(RelationshipManagerBase):
                 raise Error("Unable to fetch the peer, id and/or typename are not defined")
             ids_per_kind_map[peer.typename].append(peer.id)
 
-        # Unlike Async, no need to create a new batch from scratch because we are not using a semaphore
         batch = self.client.create_batch()
         for kind, ids in ids_per_kind_map.items():
             batch.add(