Merge pull request #4845 from opsmill/lgu-migrate-req-gql-group-update

LucasG0 · web-flow · commit 5f20627fa208 · 2024-11-05T16:30:04.000+01:00
Migrate RequestGraphQLQueryGroupUpdate to prefect
diff --git a/backend/infrahub/api/query.py b/backend/infrahub/api/query.py
@@ -24,8 +24,9 @@
     GRAPHQL_TOP_LEVEL_QUERIES_METRICS,
 )
 from infrahub.graphql.utils import extract_data
+from infrahub.groups.models import RequestGraphQLQueryGroupUpdate
 from infrahub.log import get_logger
-from infrahub.message_bus import messages
+from infrahub.workflows.catalogue import UPDATE_GRAPHQL_QUERY_GROUP
 
 if TYPE_CHECKING:
     from infrahub.auth import AccountSession
@@ -106,16 +107,15 @@ async def execute_query(
 
     if update_group:
         service: InfrahubServices = request.app.state.service
-        await service.send(
-            message=messages.RequestGraphQLQueryGroupUpdate(
-                branch=branch_params.branch.name,
-                query_id=gql_query.id,
-                query_name=gql_query.name.value,
-                related_node_ids=sorted(list(related_node_ids)),
-                subscribers=sorted(subscribers),
-                params=params,
-            )
+        model = RequestGraphQLQueryGroupUpdate(
+            branch=branch_params.branch.name,
+            query_id=gql_query.id,
+            query_name=gql_query.name.value,
+            related_node_ids=sorted(list(related_node_ids)),
+            subscribers=sorted(subscribers),
+            params=params,
         )
+        await service.workflow.submit_workflow(workflow=UPDATE_GRAPHQL_QUERY_GROUP, parameters={"model": model})
 
     return response_payload
 
diff --git a/backend/infrahub/groups/__init__.py b/backend/infrahub/groups/__init__.py
diff --git a/backend/infrahub/groups/models.py b/backend/infrahub/groups/models.py
@@ -1,9 +1,7 @@
-from pydantic import Field
+from pydantic import BaseModel, Field
 
-from infrahub.message_bus import InfrahubMessage
 
-
-class RequestGraphQLQueryGroupUpdate(InfrahubMessage):
+class RequestGraphQLQueryGroupUpdate(BaseModel):
     """Sent to create or update a GraphQLQueryGroup associated with a given GraphQLQuery."""
 
     branch: str = Field(..., description="The branch to target")
diff --git a/backend/infrahub/groups/tasks.py b/backend/infrahub/groups/tasks.py
@@ -0,0 +1,36 @@
+from infrahub_sdk.groups import group_add_subscriber
+from infrahub_sdk.utils import dict_hash
+from prefect import flow
+
+from infrahub.core.constants import InfrahubKind
+from infrahub.groups.models import RequestGraphQLQueryGroupUpdate
+from infrahub.services import services
+from infrahub.workflows.utils import add_branch_tag
+
+
+@flow(name="update_graphql_query_group")
+async def update_graphql_query_group(model: RequestGraphQLQueryGroupUpdate) -> None:
+    """Create or Update a GraphQLQueryGroup."""
+
+    await add_branch_tag(branch_name=model.branch)
+    service = services.service
+
+    params_hash = dict_hash(model.params)
+    group_name = f"{model.query_name}__{params_hash}"
+    group_label = f"Query {model.query_name} Hash({params_hash[:8]})"
+    group = await service.client.create(
+        kind=InfrahubKind.GRAPHQLQUERYGROUP,
+        branch=model.branch,
+        name=group_name,
+        label=group_label,
+        group_type="internal",
+        query=model.query_id,
+        parameters=model.params,
+        members=model.related_node_ids,
+    )
+    await group.save(allow_upsert=True)
+
+    if model.subscribers:
+        await group_add_subscriber(
+            client=service.client, group=group, subscribers=model.subscribers, branch=model.branch
+        )
diff --git a/backend/infrahub/message_bus/messages/__init__.py b/backend/infrahub/message_bus/messages/__init__.py
@@ -30,7 +30,6 @@
 from .refresh_webhook_configuration import RefreshWebhookConfiguration
 from .request_artifactdefinition_check import RequestArtifactDefinitionCheck
 from .request_generatordefinition_check import RequestGeneratorDefinitionCheck
-from .request_graphqlquerygroup_update import RequestGraphQLQueryGroupUpdate
 from .request_proposedchange_pipeline import RequestProposedChangePipeline
 from .request_repository_checks import RequestRepositoryChecks
 from .request_repository_userchecks import RequestRepositoryUserChecks
@@ -66,7 +65,6 @@
     "refresh.webhook.configuration": RefreshWebhookConfiguration,
     "request.artifact_definition.check": RequestArtifactDefinitionCheck,
     "request.generator_definition.check": RequestGeneratorDefinitionCheck,
-    "request.graphql_query_group.update": RequestGraphQLQueryGroupUpdate,
     "request.proposed_change.data_integrity": RequestProposedChangeDataIntegrity,
     "request.proposed_change.pipeline": RequestProposedChangePipeline,
     "request.proposed_change.refresh_artifacts": RequestProposedChangeRefreshArtifacts,
diff --git a/backend/infrahub/message_bus/operations/__init__.py b/backend/infrahub/message_bus/operations/__init__.py
@@ -42,7 +42,6 @@
     "refresh.registry.rebased_branch": refresh.registry.rebased_branch,
     "refresh.webhook.configuration": refresh.webhook.configuration,
     "request.generator_definition.check": requests.generator_definition.check,
-    "request.graphql_query_group.update": requests.graphql_query_group.update,
     "request.artifact_definition.check": requests.artifact_definition.check,
     "request.proposed_change.data_integrity": requests.proposed_change.data_integrity,
     "request.proposed_change.pipeline": requests.proposed_change.pipeline,
diff --git a/backend/infrahub/message_bus/operations/requests/__init__.py b/backend/infrahub/message_bus/operations/requests/__init__.py
@@ -1,15 +1,13 @@
 from . import (
     artifact_definition,
     generator_definition,
-    graphql_query_group,
     proposed_change,
     repository,
 )
 
 __all__ = [
     "artifact_definition",
     "generator_definition",
-    "graphql_query_group",
     "proposed_change",
     "repository",
 ]
diff --git a/backend/infrahub/message_bus/operations/requests/graphql_query_group.py b/backend/infrahub/message_bus/operations/requests/graphql_query_group.py
diff --git a/backend/infrahub/workflows/catalogue.py b/backend/infrahub/workflows/catalogue.py
@@ -187,6 +187,14 @@
     function="cancel_proposed_changes_branch",
 )
 
+UPDATE_GRAPHQL_QUERY_GROUP = WorkflowDefinition(
+    name="update_graphql_query_group",
+    type=WorkflowType.INTERNAL,
+    module="infrahub.groups.tasks",
+    function="update_graphql_query_group",
+    branch_support=BranchSupportType.AWARE,
+)
+
 worker_pools = [INFRAHUB_WORKER_POOL]
 
 workflows = [
@@ -213,4 +221,5 @@
     TRIGGER_GENERATOR_DEFINITION_RUN,
     BRANCH_CANCEL_PROPOSED_CHANGES,
     REQUEST_GENERATOR_DEFINITION_RUN,
+    UPDATE_GRAPHQL_QUERY_GROUP,
 ]
diff --git a/backend/tests/helpers/utils.py b/backend/tests/helpers/utils.py
@@ -1,6 +1,12 @@
+from contextlib import contextmanager
+from typing import Generator
+
+from infrahub_sdk import InfrahubClient
 from testcontainers.core.container import DockerContainer
 from testcontainers.core.waiting_utils import wait_for_logs
 
+from infrahub.services import InfrahubServices, services
+from infrahub.services.adapters.workflow.local import WorkflowLocalExecution
 from tests.helpers.constants import PORT_BOLT_NEO4J, PORT_HTTP_NEO4J
 
 
@@ -29,3 +35,19 @@ def start_neo4j_container(neo4j_image: str) -> DockerContainer:
     container.start()
     wait_for_logs(container, "Started.")  # wait_container_is_ready does not seem to be enough
     return container
+
+
+@contextmanager
+def init_service_with_client(client: InfrahubClient) -> Generator:
+    """
+    This helper is needed for tests defining a specific client while `service` still needs to be accessed
+    through a global variable within prefect tasks.
+    """
+
+    original = services.service
+    service = InfrahubServices(client=client, workflow=WorkflowLocalExecution())
+    services.service = service
+    try:
+        yield service
+    finally:
+        services.service = original
diff --git a/backend/tests/unit/api/test_05_query_api.py b/backend/tests/unit/api/test_05_query_api.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING
+from unittest.mock import call, patch
 
 import pytest
 
 from infrahub.core.initialization import create_branch
-from infrahub.message_bus import messages
+from infrahub.groups.models import RequestGraphQLQueryGroupUpdate
+from infrahub.workflows.catalogue import UPDATE_GRAPHQL_QUERY_GROUP
 
 if TYPE_CHECKING:
     from fastapi.testclient import TestClient
@@ -32,70 +34,90 @@ async def test_query_endpoint_group_no_params(
     patch_services,
 ):
     # Must execute in a with block to execute the startup/shutdown events
-    with client:
+    with (
+        client,
+        patch(
+            "infrahub.services.adapters.workflow.local.WorkflowLocalExecution.submit_workflow"
+        ) as mock_submit_workflow,
+    ):
         response = client.get(
             "/api/query/query01?update_group=true&subscribers=AAAAAA&subscribers=BBBBBB", headers=admin_headers
         )
 
-    assert "errors" not in response.json()
-    assert response.status_code == 200
-    assert response.json()["data"] is not None
-    result = response.json()["data"]
+        assert "errors" not in response.json()
+        assert response.status_code == 200
+        assert response.json()["data"] is not None
+        result = response.json()["data"]
 
-    result_per_name = {result["node"]["name"]["value"]: result for result in result["TestPerson"]["edges"]}
-    assert sorted(result_per_name.keys()) == ["Jane", "John"]
-    assert len(result_per_name["John"]["node"]["cars"]["edges"]) == 2
-    assert len(result_per_name["Jane"]["node"]["cars"]["edges"]) == 1
+        result_per_name = {result["node"]["name"]["value"]: result for result in result["TestPerson"]["edges"]}
+        assert sorted(result_per_name.keys()) == ["Jane", "John"]
+        assert len(result_per_name["John"]["node"]["cars"]["edges"]) == 2
+        assert len(result_per_name["Jane"]["node"]["cars"]["edges"]) == 1
 
-    q1 = car_person_data["q1"]
-    p1 = car_person_data["p1"]
-    p2 = car_person_data["p2"]
-    c1 = car_person_data["c1"]
-    c2 = car_person_data["c2"]
-    c3 = car_person_data["c3"]
+        q1 = car_person_data["q1"]
+        p1 = car_person_data["p1"]
+        p2 = car_person_data["p2"]
+        c1 = car_person_data["c1"]
+        c2 = car_person_data["c2"]
+        c3 = car_person_data["c3"]
 
-    assert (
-        messages.RequestGraphQLQueryGroupUpdate(
+        model = RequestGraphQLQueryGroupUpdate(
             query_id=q1.id,
             query_name="query01",
             branch="main",
             related_node_ids=sorted([p1.id, p2.id, c1.id, c2.id, c3.id]),
             subscribers=sorted(["AAAAAA", "BBBBBB"]),
             params={},
         )
-        in client.app.state.service.message_bus.messages
-    )
+
+        expected_calls = [
+            call(
+                workflow=UPDATE_GRAPHQL_QUERY_GROUP,
+                parameters={"model": model},
+            ),
+        ]
+        mock_submit_workflow.assert_has_calls(expected_calls)
 
 
 async def test_query_endpoint_group_params(
     db: InfrahubDatabase, client: TestClient, admin_headers, default_branch, create_test_admin, car_person_data
 ):
     # Must execute in a with block to execute the startup/shutdown events
-    with client:
+    with (
+        client,
+        patch(
+            "infrahub.services.adapters.workflow.local.WorkflowLocalExecution.submit_workflow"
+        ) as mock_submit_workflow,
+    ):
         response = client.get("/api/query/query02?update_group=true&person=John", headers=admin_headers)
 
-    assert "errors" not in response.json()
-    assert response.status_code == 200
-    assert response.json()["data"] is not None
-    result = response.json()["data"]
+        assert "errors" not in response.json()
+        assert response.status_code == 200
+        assert response.json()["data"] is not None
+        result = response.json()["data"]
 
-    result_per_name = {result["node"]["name"]["value"]: result for result in result["TestPerson"]["edges"]}
-    assert sorted(result_per_name.keys()) == ["John"]
+        result_per_name = {result["node"]["name"]["value"]: result for result in result["TestPerson"]["edges"]}
+        assert sorted(result_per_name.keys()) == ["John"]
 
-    q2 = car_person_data["q2"]
-    p1 = car_person_data["p1"]
+        q2 = car_person_data["q2"]
+        p1 = car_person_data["p1"]
 
-    assert (
-        messages.RequestGraphQLQueryGroupUpdate(
+        model = RequestGraphQLQueryGroupUpdate(
             query_id=q2.id,
             query_name="query02",
             branch="main",
             related_node_ids={p1.id},
             subscribers=[],
             params={"person": "John"},
         )
-        in client.app.state.service.message_bus.messages
-    )
+
+        expected_calls = [
+            call(
+                workflow=UPDATE_GRAPHQL_QUERY_GROUP,
+                parameters={"model": model},
+            ),
+        ]
+        mock_submit_workflow.assert_has_calls(expected_calls)
 
 
 async def test_query_endpoint_get_default_branch(
diff --git a/backend/tests/unit/conftest.py b/backend/tests/unit/conftest.py
@@ -2943,9 +2943,6 @@ def workflow_local():
 @pytest.fixture
 def init_service(db: InfrahubDatabase):
     original = services.service
-    database = db
-    workflow = WorkflowLocalExecution()
-    service = InfrahubServices(database=database, workflow=workflow)
-    services.service = service
-    yield service
+    services.service = InfrahubServices(database=db, workflow=WorkflowLocalExecution())
+    yield services.service
     services.service = original
diff --git a/backend/tests/unit/message_bus/operations/requests/test_graphql_query_group.py b/backend/tests/unit/message_bus/operations/requests/test_graphql_query_group.py
diff --git a/docs/docs/reference/message-bus-events.mdx b/docs/docs/reference/message-bus-events.mdx
diff --git a/python_sdk b/python_sdk
