Merge pull request #4918 from opsmill/pog-schema-updated-event-IFC-857

Migrate EventSchemaUpdate to InfrahubEvent

Author: ogenstad

backend/infrahub/api/schema.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, Any, Optional, Union
 
-from fastapi import APIRouter, BackgroundTasks, Depends, Query, Request
+from fastapi import APIRouter, Depends, Query, Request
 from pydantic import (
     BaseModel,
     Field,
@@ -12,7 +12,7 @@
 )
 from starlette.responses import JSONResponse
 
-from infrahub import config, lock
+from infrahub import lock
 from infrahub.api.dependencies import get_branch_dep, get_current_user, get_db
 from infrahub.api.exceptions import SchemaNotValidError
 from infrahub.core import registry
@@ -29,13 +29,13 @@
 from infrahub.core.schema.constants import SchemaNamespace  # noqa: TCH001
 from infrahub.core.validators.models.validate_migration import SchemaValidateMigrationData
 from infrahub.database import InfrahubDatabase  # noqa: TCH001
+from infrahub.events import EventMeta
+from infrahub.events.schema_action import SchemaUpdatedEvent
 from infrahub.exceptions import MigrationError, PermissionDeniedError
-from infrahub.log import get_logger
-from infrahub.message_bus import Meta, messages
-from infrahub.services import services
+from infrahub.log import get_log_data, get_logger
 from infrahub.types import ATTRIBUTE_PYTHON_TYPES
 from infrahub.worker import WORKER_IDENTITY
-from infrahub.workflows.catalogue import COMPUTED_ATTRIBUTE_SETUP, SCHEMA_APPLY_MIGRATION, SCHEMA_VALIDATE_MIGRATION
+from infrahub.workflows.catalogue import SCHEMA_APPLY_MIGRATION, SCHEMA_VALIDATE_MIGRATION
 
 if TYPE_CHECKING:
     from typing_extensions import Self
@@ -239,7 +239,6 @@ async def get_json_schema_by_kind(schema_kind: str, branch: Branch = Depends(get
 async def load_schema(
     request: Request,
     schemas: SchemasLoadAPI,
-    background_tasks: BackgroundTasks,
     db: InfrahubDatabase = Depends(get_db),
     branch: Branch = Depends(get_branch_dep),
     account_session: AccountSession = Depends(get_current_user),
@@ -339,18 +338,17 @@ async def load_schema(
         if migration_error_msgs:
             raise MigrationError(message=",\n".join(migration_error_msgs))
 
-        if config.SETTINGS.broker.enable:
-            message = messages.EventSchemaUpdate(
-                branch=branch.name,
-                meta=Meta(initiator_id=WORKER_IDENTITY),
-            )
-            background_tasks.add_task(services.send, message)
+    log_data = get_log_data()
+    request_id = log_data.get("request_id", "")
+    event = SchemaUpdatedEvent(
+        branch=branch.name,
+        schema_hash=branch.active_schema_hash.main,
+        meta=EventMeta(initiator_id=WORKER_IDENTITY, request_id=request_id, account_id=account_session.account_id),
+    )
+    await service.event.send(event=event)
 
     await service.component.refresh_schema_hash(branches=[branch.name])
 
-    # Temporary workaround, we need to emit a proper event
-    await service.workflow.submit_workflow(workflow=COMPUTED_ATTRIBUTE_SETUP)
-
     return SchemaUpdate(hash=updated_hash, previous_hash=original_hash, diff=result.diff)
 
 

backend/infrahub/events/schema_action.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+from pydantic import Field
+
+from infrahub.message_bus import InfrahubMessage
+from infrahub.message_bus.messages.event_schema_update import EventSchemaUpdate
+
+from .models import InfrahubBranchEvent
+
+
+class SchemaUpdatedEvent(InfrahubBranchEvent):
+    """Event generated when the schema within a branch has been updated."""
+
+    schema_hash: str = Field(..., description="Schema hash after the update")
+
+    def get_name(self) -> str:
+        return f"{self.get_event_namespace()}.schema.update"
+
+    def get_resource(self) -> dict[str, str]:
+        return {
+            "prefect.resource.id": f"infrahub.schema_branch.{self.branch}",
+            "infrahub.branch.name": self.branch,
+            "infrahub.branch.schema_hash": self.schema_hash,
+        }
+
+    def get_payload(self) -> dict[str, Any]:
+        return {"branch": self.branch, "schema_hash": self.schema_hash}
+
+    def get_messages(self) -> list[InfrahubMessage]:
+        return [
+            EventSchemaUpdate(
+                branch=self.branch,
+                meta=self.get_message_meta(),
+            )
+        ]

backend/infrahub/graphql/mutations/schema.py

@@ -1,24 +1,29 @@
-from typing import TYPE_CHECKING, Union
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Self, Union
 
 from graphene import Boolean, Field, InputObjectType, Mutation, String
-from graphql import GraphQLResolveInfo
 
-from infrahub import config, lock
+from infrahub import lock
 from infrahub.core import registry
-from infrahub.core.branch import Branch
 from infrahub.core.constants import RESTRICTED_NAMESPACES
 from infrahub.core.manager import NodeManager
 from infrahub.core.schema import DropdownChoice, GenericSchema, NodeSchema
 from infrahub.database import InfrahubDatabase, retry_db_transaction
+from infrahub.events import EventMeta
+from infrahub.events.schema_action import SchemaUpdatedEvent
 from infrahub.exceptions import ValidationError
-from infrahub.log import get_logger
-from infrahub.message_bus import Meta, messages
-from infrahub.services import services
+from infrahub.log import get_log_data, get_logger
 from infrahub.worker import WORKER_IDENTITY
 
 from ..types import DropdownFields
 
 if TYPE_CHECKING:
+    from graphql import GraphQLResolveInfo
+
+    from infrahub.core.branch import Branch
+    from infrahub.services import InfrahubServices
+
     from ..initialization import GraphqlContext
 
 log = get_logger()
@@ -56,7 +61,7 @@ async def mutate(
         root: dict,  # pylint: disable=unused-argument
         info: GraphQLResolveInfo,
         data: SchemaDropdownAddInput,
-    ):
+    ) -> Self:
         context: GraphqlContext = info.context
 
         kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
@@ -73,7 +78,13 @@ async def mutate(
                 )
             attrib.choices.append(choice)
 
-        await update_registry(kind=kind, branch=context.branch, db=context.db)
+        await update_registry(
+            kind=kind,
+            branch=context.branch,
+            db=context.db,
+            account_id=context.active_account_session.account_id,
+            service=context.active_service,
+        )
 
         kind = context.db.schema.get(name=str(data.kind), branch=context.branch.name)
         attrib = kind.get_attribute(attribute)
@@ -130,7 +141,13 @@ async def mutate(
                 raise ValidationError(f"Unable to remove the last dropdown on {kind.kind} in attribute {attribute}")
             attrib.choices = [entry for entry in attrib.choices if dropdown != entry.name]
 
-        await update_registry(kind=kind, branch=context.branch, db=context.db)
+        await update_registry(
+            kind=kind,
+            branch=context.branch,
+            db=context.db,
+            account_id=context.active_account_session.account_id,
+            service=context.active_service,
+        )
 
         return {"ok": True}
 
@@ -165,7 +182,13 @@ async def mutate(
                     )
                 attrib.enum.append(enum)
 
-        await update_registry(kind=kind, branch=context.branch, db=context.db)
+        await update_registry(
+            kind=kind,
+            branch=context.branch,
+            db=context.db,
+            account_id=context.active_account_session.account_id,
+            service=context.active_service,
+        )
 
         return {"ok": True}
 
@@ -206,7 +229,13 @@ async def mutate(
                     raise ValidationError(f"Unable to remove the last enum on {kind.kind} in attribute {attribute}")
                 attrib.enum = [entry for entry in attrib.enum if entry != enum]
 
-        await update_registry(kind=kind, branch=branch, db=db)
+        await update_registry(
+            kind=kind,
+            branch=branch,
+            db=db,
+            account_id=context.active_account_session.account_id,
+            service=context.active_service,
+        )
 
         return {"ok": True}
 
@@ -237,7 +266,9 @@ def validate_kind(kind: Union[GenericSchema, NodeSchema], attribute: str) -> Non
         raise ValidationError(f"Attribute {attribute} on {kind.kind} is inherited and must be changed on the generic")
 
 
-async def update_registry(kind: NodeSchema, branch: Branch, db: InfrahubDatabase) -> None:
+async def update_registry(
+    kind: NodeSchema, db: InfrahubDatabase, branch: Branch, account_id: str, service: InfrahubServices
+) -> None:
     async with lock.registry.global_schema_lock():
         branch_schema = registry.schema.get_schema_branch(name=branch.name)
 
@@ -256,13 +287,20 @@ async def update_registry(kind: NodeSchema, branch: Branch, db: InfrahubDatabase
                     schema=tmp_schema, db=dbt, branch=branch.name, limit=diff.all, update_db=True
                 )
                 branch.update_schema_hash()
-                log.info("Schema has been updated", branch=branch.name, hash=branch.schema_hash.main)
+                log.info("Schema has been updated", branch=branch.name, hash=branch.active_schema_hash.main)
                 await branch.save(db=dbt)
 
-            if config.SETTINGS.broker.enable:
-                message = messages.EventSchemaUpdate(
-                    branch=branch.name,
-                    meta=Meta(initiator_id=WORKER_IDENTITY),
-                )
-                await services.send(message)
-            await services.service.component.refresh_schema_hash(branches=[branch.name])
+            log_data = get_log_data()
+            request_id = log_data.get("request_id", "")
+            event = SchemaUpdatedEvent(
+                branch=branch.name,
+                schema_hash=branch.active_schema_hash.main,
+                meta=EventMeta(
+                    initiator_id=WORKER_IDENTITY,
+                    request_id=request_id,
+                    account_id=account_id,
+                ),
+            )
+            await service.event.send(event=event)
+
+            await service.component.refresh_schema_hash(branches=[branch.name])

backend/infrahub/schema/__init__.py

@@ -0,0 +1 @@
+AUTOMATION_NAME = "Trigger-schema-update-event"

backend/infrahub/schema/constants.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+from datetime import timedelta
+
+from prefect import flow
+from prefect.automations import AutomationCore
+from prefect.client.orchestration import get_client
+from prefect.client.schemas.filters import DeploymentFilter, DeploymentFilterName
+from prefect.events.actions import RunDeployment
+from prefect.events.schemas.automations import EventTrigger, Posture
+from prefect.logging import get_run_logger
+
+from infrahub.workflows.catalogue import COMPUTED_ATTRIBUTE_SETUP
+
+from .constants import AUTOMATION_NAME
+
+
+@flow(name="schema-updated-setup", flow_run_name="Setup schema updated event in task-manager")
+async def schema_updated_setup() -> None:
+    log = get_run_logger()
+
+    async with get_client(sync_client=False) as client:
+        deployments = {
+            item.name: item
+            for item in await client.read_deployments(
+                deployment_filter=DeploymentFilter(
+                    name=DeploymentFilterName(
+                        any_=[
+                            COMPUTED_ATTRIBUTE_SETUP.name,
+                        ]
+                    )
+                )
+            )
+        }
+        if COMPUTED_ATTRIBUTE_SETUP.name not in deployments:
+            raise ValueError("Unable to find the deployment for PROCESS_COMPUTED_MACRO")
+
+        deployment_id_computed_attribute_setup = deployments[COMPUTED_ATTRIBUTE_SETUP.name].id
+
+        schema_update_automation = await client.find_automation(id_or_name=AUTOMATION_NAME)
+
+        automation = AutomationCore(
+            name=AUTOMATION_NAME,
+            description="Trigger actions on schema update event",
+            enabled=True,
+            trigger=EventTrigger(
+                posture=Posture.Reactive,
+                expect={"infrahub.schema.update"},
+                within=timedelta(0),
+                threshold=1,
+            ),
+            actions=[
+                RunDeployment(
+                    source="selected",
+                    deployment_id=deployment_id_computed_attribute_setup,
+                    parameters={},
+                    job_variables={},
+                )
+            ],
+        )
+
+        if schema_update_automation:
+            await client.update_automation(automation_id=schema_update_automation.id, automation=automation)
+            log.info(f"{AUTOMATION_NAME} Updated")
+        else:
+            await client.create_automation(automation=automation)
+            log.info(f"{AUTOMATION_NAME} Created")

backend/infrahub/schema/tasks.py

@@ -2,7 +2,7 @@
 
 from typing import TYPE_CHECKING, Any, overload
 
-from prefect.client.schemas import StateType
+from prefect.client.schemas.objects import StateType
 from prefect.deployments import run_deployment
 
 from infrahub.workflows.initialization import setup_task_manager

backend/infrahub/services/adapters/workflow/worker.py

@@ -247,6 +247,13 @@
     function="run_proposed_change_data_integrity_check",
 )
 
+SCHEMA_UPDATED_SETUP = WorkflowDefinition(
+    name="schema-updated-setup",
+    type=WorkflowType.INTERNAL,
+    module="infrahub.schema.tasks",
+    function="schema_updated_setup",
+)
+
 
 worker_pools = [INFRAHUB_WORKER_POOL]
 
@@ -284,4 +291,5 @@
     COMPUTED_ATTRIBUTE_SETUP,
     UPDATE_COMPUTED_ATTRIBUTE_TRANSFORM,
     REQUEST_PROPOSED_CHANGE_DATA_INTEGRITY,
+    SCHEMA_UPDATED_SETUP,
 ]

backend/infrahub/workflows/catalogue.py

@@ -7,7 +7,7 @@
 
 from infrahub import config
 
-from .catalogue import worker_pools, workflows
+from .catalogue import SCHEMA_UPDATED_SETUP, worker_pools, workflows
 from .models import TASK_RESULT_STORAGE_NAME
 
 
@@ -37,6 +37,9 @@ async def setup_deployments(client: PrefectClient) -> None:
         await workflow.save(client=client, work_pool=work_pool)
         log.info(f"Flow {workflow.name}, created successfully ... ")
 
+    schema_update_setup = SCHEMA_UPDATED_SETUP.get_function()
+    await schema_update_setup()
+
 
 @task(name="task-manager-setup-blocks", task_run_name="Setup Blocks")
 async def setup_blocks() -> None:

backend/infrahub/workflows/initialization.py

@@ -330,7 +330,12 @@ ignore_errors = true
 
 [[tool.mypy.overrides]]
 module = "infrahub.graphql.mutations.schema"
-ignore_errors = true
+disable_error_code = [
+    "arg-type",
+    "operator",
+    "union-attr",
+]
+
 
 [[tool.mypy.overrides]]
 module = "infrahub.graphql.resolver"