Merge pull request #4902 from opsmill/dga-20241109-fix-computed-attr-serialization

Fix serialization for ComputedAttribute in schema

Author: dgarros

backend/infrahub/core/schema/attribute_schema.py

@@ -92,6 +92,19 @@ def update_from_generic(self, other: AttributeSchema) -> None:
             if getattr(self, name) != getattr(other, name):
                 setattr(self, name, getattr(other, name))
 
+    def to_node(self) -> dict[str, Any]:
+        fields_to_exclude = {"id", "state", "filters"}
+        fields_to_json = {"computed_attribute"}
+        data = self.model_dump(exclude=fields_to_exclude | fields_to_json)
+
+        for field_name in fields_to_json:
+            if field := getattr(self, field_name):
+                data[field_name] = {"value": field.model_dump()}
+            else:
+                data[field_name] = None
+
+        return data
+
     async def get_query_filter(
         self,
         name: str,

backend/infrahub/core/schema/computed_attribute.py

@@ -1,6 +1,6 @@
-from typing import Optional
+from typing import Any, Optional
 
-from pydantic import ConfigDict, Field
+from pydantic import ConfigDict, Field, model_serializer
 
 from infrahub.core.constants import ComputedAttributeKind
 from infrahub.core.models import HashableModel
@@ -15,6 +15,10 @@ class ComputedAttribute(HashableModel):
         default=None, description="The Python Transform name or ID, required when assignment_type=transform"
     )
 
+    @model_serializer
+    def ser_model(self) -> dict[str, Any]:
+        return {"kind": self.kind.value, "jinja2_template": self.jinja2_template, "transform": self.transform}
+
     model_config = ConfigDict(
         extra="forbid",
         json_schema_extra={

backend/infrahub/core/schema/manager.py

@@ -537,7 +537,7 @@ async def create_attribute_in_db(
         schema: NodeSchema, item: AttributeSchema, branch: Branch, parent: Node, db: InfrahubDatabase
     ) -> AttributeSchema:
         obj = await Node.init(schema=schema, branch=branch, db=db)
-        await obj.new(**item.model_dump(exclude={"id", "state", "filters"}), node=parent, db=db)
+        await obj.new(**item.to_node(), node=parent, db=db)
         await obj.save(db=db)
         new_item = item.duplicate()
         new_item.id = obj.id

backend/tests/helpers/schema/__init__.py

@@ -20,15 +20,17 @@
 CAR_SCHEMA = SchemaRoot(nodes=[CAR, MANUFACTURER, PERSON])
 
 
-async def load_schema(db: InfrahubDatabase, schema: SchemaRoot, branch_name: str | None = None) -> None:
+async def load_schema(
+    db: InfrahubDatabase, schema: SchemaRoot, branch_name: str | None = None, update_db: bool = False
+) -> None:
     default_branch_name = registry.default_branch
     branch_schema = registry.schema.get_schema_branch(name=branch_name or default_branch_name)
     tmp_schema = branch_schema.duplicate()
     tmp_schema.load_schema(schema=schema)
     tmp_schema.process()
 
     await registry.schema.update_schema_branch(
-        schema=tmp_schema, db=db, branch=branch_name or default_branch_name, update_db=True
+        schema=tmp_schema, db=db, branch=branch_name or default_branch_name, update_db=update_db
     )
 
 

backend/tests/integration/diff/test_diff_rebase.py

@@ -78,7 +78,7 @@ async def initial_dataset(
         client: InfrahubClient,
         bus_simulator: BusSimulator,
     ) -> dict[str, Node]:
-        await load_schema(db, schema=CAR_SCHEMA)
+        await load_schema(db, schema=CAR_SCHEMA, update_db=True)
         john = await Node.init(schema=TestKind.PERSON, db=db)
         await john.new(db=db, name="John", height=175, description="The famous Joe Doe")
         await john.save(db=db)

backend/tests/unit/core/schema_manager/test_manager_schema.py

@@ -18,16 +18,20 @@
     SchemaPathType,
 )
 from infrahub.core.schema import (
+    AttributeSchema,
     GenericSchema,
     NodeSchema,
+    RelationshipSchema,
     SchemaRoot,
     core_models,
     internal_schema,
 )
+from infrahub.core.schema.computed_attribute import ComputedAttribute
 from infrahub.core.schema.manager import SchemaManager
 from infrahub.core.schema.schema_branch import SchemaBranch
 from infrahub.database import InfrahubDatabase
 from infrahub.exceptions import SchemaNotFoundError, ValidationError
+from tests.helpers.schema import CHILD, THING
 
 from .conftest import _get_schema_by_kind
 
@@ -1416,10 +1420,11 @@ async def test_schema_branch_validate_count_against_cardinality_invalid(relation
 
 
 async def test_schema_branch_from_dict_schema_object():
-    schema = SchemaRoot(**core_models)
-
     schema_branch = SchemaBranch(cache={}, name="test")
-    schema_branch.load_schema(schema=schema)
+
+    # Load the core models and a model with a computed_attribute
+    schema_branch.load_schema(schema=SchemaRoot(**core_models))
+    schema_branch.load_schema(schema=SchemaRoot(nodes=[CHILD, THING]))
 
     exported = schema_branch.to_dict_schema_object()
 
@@ -2073,31 +2078,32 @@ async def test_load_node_to_db_node_schema(db: InfrahubDatabase, default_branch:
     registry.schema = SchemaManager()
     registry.schema.register_schema(schema=SchemaRoot(**internal_schema), branch=default_branch.name)
 
-    SCHEMA = {
-        "name": "Criticality",
-        "namespace": "Builtin",
-        "default_filter": "name__value",
-        "attributes": [
-            {"name": "name", "kind": "Text", "unique": True},
-            {"name": "level", "kind": "Number"},
-            {"name": "color", "kind": "Text", "default_value": "#444444"},
-            {"name": "description", "kind": "Text", "optional": True},
-        ],
-        "relationships": [
-            {"name": "others", "peer": "BuiltinCriticality", "optional": True, "cardinality": "many"},
+    node = NodeSchema(
+        name="Criticality",
+        namespace="Builtin",
+        default_filter="name__value",
+        attributes=[
+            AttributeSchema(name="name", kind="Text", unique=True),
+            AttributeSchema(name="level", kind="Number"),
+            AttributeSchema(name="color", kind="Text", default_value="default_value"),
+            AttributeSchema(
+                name="description",
+                kind="Text",
+                optional=True,
+                computed_attribute=ComputedAttribute(kind="Jinja2", jinja2_template="{{ name__value }}"),
+            ),
         ],
-    }
-    node = NodeSchema(**SCHEMA)
-
+        relationships=[RelationshipSchema(name="others", peer="BuiltinCriticality", optional=True, cardinality="many")],
+    )
     await registry.schema.load_node_to_db(node=node, db=db, branch=default_branch)
 
     node2 = registry.schema.get(name=node.kind, branch=default_branch)
     assert node2.id
     assert node2.relationships[0].id
     assert node2.attributes[0].id
 
-    results = await SchemaManager.query(schema="SchemaNode", branch=default_branch, db=db)
-    assert len(results) == 1
+    node_from_db = await SchemaManager.get_one(db=db, id=node2.id, branch=default_branch)
+    assert node_from_db
 
 
 async def test_load_node_to_db_generic_schema(db: InfrahubDatabase, default_branch):
@@ -2249,7 +2255,14 @@ async def test_load_schema_from_db(
                     {"name": "name", "kind": "Text", "label": "Name", "unique": True},
                     {"name": "level", "kind": "Number", "label": "Level"},
                     {"name": "color", "kind": "Text", "label": "Color", "default_value": "#444444"},
-                    {"name": "description", "kind": "Text", "label": "Description", "optional": True},
+                    {
+                        "name": "description",
+                        "kind": "Text",
+                        "label": "Description",
+                        "optional": True,
+                        "read_only": True,
+                        "computed_attribute": {"kind": "Jinja2", "jinja2_template": "{{ name__value }}"},
+                    },
                 ],
                 "relationships": [
                     {
@@ -2320,6 +2333,10 @@ async def test_load_schema_from_db(
     assert schema11.get(name=InfrahubKind.TAG).get_hash() == schema2.get(name="BuiltinTag").get_hash()
     assert schema11.get(name="TestGenericInterface").get_hash() == schema2.get(name="TestGenericInterface").get_hash()
 
+    description_schema = crit_schema.get_attribute("description")
+    assert description_schema.computed_attribute is not None
+    assert description_schema.computed_attribute.jinja2_template == "{{ name__value }}"
+
 
 async def test_load_schema(
     db: InfrahubDatabase, reset_registry, default_branch: Branch, register_internal_models_schema

backend/tests/unit/core/test_node.py

@@ -445,7 +445,17 @@ async def test_node_create_local_attrs(db: InfrahubDatabase, default_branch: Bra
     assert obj.json_no_default.value is None
 
     obj = await Node.init(db=db, schema=criticality_schema)
-    await obj.new(db=db, name="medium", level=3, description="My desc", is_true=False, is_false=True, color="#333333")
+    await obj.new(
+        db=db,
+        name="medium",
+        level=3,
+        description="My desc",
+        is_true=False,
+        is_false=True,
+        color="#333333",
+        json_default={"value": {"value": "xxxxx"}},
+        json_no_default={"value": {"testing": True}},
+    )
     await obj.save(db=db)
 
     assert obj.id
@@ -460,6 +470,8 @@ async def test_node_create_local_attrs(db: InfrahubDatabase, default_branch: Bra
     assert obj.color.id
     assert obj.is_true.value is False
     assert obj.is_false.value is True
+    assert obj.json_default.value == {"value": "xxxxx"}
+    assert obj.json_no_default.value == {"testing": True}
 
 
 async def test_node_create_attribute_with_source(

backend/tests/unit/graphql/test_mutation_create_macros.py

@@ -24,9 +24,7 @@ async def test_create_related_macro(
     default_branch.update_schema_hash()
     await default_branch.save(db=db)
     await load_schema(db, schema=SchemaRoot(nodes=[CHILD, THING]))
-    # FIXME: This is a workaround hack as it's not posible to have Schema attributes in JSON format
-    # on creation at the moment so we need to get the schema updated after initial creation..
-    await load_schema(db, schema=SchemaRoot(nodes=[CHILD, THING]))
+
     fred = await Node.init(schema=TestKind.CHILD, db=db)
     await fred.new(db=db, name="Fred", height=110)
     await fred.save(db=db)