fix(backend): lock before beginning transactions

WIP: still need to fix the uniqueness constraint lock on update.

Signed-off-by: Fatih Acar <[email protected]>

Author: fatih-acar

backend/infrahub/core/node/__init__.py

@@ -257,7 +257,13 @@ async def init(
 
         return cls(**attrs)
 
-    async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, errors: list) -> None:
+    async def handle_pool(
+        self,
+        db: InfrahubDatabase,
+        attribute: BaseAttribute,
+        errors: list,
+        allocate_resources: bool = True,
+    ) -> None:
         """Evaluate if a resource has been requested from a pool and apply the resource
 
         This method only works on number pools, currently Integer is the only type that has the from_pool
@@ -268,7 +274,7 @@ async def handle_pool(self, db: InfrahubDatabase, attribute: BaseAttribute, erro
             attribute.from_pool = {"id": attribute.schema.parameters.number_pool_id}
             attribute.is_default = False
 
-        if not attribute.from_pool:
+        if not attribute.from_pool or not allocate_resources:
             return
 
         try:
@@ -426,7 +432,12 @@ async def handle_object_template(self, fields: dict, db: InfrahubDatabase, error
             elif relationship_peers := await relationship.get_peers(db=db):
                 fields[relationship_name] = [{"id": peer_id} for peer_id in relationship_peers]
 
-    async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
+    async def _process_fields(
+        self,
+        fields: dict,
+        db: InfrahubDatabase,
+        process_pools: bool = True,
+    ) -> None:
         errors = []
 
         if "_source" in fields.keys():
@@ -480,7 +491,7 @@ async def _process_fields(self, fields: dict, db: InfrahubDatabase) -> None:
         # Generate Attribute and Relationship and assign them
         # -------------------------------------------
         errors.extend(await self._process_fields_relationships(fields=fields, db=db))
-        errors.extend(await self._process_fields_attributes(fields=fields, db=db))
+        errors.extend(await self._process_fields_attributes(fields=fields, db=db, process_pools=process_pools))
 
         if errors:
             raise ValidationError(errors)
@@ -517,7 +528,12 @@ async def _process_fields_relationships(self, fields: dict, db: InfrahubDatabase
 
         return errors
 
-    async def _process_fields_attributes(self, fields: dict, db: InfrahubDatabase) -> list[ValidationError]:
+    async def _process_fields_attributes(
+        self,
+        fields: dict,
+        db: InfrahubDatabase,
+        process_pools: bool,
+    ) -> list[ValidationError]:
         errors: list[ValidationError] = []
 
         for attr_schema in self._schema.attributes:
@@ -542,9 +558,15 @@ async def _process_fields_attributes(self, fields: dict, db: InfrahubDatabase) -
                 )
                 if not self._existing:
                     attribute: BaseAttribute = getattr(self, attr_schema.name)
-                    await self.handle_pool(db=db, attribute=attribute, errors=errors)
+                    await self.handle_pool(
+                        db=db,
+                        attribute=attribute,
+                        errors=errors,
+                        allocate_resources=process_pools,
+                    )
 
-                    attribute.validate(value=attribute.value, name=attribute.name, schema=attribute.schema)
+                    if process_pools or attribute.schema.kind != "NumberPool" or attribute.from_pool is None:
+                        attribute.validate(value=attribute.value, name=attribute.name, schema=attribute.schema)
             except ValidationError as exc:
                 errors.append(exc)
 
@@ -672,7 +694,13 @@ async def process_label(self, db: InfrahubDatabase | None = None) -> None:  # no
                 self.label.value = " ".join([word.title() for word in self.name.value.split("_")])
                 self.label.is_default = False
 
-    async def new(self, db: InfrahubDatabase, id: str | None = None, **kwargs: Any) -> Self:
+    async def new(
+        self,
+        db: InfrahubDatabase,
+        id: str | None = None,
+        process_pools: bool = True,
+        **kwargs: Any,
+    ) -> Self:
         if id and not is_valid_uuid(id):
             raise ValidationError({"id": f"{id} is not a valid UUID"})
         if id:
@@ -682,7 +710,7 @@ async def new(self, db: InfrahubDatabase, id: str | None = None, **kwargs: Any)
 
         self.id = id or str(UUIDT())
 
-        await self._process_fields(db=db, fields=kwargs)
+        await self._process_fields(db=db, fields=kwargs, process_pools=process_pools)
         await self._process_macros(db=db)
 
         return self

backend/infrahub/core/node/create.py

@@ -1,7 +1,8 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Mapping
+from typing import TYPE_CHECKING, Any, Mapping, Sequence
 
+from infrahub import lock
 from infrahub.core import registry
 from infrahub.core.constants import RelationshipCardinality, RelationshipKind
 from infrahub.core.constraint.node.runner import NodeConstraintRunner
@@ -10,6 +11,8 @@
 from infrahub.core.node.save import run_constraints_and_save
 from infrahub.core.protocols import CoreObjectTemplate
 from infrahub.dependencies.registry import get_component_registry
+from infrahub.lock import InfrahubMultiLock
+from infrahub.lock_getter import get_lock_names_on_object_mutation
 
 if TYPE_CHECKING:
     from infrahub.core.branch import Branch
@@ -142,24 +145,22 @@ async def refresh_for_profile_update(
 
 
 async def _do_create_node(
-    node_class: type[Node],
+    obj: Node,
     db: InfrahubDatabase,
-    data: dict,
-    schema: NonGenericSchemaTypes,
-    fields_to_validate: list,
     branch: Branch,
+    fields_to_validate: list[str],
     node_constraint_runner: NodeConstraintRunner,
+    lock_names: Sequence[str],
 ) -> Node:
-    obj = await node_class.init(db=db, schema=schema, branch=branch)
-    await obj.new(db=db, **data)
-
     await run_constraints_and_save(
         node=obj,
         node_constraint_runner=node_constraint_runner,
         fields_to_validate=fields_to_validate,
         fields_to_save=[],
         db=db,
         branch=branch,
+        lock_names=lock_names,
+        manage_lock=False,
     )
 
     object_template = await obj.get_object_template(db=db)
@@ -170,6 +171,7 @@ async def _do_create_node(
             template=object_template,
             obj=obj,
             fields=fields_to_validate,
+            constraint_runner=node_constraint_runner,
         )
     return obj
 
@@ -183,35 +185,39 @@ async def create_node(
     """Create a node in the database if constraint checks succeed."""
 
     component_registry = get_component_registry()
-    node_constraint_runner = await component_registry.get_component(
-        NodeConstraintRunner, db=db.start_session() if not db.is_transaction else db, branch=branch
-    )
     node_class = Node
     if schema.kind in registry.node:
         node_class = registry.node[schema.kind]
 
     fields_to_validate = list(data)
-    if db.is_transaction:
-        obj = await _do_create_node(
-            node_class=node_class,
-            node_constraint_runner=node_constraint_runner,
-            db=db,
-            schema=schema,
+    preview_obj = await node_class.init(db=db, schema=schema, branch=branch)
+    await preview_obj.new(db=db, process_pools=False, **data)
+    schema_branch = db.schema.get_schema_branch(name=branch.name)
+    lock_names = get_lock_names_on_object_mutation(node=preview_obj, branch=branch, schema_branch=schema_branch)
+
+    async def _persist(current_db: InfrahubDatabase) -> Node:
+        node_constraint_runner = await component_registry.get_component(
+            NodeConstraintRunner, db=current_db, branch=branch
+        )
+        obj = await node_class.init(db=current_db, schema=schema, branch=branch)
+        await obj.new(db=current_db, **data)
+
+        return await _do_create_node(
+            obj=obj,
+            db=current_db,
             branch=branch,
             fields_to_validate=fields_to_validate,
-            data=data,
+            node_constraint_runner=node_constraint_runner,
+            lock_names=lock_names,
         )
-    else:
-        async with db.start_transaction() as dbt:
-            obj = await _do_create_node(
-                node_class=node_class,
-                node_constraint_runner=node_constraint_runner,
-                db=dbt,
-                schema=schema,
-                branch=branch,
-                fields_to_validate=fields_to_validate,
-                data=data,
-            )
+
+    obj: Node
+    async with InfrahubMultiLock(lock_registry=lock.registry, locks=lock_names):
+        if db.is_transaction:
+            obj = await _persist(db)
+        else:
+            async with db.start_transaction() as dbt:
+                obj = await _persist(dbt)
 
     if await get_profile_ids(db=db, obj=obj):
         obj = await refresh_for_profile_update(db=db, branch=branch, schema=schema, obj=obj)

backend/infrahub/core/node/save.py

@@ -1,3 +1,5 @@
+from collections.abc import Sequence
+
 from infrahub import lock
 from infrahub.core.branch import Branch
 from infrahub.core.constraint.node.runner import NodeConstraintRunner
@@ -15,11 +17,41 @@ async def run_constraints_and_save(
     db: InfrahubDatabase,
     branch: Branch,
     skip_uniqueness_check: bool = False,
+    lock_names: Sequence[str] | None = None,
+    manage_lock: bool = True,
 ) -> None:
+    """Validate a node and persist it, optionally reusing an existing lock context.
+
+    Args:
+        node: The node instance to validate and persist.
+        node_constraint_runner: Runner executing node-level constraints.
+        fields_to_validate: Field names that must be validated.
+        fields_to_save: Field names that must be persisted.
+        db: Database connection or transaction to use for persistence.
+        branch: Branch associated with the mutation.
+        skip_uniqueness_check: Whether to skip uniqueness constraints.
+        lock_names: Precomputed lock identifiers to reuse when ``manage_lock`` is False.
+        manage_lock: Whether this helper should acquire and release locks itself.
+    """
+
+    if not manage_lock and lock_names is None:
+        raise ValueError("lock_names must be provided when manage_lock is False")
+
     schema_branch = db.schema.get_schema_branch(name=branch.name)
-    lock_names = get_lock_names_on_object_mutation(node=node, branch=branch, schema_branch=schema_branch)
-    async with InfrahubMultiLock(lock_registry=lock.registry, locks=lock_names):
+    locks = (
+        list(lock_names)
+        if lock_names is not None
+        else get_lock_names_on_object_mutation(node=node, branch=branch, schema_branch=schema_branch)
+    )
+
+    async def _persist() -> None:
         await node_constraint_runner.check(
             node=node, field_filters=fields_to_validate, skip_uniqueness_check=skip_uniqueness_check
         )
         await node.save(db=db, fields=fields_to_save)
+
+    if manage_lock:
+        async with InfrahubMultiLock(lock_registry=lock.registry, locks=locks):
+            await _persist()
+    else:
+        await _persist()