Merge pull request #6572 from opsmill/pog-numberpool-deletion

Delete unused schema defined NumberPools when they are no longer in use

Author: ogenstad

backend/infrahub/branch/__init__.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from prefect import flow
+from prefect.logging import get_run_logger
+
+from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
+from infrahub.core.registry import registry
+from infrahub.pools.tasks import validate_schema_number_pools
+from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
+from infrahub.workflows.utils import wait_for_schema_to_converge
+
+
+@flow(
+    name="branch-merged",
+    flow_run_name="Running actions after '{source_branch}' was merged",
+)
+async def branch_merged(
+    source_branch: str,  # noqa: ARG001
+    context: InfrahubContext,
+    service: InfrahubServices,
+    target_branch: str | None = None,
+) -> None:
+    target_branch = target_branch or registry.default_branch
+    log = get_run_logger()
+    await wait_for_schema_to_converge(
+        branch_name=target_branch, component=service.component, db=service.database, log=log
+    )
+
+    await validate_schema_number_pools(branch_name=target_branch, context=context, service=service)

backend/infrahub/branch/tasks.py

@@ -0,0 +1,22 @@
+from infrahub.events.branch_action import BranchMergedEvent
+from infrahub.trigger.models import BuiltinTriggerDefinition, EventTrigger, ExecuteWorkflow
+from infrahub.workflows.catalogue import BRANCH_MERGED
+
+TRIGGER_BRANCH_MERGED = BuiltinTriggerDefinition(
+    name="branch-merged-trigger",
+    trigger=EventTrigger(
+        events={BranchMergedEvent.event_name},
+    ),
+    actions=[
+        ExecuteWorkflow(
+            workflow=BRANCH_MERGED,
+            parameters={
+                "source_branch": "{{ event.payload['data']['branch_name'] }}",
+                "context": {
+                    "__prefect_kind": "json",
+                    "value": {"__prefect_kind": "jinja", "template": "{{ event.payload['context'] | tojson }}"},
+                },
+            },
+        ),
+    ],
+)

backend/infrahub/branch/triggers.py

@@ -112,6 +112,11 @@ class AccountType(InfrahubStringEnum):
     Git = "Git"
 
 
+class NumberPoolType(InfrahubStringEnum):
+    USER = "User"
+    SCHEMA = "Schema"
+
+
 class AccountStatus(InfrahubStringEnum):
     ACTIVE = "active"
     INACTIVE = "inactive"

backend/infrahub/core/constants/__init__.py

@@ -16,6 +16,7 @@
     BranchSupportType,
     ComputedAttributeKind,
     InfrahubKind,
+    NumberPoolType,
     RelationshipCardinality,
     RelationshipKind,
 )
@@ -328,6 +329,7 @@ async def _create_number_pool(
             node_attribute=attribute.schema.name,
             start_range=number_pool_parameters.start_range,
             end_range=number_pool_parameters.end_range,
+            pool_type=NumberPoolType.SCHEMA.value,
         )
         await number_pool.save(db=db)
         # Do a lookup of the number pool to get the correct mapped type from the registry

backend/infrahub/core/node/__init__.py

@@ -455,6 +455,7 @@ class CoreNumberPool(CoreResourcePool, LineageSource):
     node_attribute: String
     start_range: Integer
     end_range: Integer
+    pool_type: Enum
 
 
 class CoreObjectPermission(CoreBasePermission):

backend/infrahub/core/protocols.py

@@ -1,6 +1,7 @@
 from infrahub.core.constants import (
     BranchSupportType,
     InfrahubKind,
+    NumberPoolType,
 )
 from infrahub.core.constants import RelationshipCardinality as Cardinality
 from infrahub.core.constants import RelationshipKind as RelKind
@@ -186,5 +187,13 @@
         Attr(
             name="end_range", kind="Number", optional=False, description="The end range for the pool", order_weight=6000
         ),
+        Attr(
+            name="pool_type",
+            kind="Text",
+            description="Defines how this number pool was created",
+            default_value=NumberPoolType.USER.value,
+            enum=NumberPoolType.available_types(),
+            read_only=True,
+        ),
     ],
 )

backend/infrahub/core/schema/definitions/core/resource_pool.py

@@ -14,6 +14,7 @@
 from infrahub.core.schema.attribute_parameters import NumberAttributeParameters
 from infrahub.database import retry_db_transaction
 from infrahub.exceptions import QueryValidationError, SchemaNotFoundError, ValidationError
+from infrahub.pools.registration import get_branches_with_schema_number_pool
 
 from ..queries.resource_manager import PoolAllocatedNode
 from .main import DeleteResult, InfrahubMutationMixin, InfrahubMutationOptions
@@ -257,18 +258,9 @@ async def mutate_delete(
             branch=branch,
         )
 
-        active_branches = registry.schema.get_branches()
-        violating_branches = []
-        for active_branch in active_branches:
-            try:
-                schema = registry.schema.get(name=number_pool.node.value, branch=active_branch)
-            except SchemaNotFoundError:
-                continue
-
-            if number_pool.node_attribute.value in schema.attribute_names:
-                attribute = schema.get_attribute(name=number_pool.node_attribute.value)
-                if attribute.kind == "NumberPool":
-                    violating_branches.append(active_branch)
+        violating_branches = get_branches_with_schema_number_pool(
+            kind=number_pool.node.value, attribute_name=number_pool.node_attribute.value
+        )
 
         if violating_branches:
             raise ValidationError(

backend/infrahub/graphql/mutations/resource_manager.py

@@ -0,0 +1,22 @@
+from infrahub.core.registry import registry
+from infrahub.exceptions import SchemaNotFoundError
+
+
+def get_branches_with_schema_number_pool(kind: str, attribute_name: str) -> list[str]:
+    """Return branches where schema defined NumberPool exists"""
+
+    registered_branches = []
+    active_branches = registry.schema.get_branches()
+
+    for active_branch in active_branches:
+        try:
+            schema = registry.schema.get(name=kind, branch=active_branch)
+        except SchemaNotFoundError:
+            continue
+
+        if attribute_name in schema.attribute_names:
+            attribute = schema.get_attribute(name=attribute_name)
+            if attribute.kind == "NumberPool":
+                registered_branches.append(active_branch)
+
+    return registered_branches

backend/infrahub/pools/registration.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from prefect import flow
+from prefect.logging import get_run_logger
+
+from infrahub.context import InfrahubContext  # noqa: TC001  needed for prefect flow
+from infrahub.core.constants import NumberPoolType
+from infrahub.core.manager import NodeManager
+from infrahub.core.protocols import CoreNumberPool
+from infrahub.core.registry import registry
+from infrahub.core.schema.attribute_parameters import NumberPoolParameters
+from infrahub.pools.registration import get_branches_with_schema_number_pool
+from infrahub.services import InfrahubServices  # noqa: TC001  needed for prefect flow
+
+
+@flow(
+    name="validate-schema-number-pools",
+    flow_run_name="Validate schema number pools on {branch_name}",
+)
+async def validate_schema_number_pools(
+    branch_name: str,  # noqa: ARG001
+    context: InfrahubContext,  # noqa: ARG001
+    service: InfrahubServices,
+) -> None:
+    log = get_run_logger()
+
+    async with service.database.start_session() as dbs:
+        schema_number_pools = await NodeManager.query(
+            db=dbs, schema=CoreNumberPool, filters={"pool_type__value": NumberPoolType.SCHEMA.value}
+        )
+
+    for schema_number_pool in list(schema_number_pools):
+        defined_on_branches = get_branches_with_schema_number_pool(
+            kind=schema_number_pool.node.value, attribute_name=schema_number_pool.node_attribute.value
+        )
+        if registry.default_branch in defined_on_branches:
+            schema = registry.schema.get(name=schema_number_pool.node.value, branch=registry.default_branch)
+            attribute = schema.get_attribute(name=schema_number_pool.node_attribute.value)
+            number_pool_updated = False
+            if isinstance(attribute.parameters, NumberPoolParameters):
+                if schema_number_pool.start_range.value != attribute.parameters.start_range:
+                    schema_number_pool.start_range.value = attribute.parameters.start_range
+                    number_pool_updated = True
+                if schema_number_pool.end_range.value != attribute.parameters.end_range:
+                    schema_number_pool.end_range.value = attribute.parameters.end_range
+                    number_pool_updated = True
+
+            if number_pool_updated:
+                log.info(
+                    f"Updating NumberPool={schema_number_pool.id} based on changes in the schema on {registry.default_branch}"
+                )
+                await schema_number_pool.save(db=service.database)
+
+        elif not defined_on_branches:
+            log.info(f"Deleting number pool (id={schema_number_pool.id}) as it is no longer defined in the schema")
+            await schema_number_pool.delete(db=service.database)