Merge pull request #5819 from opsmill/dga-20250222-diff-update-task

dgarros · web-flow · commit 7a0b12d477af · 2025-02-24T12:50:56.000+01:00
Update DiffUpdate mutation to return the id of the task when running asynchronously
diff --git a/backend/infrahub/core/diff/tasks.py b/backend/infrahub/core/diff/tasks.py
@@ -8,15 +8,15 @@
 from infrahub.log import get_logger
 from infrahub.services import services
 from infrahub.workflows.catalogue import DIFF_REFRESH
-from infrahub.workflows.utils import add_branch_tag
+from infrahub.workflows.utils import add_tags
 
 log = get_logger()
 
 
 @flow(name="diff-update", flow_run_name="Update diff for branch {model.branch_name}")
 async def update_diff(model: RequestDiffUpdate) -> None:
     service = services.service
-    await add_branch_tag(branch_name=model.branch_name)
+    await add_tags(branches=[model.branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -37,7 +37,7 @@ async def update_diff(model: RequestDiffUpdate) -> None:
 @flow(name="diff-refresh", flow_run_name="Recreate diff for branch {branch_name}")
 async def refresh_diff(branch_name: str, diff_id: str) -> None:
     service = services.service
-    await add_branch_tag(branch_name=branch_name)
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
@@ -51,7 +51,7 @@ async def refresh_diff(branch_name: str, diff_id: str) -> None:
 @flow(name="diff-refresh-all", flow_run_name="Recreate all diffs for branch {branch_name}")
 async def refresh_diff_all(branch_name: str) -> None:
     service = services.service
-    await add_branch_tag(branch_name=branch_name)
+    await add_tags(branches=[branch_name])
 
     async with service.database.start_session() as db:
         component_registry = get_component_registry()
diff --git a/backend/infrahub/graphql/mutations/diff.py b/backend/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
@@ -14,6 +14,8 @@
 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,14 +25,16 @@ 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")
@@ -39,9 +43,13 @@ async def mutate(
         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 +61,11 @@ async def mutate(
         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 +76,7 @@ async def mutate(
             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 +97,7 @@ async def mutate(
             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}
diff --git a/backend/tests/integration/diff/test_diff_update.py b/backend/tests/integration/diff/test_diff_update.py
@@ -32,8 +32,8 @@
 BRANCH_NAME = "branch1"
 PROPOSED_CHANGE_NAME = "branch1-pc"
 DIFF_UPDATE_QUERY = """
-mutation DiffUpdate($branch_name: String!, $wait_for_completion: Boolean) {
-    DiffUpdate(data: { branch: $branch_name, wait_for_completion: $wait_for_completion }) {
+mutation DiffUpdate($branch_name: String!, $wait_for_completion: Boolean = true) {
+    DiffUpdate(data: { branch: $branch_name }, wait_until_completion: $wait_for_completion) {
         ok
     }
 }
diff --git a/backend/tests/unit/graphql/diff/test_diff_update_mutation.py b/backend/tests/unit/graphql/diff/test_diff_update_mutation.py
@@ -8,12 +8,27 @@
 from infrahub.core.timestamp import Timestamp
 from infrahub.database import InfrahubDatabase
 from infrahub.graphql.initialization import prepare_graphql_params
+from infrahub.services import InfrahubServices, services
+from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
+from tests.adapters.cache import MemoryCache
+from tests.adapters.message_bus import BusRecorder
 from tests.helpers.graphql import graphql
 
 DIFF_UPDATE_MUTATION = """
-    mutation($branch: String!, $name: String, $from_time: DateTime, $to_time: DateTime) {
-        DiffUpdate(data: {branch: $branch, name: $name, from_time: $from_time, to_time: $to_time, wait_for_completion: true}) {
+    mutation($branch: String!, $name: String, $from_time: DateTime, $to_time: DateTime, $wait_until_completion: Boolean = true) {
+        DiffUpdate(
+            data: {
+                branch: $branch,
+                name: $name,
+                from_time: $from_time,
+                to_time: $to_time
+            },
+            wait_until_completion: $wait_until_completion
+        ) {
             ok
+            task {
+                id
+            }
         }
     }
 """
@@ -22,16 +37,34 @@
 class TestDiffUpdateMutation:
     diff_name = "CountDiffula"
 
+    @pytest.fixture
+    def service_testing(self, db: InfrahubDatabase):
+        original = services.service
+        service = InfrahubServices(
+            database=db, message_bus=BusRecorder(), workflow=WorkflowLocalExecution(), cache=MemoryCache()
+        )
+        services.service = service
+        services.prepare(service=services.service)
+        yield service
+        services.service = original
+        services.prepare(service=services.service)
+
     @pytest.fixture
     async def diff_branch(self, db: InfrahubDatabase) -> Branch:
         return await create_branch(db=db, branch_name="branch")
 
     @pytest.fixture
     async def named_diff(
-        self, db: InfrahubDatabase, default_branch: Branch, criticality_schema, diff_branch: Branch
+        self,
+        db: InfrahubDatabase,
+        default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
+        criticality_schema,
+        diff_branch: Branch,
     ) -> EnrichedDiffRootMetadata:
         params = await prepare_graphql_params(
-            db=db, include_mutation=True, include_subscription=False, branch=default_branch
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
         )
         result = await graphql(
             schema=params.schema,
@@ -53,11 +86,17 @@ async def named_diff(
         )[0]
 
     async def test_create_diff_before_branched_from_fails(
-        self, db: InfrahubDatabase, default_branch: Branch, criticality_schema, diff_branch: Branch
+        self,
+        db: InfrahubDatabase,
+        default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
+        criticality_schema,
+        diff_branch: Branch,
     ):
         branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
         params = await prepare_graphql_params(
-            db=db, include_mutation=True, include_subscription=False, branch=default_branch
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
         )
         result = await graphql(
             schema=params.schema,
@@ -74,11 +113,17 @@ async def test_create_diff_before_branched_from_fails(
         assert result.data["DiffUpdate"]["ok"] is True
 
     async def test_create_time_range_diff_without_name_fails(
-        self, db: InfrahubDatabase, default_branch: Branch, criticality_schema, diff_branch: Branch
+        self,
+        db: InfrahubDatabase,
+        default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
+        criticality_schema,
+        diff_branch: Branch,
     ):
         branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
         params = await prepare_graphql_params(
-            db=db, include_mutation=True, include_subscription=False, branch=default_branch
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
         )
         result = await graphql(
             schema=params.schema,
@@ -99,12 +144,14 @@ async def test_create_diff_with_illegal_times_fails(
         self,
         db: InfrahubDatabase,
         default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
         criticality_schema,
         diff_branch: Branch,
         named_diff: EnrichedDiffRootMetadata,
     ):
         params = await prepare_graphql_params(
-            db=db, include_mutation=True, include_subscription=False, branch=default_branch
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
         )
         result = await graphql(
             schema=params.schema,
@@ -140,13 +187,44 @@ async def test_create_named_diff_with_legal_times_succeeds(
         self,
         db: InfrahubDatabase,
         default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
+        criticality_schema,
+        diff_branch: Branch,
+        named_diff: EnrichedDiffRootMetadata,
+    ):
+        branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
+        params = await prepare_graphql_params(
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
+        )
+        result = await graphql(
+            schema=params.schema,
+            source=DIFF_UPDATE_MUTATION,
+            context_value=params.context,
+            root_value=None,
+            variable_values={
+                "branch": diff_branch.name,
+                "from_time": branched_from_timestamp.to_string(),
+                "to_time": Timestamp().to_string(),
+                "name": self.diff_name,
+            },
+        )
+        assert result.errors is None
+        assert result.data["DiffUpdate"]["ok"] is True
+
+    async def test_retrieve_task_id(
+        self,
+        db: InfrahubDatabase,
+        default_branch: Branch,
+        prefect_test_fixture,
+        service_testing: InfrahubServices,
         criticality_schema,
         diff_branch: Branch,
         named_diff: EnrichedDiffRootMetadata,
     ):
         branched_from_timestamp = Timestamp(diff_branch.get_branched_from())
         params = await prepare_graphql_params(
-            db=db, include_mutation=True, include_subscription=False, branch=default_branch
+            db=db, include_mutation=True, include_subscription=False, branch=default_branch, service=service_testing
         )
         result = await graphql(
             schema=params.schema,
@@ -158,7 +236,9 @@ async def test_create_named_diff_with_legal_times_succeeds(
                 "from_time": branched_from_timestamp.to_string(),
                 "to_time": Timestamp().to_string(),
                 "name": self.diff_name,
+                "wait_until_completion": False,
             },
         )
         assert result.errors is None
         assert result.data["DiffUpdate"]["ok"] is True
+        assert result.data["DiffUpdate"]["task"]["id"] is not None
diff --git a/changelog/+diff-update.changed.md b/changelog/+diff-update.changed.md
@@ -0,0 +1 @@
+Update DiffUpdate mutation to return the id of the task when `wait_until_completion` is False. Also the argument `wait_for_completion` under data is deprecated and it has been replaced with `wait_until_completion` at the root of the mutation instead to align with the format of the other mutations.

Original file line number	Diff line number	Diff line change
`@@ -32,8 +32,8 @@`
`32`	`32`	`BRANCH_NAME = "branch1"`
`33`	`33`	`PROPOSED_CHANGE_NAME = "branch1-pc"`
`34`	`34`	`DIFF_UPDATE_QUERY = """`
`35`		`-mutation DiffUpdate($branch_name: String!, $wait_for_completion: Boolean) {`
`36`		`- DiffUpdate(data: { branch: $branch_name, wait_for_completion: $wait_for_completion }) {`
	`35`	`+mutation DiffUpdate($branch_name: String!, $wait_for_completion: Boolean = true) {`
	`36`	`+ DiffUpdate(data: { branch: $branch_name }, wait_until_completion: $wait_for_completion) {`
`37`	`37`	`ok`
`38`	`38`	`}`
`39`	`39`	`}`
Original file line number	Diff line number	Diff line change
`@@ -0,0 +1 @@`
	`1`	+Update DiffUpdate mutation to return the id of the task when `wait_until_completion` is False. Also the argument `wait_for_completion` under data is deprecated and it has been replaced with `wait_until_completion` at the root of the mutation instead to align with the format of the other mutations.