Merge pull request #4652 from opsmill/dga-20241017-convert-transform

Convert Transform tasks to prefect

Author: dgarros

backend/infrahub/api/transformation.py

@@ -23,12 +23,8 @@
 from infrahub.exceptions import TransformError
 from infrahub.graphql.initialization import prepare_graphql_params
 from infrahub.graphql.utils import extract_data
-from infrahub.message_bus.messages import (
-    TransformPythonData,
-    TransformPythonDataResponse,
-)
-from infrahub.message_bus.messages.transform_jinja_template import TransformJinjaTemplateData
-from infrahub.workflows.catalogue import TRANSFORM_JINJA2_RENDER
+from infrahub.transformations.models import TransformJinjaTemplateData, TransformPythonData
+from infrahub.workflows.catalogue import TRANSFORM_JINJA2_RENDER, TRANSFORM_PYTHON_RENDER
 
 if TYPE_CHECKING:
     from infrahub.services import InfrahubServices
@@ -89,8 +85,10 @@ async def transform_python(
         data=data,
     )
 
-    response = await service.message_bus.rpc(message=message, response_class=TransformPythonDataResponse)
-    return JSONResponse(content=response.data.transformed_data)
+    response = await service.workflow.execute_workflow(
+        workflow=TRANSFORM_PYTHON_RENDER, parameters={"message": message}
+    )
+    return JSONResponse(content=response)
 
 
 @router.get("/transform/jinja2/{transform_id}", response_class=PlainTextResponse)

backend/infrahub/message_bus/messages/__init__.py

@@ -49,8 +49,6 @@
 from .send_echo_request import SendEchoRequest, SendEchoRequestResponse
 from .send_telemetry_push import SendTelemetryPush
 from .send_webhook_event import SendWebhookEvent
-from .transform_jinja_template import TransformJinjaTemplate, TransformJinjaTemplateResponse
-from .transform_python_data import TransformPythonData, TransformPythonDataResponse
 from .trigger_artifact_definition_generate import TriggerArtifactDefinitionGenerate
 from .trigger_generatordefinition_run import TriggerGeneratorDefinitionRun
 from .trigger_proposed_change_cancel import TriggerProposedChangeCancel
@@ -106,17 +104,13 @@
     "send.echo.request": SendEchoRequest,
     "send.webhook.event": SendWebhookEvent,
     "send.telemetry.push": SendTelemetryPush,
-    "transform.jinja.template": TransformJinjaTemplate,
-    "transform.python.data": TransformPythonData,
     "trigger.artifact_definition.generate": TriggerArtifactDefinitionGenerate,
     "trigger.generator_definition.run": TriggerGeneratorDefinitionRun,
     "trigger.proposed_change.cancel": TriggerProposedChangeCancel,
     "trigger.webhook.actions": TriggerWebhookActions,
 }
 
 RESPONSE_MAP: dict[str, type[InfrahubResponse]] = {
-    "transform.jinja.template": TransformJinjaTemplateResponse,
-    "transform.python.data": TransformPythonDataResponse,
     "git.diff.names_only": GitDiffNamesOnlyResponse,
     "git.file.get": GitFileGetResponse,
     "send.echo.request": SendEchoRequestResponse,

backend/infrahub/message_bus/messages/transform_python_data.py

@@ -13,7 +13,6 @@
     requests,
     schema,
     send,
-    transform,
     trigger,
 )
 from infrahub.message_bus.types import MessageTTL
@@ -70,8 +69,6 @@
     "send.telemetry.push": send.telemetry.push,
     "schema.migration.path": schema.migration.path,
     "schema.validator.path": schema.validator.path,
-    "transform.jinja.template": transform.jinja.template,
-    "transform.python.data": transform.python.data,
     "trigger.artifact_definition.generate": trigger.artifact_definition.generate,
     "trigger.generator_definition.run": trigger.generator_definition.run,
     "trigger.proposed_change.cancel": trigger.proposed_change.cancel,

backend/infrahub/message_bus/operations/__init__.py

@@ -1,10 +1,16 @@
-from typing import Optional
-
 from pydantic import BaseModel, Field
 
-from infrahub.message_bus import InfrahubMessage, InfrahubResponse, InfrahubResponseData
 
-ROUTING_KEY = "transform.jinja.template"
+class TransformPythonData(BaseModel):
+    """Sent to run a Python transform."""
+
+    repository_id: str = Field(..., description="The unique ID of the Repository")
+    repository_name: str = Field(..., description="The name of the repository")
+    repository_kind: str = Field(..., description="The kind of the repository")
+    data: dict = Field(..., description="Input data for the template")
+    branch: str = Field(..., description="The branch to target")
+    transform_location: str = Field(..., description="Location of the transform within the repository")
+    commit: str = Field(..., description="The commit id to use when rendering the template")
 
 
 class TransformJinjaTemplateData(BaseModel):
@@ -17,16 +23,3 @@ class TransformJinjaTemplateData(BaseModel):
     branch: str = Field(..., description="The branch to target")
     template_location: str = Field(..., description="Location of the template within the repository")
     commit: str = Field(..., description="The commit id to use when rendering the template")
-
-
-class TransformJinjaTemplate(TransformJinjaTemplateData, InfrahubMessage):
-    """Sent to trigger the checks for a repository to be executed."""
-
-
-class TransformJinjaTemplateResponseData(InfrahubResponseData):
-    rendered_template: Optional[str] = Field(None, description="Rendered template in string format")
-
-
-class TransformJinjaTemplateResponse(InfrahubResponse):
-    routing_key: str = ROUTING_KEY
-    data: TransformJinjaTemplateResponseData

backend/infrahub/message_bus/operations/transform/__init__.py

@@ -0,0 +1,51 @@
+from typing import Any
+
+from prefect import flow
+
+from infrahub.git.repository import get_initialized_repo
+from infrahub.log import get_logger
+from infrahub.services import services
+
+from .models import TransformJinjaTemplateData, TransformPythonData
+
+log = get_logger()
+
+
+@flow(name="transform-render-python")
+async def transform_python(message: TransformPythonData) -> Any:
+    service = services.service
+
+    repo = await get_initialized_repo(
+        repository_id=message.repository_id,
+        name=message.repository_name,
+        service=service,
+        repository_kind=message.repository_kind,
+    )
+
+    transformed_data = await repo.execute_python_transform(
+        branch_name=message.branch,
+        commit=message.commit,
+        location=message.transform_location,
+        data=message.data,
+        client=service.client,
+    )
+
+    return transformed_data
+
+
+@flow(name="transform-render-jinja2")
+async def transform_render_jinja2_template(message: TransformJinjaTemplateData) -> str:
+    service = services.service
+
+    repo = await get_initialized_repo(
+        repository_id=message.repository_id,
+        name=message.repository_name,
+        service=service,
+        repository_kind=message.repository_kind,
+    )
+
+    rendered_template = await repo.render_jinja2_template(
+        commit=message.commit, location=message.template_location, data={"data": message.data}
+    )
+
+    return rendered_template