infrahub-server 1.4.10__py3-none-any.whl → 1.5.0b0__py3-none-any.whl

infrahub/workflows/catalogue.py

@@ -1,6 +1,7 @@
 import random
 
 from fast_depends import Depends, inject
+from prefect.client.schemas.objects import ConcurrencyLimitStrategy
 
 from .constants import WorkflowTag, WorkflowType
 from .models import WorkerPoolDefinition, WorkflowDefinition
@@ -17,14 +18,14 @@ ACTION_ADD_NODE_TO_GROUP = WorkflowDefinition(
 
 ACTION_RUN_GENERATOR = WorkflowDefinition(
     name="action-run-generator",
-    type=WorkflowType.CORE,
+    type=WorkflowType.INTERNAL,
     module="infrahub.actions.tasks",
     function="run_generator",
 )
 
 ACTION_RUN_GENERATOR_GROUP_EVENT = WorkflowDefinition(
     name="action-run-generator-group-event",
-    type=WorkflowType.CORE,
+    type=WorkflowType.INTERNAL,
     module="infrahub.actions.tasks",
     function="run_generator_group_event",
 )
@@ -83,7 +84,7 @@ TRIGGER_ARTIFACT_DEFINITION_GENERATE = WorkflowDefinition(
 
 TRIGGER_GENERATOR_DEFINITION_RUN = WorkflowDefinition(
     name="generator-definition-run",
-    type=WorkflowType.CORE,
+    type=WorkflowType.INTERNAL,
     module="infrahub.generators.tasks",
     function="run_generator_definition",
     tags=[WorkflowTag.DATABASE_CHANGE],
@@ -179,6 +180,8 @@ GIT_REPOSITORIES_SYNC = WorkflowDefinition(
     cron="* * * * *",
     module="infrahub.git.tasks",
     function="sync_remote_repositories",
+    concurrency_limit=1,
+    concurrency_limit_strategy=ConcurrencyLimitStrategy.CANCEL_NEW,
 )
 
 GIT_REPOSITORIES_CREATE_BRANCH = WorkflowDefinition(
@@ -531,6 +534,35 @@ VALIDATE_SCHEMA_NUMBER_POOLS = WorkflowDefinition(
 )
 
 
+PROFILE_REFRESH_MULTIPLE = WorkflowDefinition(
+    name="objects-profiles-refresh-multiple",
+    type=WorkflowType.CORE,
+    module="infrahub.profiles.tasks",
+    function="objects_profiles_refresh_multiple",
+    tags=[WorkflowTag.DATABASE_CHANGE],
+)
+
+
+PROFILE_REFRESH = WorkflowDefinition(
+    name="object-profiles-refresh",
+    type=WorkflowType.CORE,
+    module="infrahub.profiles.tasks",
+    function="object_profiles_refresh",
+    tags=[WorkflowTag.DATABASE_CHANGE],
+)
+
+
+CLEAN_UP_DEADLOCKS = WorkflowDefinition(
+    name="clean-up-deadlocks",
+    type=WorkflowType.INTERNAL,
+    cron="* * * * *",
+    module="infrahub.locks.tasks",
+    function="clean_up_deadlocks",
+    concurrency_limit=1,
+    concurrency_limit_strategy=ConcurrencyLimitStrategy.CANCEL_NEW,
+)
+
+
 WORKER_POOLS = [INFRAHUB_WORKER_POOL]
 
 WORKFLOWS = [
@@ -547,6 +579,7 @@ WORKFLOWS = [
     BRANCH_MERGE_POST_PROCESS,
     BRANCH_REBASE,
     BRANCH_VALIDATE,
+    CLEAN_UP_DEADLOCKS,
     COMPUTED_ATTRIBUTE_JINJA2_UPDATE_VALUE,
     COMPUTED_ATTRIBUTE_PROCESS_JINJA2,
     COMPUTED_ATTRIBUTE_PROCESS_TRANSFORM,
@@ -572,6 +605,8 @@ WORKFLOWS = [
     GIT_REPOSITORY_USER_CHECK_RUN,
     GRAPHQL_QUERY_GROUP_UPDATE,
     IPAM_RECONCILIATION,
+    PROFILE_REFRESH,
+    PROFILE_REFRESH_MULTIPLE,
     PROPOSED_CHANGE_MERGE,
     QUERY_COMPUTED_ATTRIBUTE_TRANSFORM_TARGETS,
     REMOVE_ADD_NODE_FROM_GROUP,

infrahub/workflows/models.py

@@ -6,7 +6,7 @@ from uuid import UUID
 from prefect import Flow
 from prefect.client.orchestration import PrefectClient
 from prefect.client.schemas.actions import DeploymentScheduleCreate
-from prefect.client.schemas.objects import FlowRun
+from prefect.client.schemas.objects import ConcurrencyLimitStrategy, FlowRun
 from prefect.client.schemas.schedules import CronSchedule
 from pydantic import BaseModel, Field
 from typing_extensions import Self
@@ -48,6 +48,14 @@ class WorkflowDefinition(BaseModel):
     function: str
     cron: str | None = None
     tags: list[WorkflowTag] = Field(default_factory=list)
+    concurrency_limit: int | None = Field(
+        default=None,
+        description="The concurrency limit for the deployment.",
+    )
+    concurrency_limit_strategy: ConcurrencyLimitStrategy | None = Field(
+        default=None,
+        description="The concurrency options for the deployment.",
+    )
 
     @property
     def entrypoint(self) -> str:
@@ -60,7 +68,14 @@ class WorkflowDefinition(BaseModel):
         return f"{self.name}/{self.name}"
 
     def to_deployment(self) -> dict[str, Any]:
-        payload: dict[str, Any] = {"name": self.name, "entrypoint": self.entrypoint, "tags": self.get_tags()}
+        payload: dict[str, Any] = {
+            "name": self.name,
+            "entrypoint": self.entrypoint,
+            "tags": self.get_tags(),
+            "concurrency_limit": self.concurrency_limit,
+        }
+        if self.concurrency_limit_strategy:
+            payload["concurrency_options"] = {"collision_strategy": self.concurrency_limit_strategy}
         if self.type == WorkflowType.CORE:
             payload["version"] = __version__
         if self.cron:

infrahub_sdk/branch.py

@@ -188,9 +188,7 @@ class InfrahubBranchManager(InfraHubBranchManagerBase):
         query = Query(name="GetAllBranch", query=QUERY_ALL_BRANCHES_DATA)
         data = await self.client.execute_graphql(query=query.render(), tracker="query-branch-all")
 
-        branches = {branch["name"]: BranchData(**branch) for branch in data["Branch"]}
-
-        return branches
+        return {branch["name"]: BranchData(**branch) for branch in data["Branch"]}
 
     async def get(self, branch_name: str) -> BranchData:
         query = Query(name="GetBranch", query=QUERY_ONE_BRANCH_DATA, variables={"branch_name": str})
@@ -230,9 +228,7 @@ class InfrahubBranchManagerSync(InfraHubBranchManagerBase):
         query = Query(name="GetAllBranch", query=QUERY_ALL_BRANCHES_DATA)
         data = self.client.execute_graphql(query=query.render(), tracker="query-branch-all")
 
-        branches = {branch["name"]: BranchData(**branch) for branch in data["Branch"]}
-
-        return branches
+        return {branch["name"]: BranchData(**branch) for branch in data["Branch"]}
 
     def get(self, branch_name: str) -> BranchData:
         query = Query(name="GetBranch", query=QUERY_ONE_BRANCH_DATA, variables={"branch_name": str})
@@ -292,13 +288,14 @@ class InfrahubBranchManagerSync(InfraHubBranchManagerBase):
             },
         }
 
-        query = Mutation(mutation="BranchCreate", input_data=input_data, query=MUTATION_QUERY_DATA)
+        mutation_query = MUTATION_QUERY_TASK if background_execution else MUTATION_QUERY_DATA
+        query = Mutation(mutation="BranchCreate", input_data=input_data, query=mutation_query)
         response = self.client.execute_graphql(query=query.render(), tracker="mutation-branch-create")
 
         # Make sure server version is recent enough to support background execution, as previously
         # using background_execution=True had no effect.
         if background_execution and "task" in response["BranchCreate"]:
-            return BranchData(**response["BranchCreate"]["task"]["id"])
+            return response["BranchCreate"]["task"]["id"]
         return BranchData(**response["BranchCreate"]["object"])
 
     def delete(self, branch_name: str) -> bool: