fix(backend): granular locking and transaction fixes

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

Author: fatih-acar

backend/infrahub/core/attribute.py

@@ -581,7 +581,7 @@ def _filter_sensitive(self, value: str, filter_sensitive: bool) -> str:
 
         return value
 
-    async def from_graphql(self, data: dict, db: InfrahubDatabase) -> bool:
+    async def from_graphql(self, data: dict, db: InfrahubDatabase, process_pools: bool = True) -> bool:
         """Update attr from GraphQL payload"""
 
         changed = False
@@ -595,7 +595,8 @@ async def from_graphql(self, data: dict, db: InfrahubDatabase) -> bool:
                 changed = True
         elif "from_pool" in data:
             self.from_pool = data["from_pool"]
-            await self.node.handle_pool(db=db, attribute=self, errors=[])
+            if process_pools:
+                await self.node.handle_pool(db=db, attribute=self, errors=[])
             changed = True
 
         if changed and self.is_from_profile:

backend/infrahub/core/constants/__init__.py

@@ -37,6 +37,7 @@
     "rels",
     "save",
     "hfid",
+    "process_pools",
 ]
 
 RESERVED_ATTR_GEN_NAMES = ["type"]

backend/infrahub/core/node/__init__.py

@@ -314,7 +314,9 @@ 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
@@ -325,7 +327,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:
@@ -485,7 +487,7 @@ 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():
@@ -539,7 +541,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)
@@ -576,7 +578,9 @@ 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:
@@ -601,9 +605,10 @@ 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.from_pool is None:
+                        attribute.validate(value=attribute.value, name=attribute.name, schema=attribute.schema)
             except ValidationError as exc:
                 errors.append(exc)
 
@@ -731,7 +736,7 @@ 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:
@@ -741,7 +746,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
@@ -1046,15 +1051,15 @@ async def to_graphql(
 
         return response
 
-    async def from_graphql(self, data: dict, db: InfrahubDatabase) -> bool:
+    async def from_graphql(self, data: dict, db: InfrahubDatabase, process_pools: bool = True) -> bool:
         """Update object from a GraphQL payload."""
 
         changed = False
 
         for key, value in data.items():
             if key in self._attributes and isinstance(value, dict):
                 attribute = getattr(self, key)
-                changed |= await attribute.from_graphql(data=value, db=db)
+                changed |= await attribute.from_graphql(data=value, db=db, process_pools=process_pools)
 
             if key in self._relationships:
                 rel: RelationshipManager = getattr(self, key)

backend/infrahub/core/node/create.py

@@ -7,7 +7,7 @@
 from infrahub.core.constants import RelationshipCardinality, RelationshipKind
 from infrahub.core.constraint.node.runner import NodeConstraintRunner
 from infrahub.core.node import Node
-from infrahub.core.node.lock_utils import get_kind_lock_names_on_object_mutation
+from infrahub.core.node.lock_utils import get_lock_names_on_object_mutation
 from infrahub.core.protocols import CoreObjectTemplate
 from infrahub.core.schema import GenericSchema
 from infrahub.dependencies.registry import get_component_registry
@@ -171,45 +171,6 @@ async def _do_create_node(
     return obj
 
 
-async def _do_create_node_with_lock(
-    node_class: type[Node],
-    node_constraint_runner: NodeConstraintRunner,
-    db: InfrahubDatabase,
-    schema: NonGenericSchemaTypes,
-    branch: Branch,
-    fields_to_validate: list[str],
-    data: dict[str, Any],
-    at: Timestamp | None = None,
-) -> Node:
-    schema_branch = registry.schema.get_schema_branch(name=branch.name)
-    lock_names = get_kind_lock_names_on_object_mutation(
-        kind=schema.kind, branch=branch, schema_branch=schema_branch, data=dict(data)
-    )
-
-    if lock_names:
-        async with InfrahubMultiLock(lock_registry=lock.registry, locks=lock_names):
-            return await _do_create_node(
-                node_class=node_class,
-                node_constraint_runner=node_constraint_runner,
-                db=db,
-                schema=schema,
-                branch=branch,
-                fields_to_validate=fields_to_validate,
-                data=data,
-                at=at,
-            )
-    return await _do_create_node(
-        node_class=node_class,
-        node_constraint_runner=node_constraint_runner,
-        db=db,
-        schema=schema,
-        branch=branch,
-        fields_to_validate=fields_to_validate,
-        data=data,
-        at=at,
-    )
-
-
 async def create_node(
     data: dict[str, Any],
     db: InfrahubDatabase,
@@ -223,37 +184,48 @@ async def create_node(
         raise ValueError(f"Node of generic schema `{schema.name=}` can not be instantiated.")
 
     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_with_lock(
-            node_class=node_class,
-            node_constraint_runner=node_constraint_runner,
-            db=db,
-            schema=schema,
-            branch=branch,
-            fields_to_validate=fields_to_validate,
-            data=data,
-            at=at,
-        )
-    else:
-        async with db.start_transaction() as dbt:
-            obj = await _do_create_node_with_lock(
+
+    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, schema_branch=schema_branch)
+
+    obj: Node
+    async with InfrahubMultiLock(lock_registry=lock.registry, locks=lock_names, metrics=False):
+        if db.is_transaction:
+            node_constraint_runner = await component_registry.get_component(NodeConstraintRunner, db=db, branch=branch)
+
+            obj = await _do_create_node(
                 node_class=node_class,
                 node_constraint_runner=node_constraint_runner,
-                db=dbt,
+                db=db,
                 schema=schema,
                 branch=branch,
                 fields_to_validate=fields_to_validate,
                 data=data,
                 at=at,
             )
+        else:
+            async with db.start_transaction() as dbt:
+                node_constraint_runner = await component_registry.get_component(
+                    NodeConstraintRunner, db=dbt, branch=branch
+                )
+
+                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,
+                    at=at,
+                )
 
     if await get_profile_ids(db=db, obj=obj):
         node_profiles_applier = NodeProfilesApplier(db=db, branch=branch)

backend/infrahub/core/node/lock_utils.py

@@ -1,12 +1,15 @@
 import hashlib
-from typing import Any
+from typing import TYPE_CHECKING
 
-from infrahub.core.branch import Branch
-from infrahub.core.constants.infrahubkind import GENERICGROUP, GRAPHQLQUERYGROUP
+from infrahub.core.node import Node
 from infrahub.core.schema import GenericSchema
 from infrahub.core.schema.schema_branch import SchemaBranch
 
-KINDS_CONCURRENT_MUTATIONS_NOT_ALLOWED = [GENERICGROUP]
+if TYPE_CHECKING:
+    from infrahub.core.relationship import RelationshipManager
+
+
+RESOURCE_POOL_LOCK_NAMESPACE = "resource_pool"
 
 
 def _get_kinds_to_lock_on_object_mutation(kind: str, schema_branch: SchemaBranch) -> list[str]:
@@ -43,55 +46,78 @@ def _get_kinds_to_lock_on_object_mutation(kind: str, schema_branch: SchemaBranch
     return kinds
 
 
-def _should_kind_be_locked_on_any_branch(kind: str, schema_branch: SchemaBranch) -> bool:
-    """
-    Check whether kind or any kind generic is in KINDS_TO_LOCK_ON_ANY_BRANCH.
-    """
-
-    if kind in KINDS_CONCURRENT_MUTATIONS_NOT_ALLOWED:
-        return True
-
-    node_schema = schema_branch.get(name=kind, duplicate=False)
-    if isinstance(node_schema, GenericSchema):
-        return False
-
-    for generic_kind in node_schema.inherit_from:
-        if generic_kind in KINDS_CONCURRENT_MUTATIONS_NOT_ALLOWED:
-            return True
-    return False
-
-
 def _hash(value: str) -> str:
     # Do not use builtin `hash` for lock names as due to randomization results would differ between
     # different processes.
     return hashlib.sha256(value.encode()).hexdigest()
 
 
-def get_kind_lock_names_on_object_mutation(
-    kind: str, branch: Branch, schema_branch: SchemaBranch, data: dict[str, Any]
-) -> list[str]:
+def get_lock_names_on_object_mutation(node: Node, schema_branch: SchemaBranch) -> list[str]:
     """
-    Return objects kind for which we want to avoid concurrent mutation (create/update). Except for some specific kinds,
-    concurrent mutations are only allowed on non-main branch as objects validations will be performed at least when merging in main branch.
+    Return lock names for object on which we want to avoid concurrent mutation (create/update).
+    Lock names include kind, some generic kinds, resource pool ids, and values of attributes of corresponding uniqueness constraints.
     """
 
-    if not branch.is_default and not _should_kind_be_locked_on_any_branch(kind=kind, schema_branch=schema_branch):
-        return []
-
-    if kind == GRAPHQLQUERYGROUP:
-        # Lock on name as well to improve performances
-        try:
-            name = data["name"].value
-            return [build_object_lock_name(kind + "." + _hash(name))]
-        except KeyError:
-            # We might reach here if we are updating a CoreGraphQLQueryGroup without updating the name,
-            # in which case we would not need to lock. This is not supposed to happen as current `update`
-            # logic first fetches the node with its name.
-            return []
-
-    lock_kinds = _get_kinds_to_lock_on_object_mutation(kind, schema_branch)
-    lock_names = [build_object_lock_name(kind) for kind in lock_kinds]
-    return lock_names
+    lock_names: set[str] = set()
+
+    # Check if node is using resource manager allocation via attributes
+    for attr_name in node.get_schema().attribute_names:
+        attribute = getattr(node, attr_name, None)
+        if attribute is not None and getattr(attribute, "from_pool", None) and "id" in attribute.from_pool:
+            lock_names.add(f"{RESOURCE_POOL_LOCK_NAMESPACE}.{attribute.from_pool['id']}")
+
+    # Check if relationships allocate resources
+    for rel_name in node._relationships:
+        rel_manager: RelationshipManager = getattr(node, rel_name)
+        for rel in rel_manager._relationships:
+            if rel.from_pool and "id" in rel.from_pool:
+                lock_names.add(f"{RESOURCE_POOL_LOCK_NAMESPACE}.{rel.from_pool['id']}")
+
+    lock_kinds = _get_kinds_to_lock_on_object_mutation(node.get_kind(), schema_branch)
+    for kind in lock_kinds:
+        schema = schema_branch.get(name=kind, duplicate=False)
+        ucs = schema.uniqueness_constraints
+        if ucs is None:
+            continue
+
+        ucs_lock_names: list[str] = []
+        uc_attributes_names = set()
+
+        for uc in ucs:
+            uc_attributes_values = []
+            # Keep only attributes constraints
+            for field_path in uc:
+                # Some attributes may exist in different uniqueness constraints, we de-duplicate them
+                if field_path in uc_attributes_names:
+                    continue
+
+                # Exclude relationships uniqueness constraints
+                schema_path = schema.parse_schema_path(path=field_path, schema=schema_branch)
+                if schema_path.related_schema is not None or schema_path.attribute_schema is None:
+                    continue
+
+                uc_attributes_names.add(field_path)
+                attr = getattr(node, schema_path.attribute_schema.name, None)
+                if attr is None or attr.value is None:
+                    # `attr.value` being None corresponds to optional unique attribute.
+                    # `attr` being None is not supposed to happen.
+                    value_hashed = _hash("")
+                else:
+                    value_hashed = _hash(str(attr.value))
+
+                uc_attributes_values.append(value_hashed)
+
+            if uc_attributes_values:
+                uc_lock_name = ".".join(uc_attributes_values)
+                ucs_lock_names.append(uc_lock_name)
+
+        if not ucs_lock_names:
+            continue
+
+        partial_lock_name = kind + "." + ".".join(ucs_lock_names)
+        lock_names.add(build_object_lock_name(partial_lock_name))
+
+    return sorted(lock_names)
 
 
 def build_object_lock_name(name: str) -> str: