Add support for convert_query_response to Python Transforms

ogenstad · ogenstad · commit 42c68558ffbc · 2025-04-30T15:59:48.000+02:00
diff --git a/backend/infrahub/api/transformation.py b/backend/infrahub/api/transformation.py
@@ -88,6 +88,7 @@ async def transform_python(
         branch=branch_params.branch.name,
         transform_location=f"{transform.file_path.value}::{transform.class_name.value}",
         timeout=transform.timeout.value,
+        convert_query_response=transform.convert_query_response.value or False,
         data=data,
     )
 
diff --git a/backend/infrahub/artifacts/models.py b/backend/infrahub/artifacts/models.py
@@ -12,6 +12,10 @@ class CheckArtifactCreate(BaseModel):
     content_type: str = Field(..., description="Content type of the artifact")
     transform_type: str = Field(..., description="The type of transform associated with this artifact")
     transform_location: str = Field(..., description="The transforms location within the repository")
+    convert_query_response: bool = Field(
+        default=False,
+        description="Indicate if the query response should be converted to InfrahubNode objects for Python transforms",
+    )
     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")
diff --git a/backend/infrahub/computed_attribute/tasks.py b/backend/infrahub/computed_attribute/tasks.py
@@ -113,6 +113,7 @@ async def process_transform(
             location=f"{transform.file_path.value}::{transform.class_name.value}",
             data=data,
             client=service.client,
+            convert_query_response=transform.convert_query_response.value,
         )  # type: ignore[misc]
 
         await service.client.execute_graphql(
diff --git a/backend/infrahub/core/protocols.py b/backend/infrahub/core/protocols.py
@@ -478,6 +478,7 @@ class CoreTransformJinja2(CoreTransformation):
 class CoreTransformPython(CoreTransformation):
     file_path: String
     class_name: String
+    convert_query_response: BooleanOptional
 
 
 class CoreUserValidator(CoreValidator):
diff --git a/backend/infrahub/core/schema/definitions/core/transform.py b/backend/infrahub/core/schema/definitions/core/transform.py
@@ -92,5 +92,6 @@
     attributes=[
         Attr(name="file_path", kind="Text"),
         Attr(name="class_name", kind="Text"),
+        Attr(name="convert_query_response", kind="Boolean", optional=True, default_value=False),
     ],
 )
diff --git a/backend/infrahub/git/integrator.py b/backend/infrahub/git/integrator.py
@@ -10,6 +10,7 @@
 import yaml
 from infrahub_sdk import InfrahubClient  # noqa: TC002
 from infrahub_sdk.exceptions import ValidationError
+from infrahub_sdk.node import InfrahubNode
 from infrahub_sdk.protocols import (
     CoreArtifact,
     CoreArtifactDefinition,
@@ -53,7 +54,6 @@
     import types
 
     from infrahub_sdk.checks import InfrahubCheck
-    from infrahub_sdk.node import InfrahubNode
     from infrahub_sdk.schema.repository import InfrahubRepositoryArtifactDefinitionConfig
     from infrahub_sdk.transforms import InfrahubTransform
 
@@ -123,6 +123,10 @@ class TransformPythonInformation(BaseModel):
     timeout: int
     """Timeout for the function."""
 
+    convert_query_response: bool = Field(
+        ..., description="Indicate if the transform should convert the query response to InfrahubNode objects"
+    )
+
 
 class InfrahubRepositoryIntegrator(InfrahubRepositoryBase):
     """
@@ -874,6 +878,7 @@ async def get_python_transforms(
                     file_path=file_path,
                     query=str(graphql_query.id),
                     timeout=transform_class.timeout,
+                    convert_query_response=transform.convert_query_response,
                 )
             )
 
@@ -1005,6 +1010,7 @@ async def create_python_transform(
             "file_path": transform.file_path,
             "class_name": transform.class_name,
             "timeout": transform.timeout,
+            "convert_query_response": transform.convert_query_response,
         }
         create_payload = self.sdk.schema.generate_payload_create(
             schema=schema,
@@ -1028,6 +1034,9 @@ async def update_python_transform(
         if existing_transform.timeout.value != local_transform.timeout:
             existing_transform.timeout.value = local_transform.timeout
 
+        if existing_transform.convert_query_response.value != local_transform.convert_query_response:
+            existing_transform.convert_query_response.value = local_transform.convert_query_response
+
         await existing_transform.save()
 
     @classmethod
@@ -1038,6 +1047,7 @@ async def compare_python_transform(
             existing_transform.query.id != local_transform.query
             or existing_transform.file_path.value != local_transform.file_path
             or existing_transform.timeout.value != local_transform.timeout
+            or existing_transform.convert_query_response.value != local_transform.convert_query_response
         ):
             return False
         return True
@@ -1129,7 +1139,13 @@ async def execute_python_check(
 
     @task(name="python-transform-execute", task_run_name="Execute Python Transform", cache_policy=NONE)  # type: ignore[arg-type]
     async def execute_python_transform(
-        self, branch_name: str, commit: str, location: str, client: InfrahubClient, data: dict | None = None
+        self,
+        branch_name: str,
+        commit: str,
+        location: str,
+        client: InfrahubClient,
+        convert_query_response: bool,
+        data: dict | None = None,
     ) -> Any:
         """Execute A Python Transform stored in the repository."""
         log = get_run_logger()
@@ -1159,7 +1175,13 @@ async def execute_python_transform(
 
             transform_class: type[InfrahubTransform] = getattr(module, class_name)
 
-            transform = transform_class(root_directory=commit_worktree.directory, branch=branch_name, client=client)
+            transform = transform_class(
+                root_directory=commit_worktree.directory,
+                branch=branch_name,
+                client=client,
+                convert_query_response=convert_query_response,
+                infrahub_node=InfrahubNode,
+            )
             return await transform.run(data=data)
 
         except ModuleNotFoundError as exc:
@@ -1216,6 +1238,7 @@ async def artifact_generate(
                 location=transformation_location,
                 data=response,
                 client=self.sdk,
+                convert_query_response=transformation.convert_query_response.value,
             )  # type: ignore[misc]
 
         if definition.content_type.value == ContentType.APPLICATION_JSON.value and isinstance(artifact_content, dict):
@@ -1275,6 +1298,7 @@ async def render_artifact(
                 location=message.transform_location,
                 data=response,
                 client=self.sdk,
+                convert_query_response=message.convert_query_response,
             )  # type: ignore[misc]
 
         if message.content_type == ContentType.APPLICATION_JSON.value and isinstance(artifact_content, dict):
diff --git a/backend/infrahub/git/models.py b/backend/infrahub/git/models.py
@@ -29,6 +29,10 @@ class RequestArtifactGenerate(BaseModel):
     repository_name: str = Field(..., description="The name of the Repository")
     repository_kind: str = Field(..., description="The kind of the Repository")
     branch_name: str = Field(..., description="The branch where the check is run")
+    convert_query_response: bool = Field(
+        default=False,
+        description="Indicate if the query response should be converted to InfrahubNode objects for Python transforms",
+    )
     target_id: str = Field(..., description="The ID of the target object for this artifact")
     target_kind: str = Field(..., description="The kind of the target object for this artifact")
     target_name: str = Field(..., description="Name of the artifact target")
diff --git a/backend/infrahub/git/tasks.py b/backend/infrahub/git/tasks.py
@@ -339,10 +339,12 @@ async def generate_request_artifact_definition(
         )
     transform_location = ""
 
+    convert_query_response = False
     if transform.typename == InfrahubKind.TRANSFORMJINJA2:
         transform_location = transform.template_path.value
     elif transform.typename == InfrahubKind.TRANSFORMPYTHON:
         transform_location = f"{transform.file_path.value}::{transform.class_name.value}"
+        convert_query_response = transform.convert_query_response.value
 
     for relationship in group.members.peers:
         member = relationship.peer
@@ -368,6 +370,7 @@ async def generate_request_artifact_definition(
             target_name=member.display_label,
             target_kind=member.get_kind(),
             timeout=transform.timeout.value,
+            convert_query_response=convert_query_response,
             context=context,
         )
 
diff --git a/backend/infrahub/message_bus/operations/requests/proposed_change.py b/backend/infrahub/message_bus/operations/requests/proposed_change.py
@@ -319,6 +319,9 @@ async def refresh_artifacts(message: messages.RequestProposedChangeRefreshArtifa
               file_path {
                 value
               }
+              convert_query_response {
+                value
+              }
             }
             repository {
               node {
@@ -526,6 +529,9 @@ def _parse_artifact_definitions(definitions: list[dict]) -> list[ProposedChangeA
         elif artifact_definition.transform_kind == InfrahubKind.TRANSFORMPYTHON:
             artifact_definition.class_name = definition["node"]["transformation"]["node"]["class_name"]["value"]
             artifact_definition.file_path = definition["node"]["transformation"]["node"]["file_path"]["value"]
+            artifact_definition.convert_query_response = definition["node"]["transformation"]["node"][
+                "convert_query_response"
+            ]["value"]
 
         parsed.append(artifact_definition)
 
diff --git a/backend/infrahub/message_bus/types.py b/backend/infrahub/message_bus/types.py
@@ -96,6 +96,9 @@ class ProposedChangeArtifactDefinition(BaseModel):
     class_name: str = Field(default="")
     content_type: str
     file_path: str = Field(default="")
+    convert_query_response: bool = Field(
+        default=False, description="Convert query response to InfrahubNode objects for Python based transforms"
+    )
     timeout: int
 
     @property
diff --git a/backend/infrahub/proposed_change/tasks.py b/backend/infrahub/proposed_change/tasks.py
@@ -607,6 +607,7 @@ async def validate_artifacts_generation(model: RequestArtifactDefinitionCheck, s
                 content_type=model.artifact_definition.content_type,
                 transform_type=model.artifact_definition.transform_kind,
                 transform_location=model.artifact_definition.transform_location,
+                convert_query_response=model.artifact_definition.convert_query_response,
                 repository_id=repository.repository_id,
                 repository_name=repository.repository_name,
                 repository_kind=repository.kind,
diff --git a/backend/infrahub/transformations/models.py b/backend/infrahub/transformations/models.py
@@ -11,6 +11,9 @@ class TransformPythonData(BaseModel):
     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 generating the artifact")
+    convert_query_response: bool = Field(
+        ..., description="Define if the GraphQL query respose should be converted into InfrahubNode objects"
+    )
     timeout: int = Field(..., description="The timeout value to use when generating the artifact")
 
 
diff --git a/backend/infrahub/transformations/tasks.py b/backend/infrahub/transformations/tasks.py
@@ -30,6 +30,7 @@ async def transform_python(message: TransformPythonData, service: InfrahubServic
         location=message.transform_location,
         data=message.data,
         client=service.client,
+        convert_query_response=message.convert_query_response,
     )  # type: ignore[misc]
 
     return transformed_data
diff --git a/backend/infrahub/webhook/models.py b/backend/infrahub/webhook/models.py
@@ -204,6 +204,7 @@ class TransformWebhook(Webhook):
     transform_class: str = Field(...)
     transform_file: str = Field(...)
     transform_timeout: int = Field(...)
+    convert_query_response: bool = Field(...)
 
     async def _prepare_payload(self, data: dict[str, Any], context: EventContext, service: InfrahubServices) -> None:
         repo: InfrahubReadOnlyRepository | InfrahubRepository
@@ -229,6 +230,7 @@ async def _prepare_payload(self, data: dict[str, Any], context: EventContext, se
             branch_name=branch,
             commit=commit,
             location=f"{self.transform_file}::{self.transform_class}",
+            convert_query_response=self.convert_query_response,
             data={"data": data, **context.model_dump()},
             client=service.client,
         )  # type: ignore[misc]
@@ -247,4 +249,5 @@ def from_object(cls, obj: CoreCustomWebhook, transform: CoreTransformPython) ->
             transform_class=transform.class_name.value,
             transform_file=transform.file_path.value,
             transform_timeout=transform.timeout.value,
+            convert_query_response=transform.convert_query_response.value or False,
         )
diff --git a/backend/tests/fixtures/repos/car-dealership/initial__main/.infrahub.yml b/backend/tests/fixtures/repos/car-dealership/initial__main/.infrahub.yml
@@ -26,6 +26,10 @@ python_transforms:
   - name: PersonWithCarsTransform
     class_name: PersonWithCarsTransform
     file_path: "transforms/person_with_cars_transform.py"
+  - name: ConvertedPersonWithCarsTransform
+    class_name: ConvertedPersonWith
+    file_path: "transforms/converted_person_with_cars.py"
+    convert_query_response: true
   - name: CarSpecMarkdown
     class_name: CarSpecMarkdown
     file_path: "transforms/car_spec_markdown.py"
@@ -62,6 +66,13 @@ artifact_definitions:
     content_type: "text/markdown"
     targets: "people"
     transformation: "CarSpecMarkdown"
+  - name: "converted-owner"
+    artifact_name: "car-converted-owner"
+    parameters:
+      name: "name__value"
+    content_type: "application/json"
+    targets: "people"
+    transformation: "ConvertedPersonWithCarsTransform"
 
 generator_definitions:
   - name: cartags
diff --git a/backend/tests/fixtures/repos/car-dealership/initial__main/templates/person_with_cars.gql b/backend/tests/fixtures/repos/car-dealership/initial__main/templates/person_with_cars.gql
@@ -2,6 +2,8 @@ query PersonWithTheirCars($name: String!) {
   TestingPerson(name__value: $name) {
     edges {
       node {
+        id
+        __typename
         name {
           value
         }
@@ -11,6 +13,8 @@ query PersonWithTheirCars($name: String!) {
         cars {
           edges {
             node {
+              id
+              __typename
               name {
                 value
               }
diff --git a/backend/tests/fixtures/repos/car-dealership/initial__main/transforms/converted_person_with_cars.py b/backend/tests/fixtures/repos/car-dealership/initial__main/transforms/converted_person_with_cars.py
@@ -0,0 +1,14 @@
+from typing import Any
+
+from infrahub_sdk.transforms import InfrahubTransform
+
+
+class ConvertedPersonWith(InfrahubTransform):
+    query = "person_with_cars"
+
+    async def transform(self, data: dict[str, Any]) -> dict[str, Any]:
+        node_id = data["TestingPerson"]["edges"][0]["node"]["id"]
+
+        person = self.store.get(key=node_id, kind="TestingPerson")
+
+        return {"name": person.name.value, "age": person.age.value}
diff --git a/backend/tests/fixtures/schemas/schema_02.json b/backend/tests/fixtures/schemas/schema_02.json
@@ -3031,6 +3031,23 @@
                     "branch": "aware",
                     "optional": false,
                     "order_weight": 7000
+                },
+                {
+                    "name": "convert_query_response",
+                    "kind": "Boolean",
+                    "namespace": "Attribute",
+                    "label": null,
+                    "description": null,
+                    "default_value": false,
+                    "enum": null,
+                    "regex": null,
+                    "max_length": null,
+                    "min_length": null,
+                    "inherited": false,
+                    "unique": false,
+                    "branch": "aware",
+                    "optional": true,
+                    "order_weight": 8000
                 }
             ],
             "relationships": [
diff --git a/backend/tests/functional/webhook/test_task.py b/backend/tests/functional/webhook/test_task.py
@@ -257,6 +257,7 @@ async def test_convert_node_to_webhook_transform(
             "name": "Webhook2",
             "repository_kind": "CoreRepository",
             "repository_name": "car-dealership",
+            "convert_query_response": False,
             "transform_class": "WebhookTransformer",
             "transform_file": "transforms/webhook_transformer.py",
             "transform_name": "WebhookTransformer",
diff --git a/backend/tests/integration/git/test_readonly_repository.py b/backend/tests/integration/git/test_readonly_repository.py
diff --git a/backend/tests/integration/proposed_change/test_proposed_change_conflict.py b/backend/tests/integration/proposed_change/test_proposed_change_conflict.py
diff --git a/backend/tests/integration/transform/test_transform.py b/backend/tests/integration/transform/test_transform.py
diff --git a/backend/tests/unit/git/conftest.py b/backend/tests/unit/git/conftest.py
diff --git a/backend/tests/unit/git/test_git_repository.py b/backend/tests/unit/git/test_git_repository.py
diff --git a/backend/tests/unit/git/test_git_transform.py b/backend/tests/unit/git/test_git_transform.py
diff --git a/changelog/6383.added.md b/changelog/6383.added.md
diff --git a/docs/docs/reference/dotinfrahub.mdx b/docs/docs/reference/dotinfrahub.mdx
diff --git a/python_sdk b/python_sdk

Original file line number	Diff line number	Diff line change
`@@ -88,6 +88,7 @@ async def transform_python(`
`88`	`88`	`branch=branch_params.branch.name,`
`89`	`89`	`transform_location=f"{transform.file_path.value}::{transform.class_name.value}",`
`90`	`90`	`timeout=transform.timeout.value,`
	`91`	`+ convert_query_response=transform.convert_query_response.value or False,`
`91`	`92`	`data=data,`
`92`	`93`	`)`
`93`	`94`
Original file line number	Diff line number	Diff line change
`@@ -92,5 +92,6 @@`
`92`	`92`	`attributes=[`
`93`	`93`	`Attr(name="file_path", kind="Text"),`
`94`	`94`	`Attr(name="class_name", kind="Text"),`
	`95`	`+ Attr(name="convert_query_response", kind="Boolean", optional=True, default_value=False),`
`95`	`96`	`],`
`96`	`97`	`)`