Implements address space cleanup when a workspace service is uninstalled

CHANGELOG.md

@@ -4,7 +4,7 @@
 * _No changes yet_
 
 ENHANCEMENTS:
-* _No changes yet_
+* `address_spaces` will now be removed from a workspace when a workspace service that uses an `address_space` is deleted to prevent IP address range exhaustion ([#4727](https://github.com/microsoft/AzureTRE/issues/4727))
 
 BUG FIXES:
 * Fix property substitution not occuring where there is only a main step in the pipeline ([#4824](https://github.com/microsoft/AzureTRE/issues/4824))

api_app/_version.py

@@ -1 +1 @@
-__version__ = "0.25.14"
+__version__ = "0.25.15"

api_app/service_bus/deployment_status_updater.py

@@ -6,7 +6,7 @@
 from pydantic import ValidationError, parse_obj_as
 
 from api.routes.resource_helpers import get_timestamp
-from models.domain.resource import Output
+from models.domain.resource import Output, ResourceType
 from db.repositories.resources_history import ResourceHistoryRepository
 from models.domain.request_action import RequestAction
 from db.repositories.resource_templates import ResourceTemplateRepository
@@ -21,6 +21,9 @@
 from models.domain.operation import DeploymentStatusUpdateMessage, Operation, OperationStep, Status
 from resources import strings
 from services.logging import logger, tracer
+from db.repositories.workspaces import WorkspaceRepository
+from models.schemas.resource import ResourcePatch
+from azure.cosmos.exceptions import CosmosAccessConditionFailedError
 
 
 class DeploymentStatusUpdater():
@@ -187,6 +190,34 @@ async def update_status_in_database(self, message: DeploymentStatusUpdateMessage
                     next_step.status = Status.UpdatingFailed
                     await self.update_overall_operation_status(operation, next_step, is_last_step)
                     await self.operations_repo.update_item(operation)
+            # If the 'main' step succeeded for an uninstall operation, free any allocated address space
+            # owned by a WorkspaceService resource. We trigger cleanup when the step with templateStepId == 'main'
+            # is successful; this ensures the primary resource has been destroyed successfully before attempting to free the ip address space
+            try:
+                # if the step that just succeeded is the main step for this operation, and this is an uninstall,
+                # proceed with post-uninstall cleanup. No need to scan the operation.steps list again.
+                if step_to_update.templateStepId == "main" and step_to_update.is_success() and operation.action == RequestAction.UnInstall:
+                    if resource_to_persist.get("resourceType") == ResourceType.WorkspaceService:
+                        address_to_free = resource_to_persist.get("properties", {}).get("address_space")
+                        parent_workspace_id = resource_to_persist.get("workspaceId")
+                        if address_to_free and parent_workspace_id:
+                            try:
+                                workspace_repo = await WorkspaceRepository.create()
+                                workspace = await workspace_repo.get_workspace_by_id(parent_workspace_id)
+                                workspace_address_spaces = workspace.properties.get("address_spaces", [])
+                                if address_to_free in workspace_address_spaces:
+                                    new_address_spaces = [a for a in workspace_address_spaces if a != address_to_free]
+                                    workspace_patch = ResourcePatch()
+                                    workspace_patch.properties = {"address_spaces": new_address_spaces}
+                                    try:
+                                        await workspace_repo.patch_workspace(workspace, workspace_patch, workspace.etag, self.resource_template_repo, self.resource_history_repo, operation.user, False)
+                                        logger.info(f"Freed address space {address_to_free} from workspace {parent_workspace_id} after successful uninstall of {resource_id}")
+                                    except CosmosAccessConditionFailedError:
+                                        logger.exception("ETag conflict when freeing workspace address space after successful uninstall")
+                            except Exception:
+                                logger.exception("Failed to free workspace address space after successful uninstall")
+            except Exception:
+                logger.exception("Unexpected error during post-uninstall address space cleanup")
 
             result = True
 

templates/workspace_services/azureml/porter.yaml

@@ -1,7 +1,7 @@
 ---
 schemaVersion: 1.0.0
 name: tre-service-azureml
-version: 1.1.2
+version: 1.1.3
 description: "An Azure TRE service for Azure Machine Learning"
 registry: azuretre
 dockerfile: Dockerfile.tmpl

templates/workspace_services/azureml/template_schema.json

@@ -417,6 +417,13 @@
       },
       {
         "stepId": "main"
+      },
+      {
+        "stepId": "f720975a-c81e-477e-854e-53fde86e5e57",
+        "stepTitle": "Upgrade to ensure workspace is aware of address space removal",
+        "resourceType": "workspace",
+        "resourceAction": "upgrade",
+        "properties": []
       }
     ]
   }

templates/workspace_services/databricks/porter.yaml

@@ -1,7 +1,7 @@
 ---
 schemaVersion: 1.0.0
 name: tre-service-databricks
-version: 1.0.14
+version: 1.0.15
 description: "An Azure TRE service for Azure Databricks."
 registry: azuretre
 dockerfile: Dockerfile.tmpl

templates/workspace_services/databricks/template_schema.json

@@ -78,7 +78,9 @@
                   "name": "databricks",
                   "description": "Communication with Azure Databricks dependancies.",
                   "source_addresses": "{{ resource.properties.databricks_address_prefixes }}",
-                  "destination_addresses": [ "AzureDatabricks"],
+                  "destination_addresses": [
+                    "AzureDatabricks"
+                  ],
                   "destination_ports": [
                     "443"
                   ],
@@ -114,9 +116,15 @@
                   "name": "AzureAD",
                   "description": "AAD access",
                   "source_addresses": "{{ resource.properties.workspace_address_spaces }}",
-                  "destination_addresses": ["AzureActiveDirectory"],
-                  "destination_ports": ["*"],
-                  "protocols": ["TCP"]
+                  "destination_addresses": [
+                    "AzureActiveDirectory"
+                  ],
+                  "destination_ports": [
+                    "*"
+                  ],
+                  "protocols": [
+                    "TCP"
+                  ]
                 }
               ]
             }
@@ -212,7 +220,9 @@
                   "name": "databricks",
                   "description": "Communication with Azure Databricks dependancies.",
                   "source_addresses": "{{ resource.properties.databricks_address_prefixes }}",
-                  "destination_addresses": [ "AzureDatabricks"],
+                  "destination_addresses": [
+                    "AzureDatabricks"
+                  ],
                   "destination_ports": [
                     "443"
                   ],
@@ -248,9 +258,15 @@
                   "name": "AzureAD",
                   "description": "AAD access",
                   "source_addresses": "{{ resource.properties.workspace_address_spaces }}",
-                  "destination_addresses": ["AzureActiveDirectory"],
-                  "destination_ports": ["*"],
-                  "protocols": ["TCP"]
+                  "destination_addresses": [
+                    "AzureActiveDirectory"
+                  ],
+                  "destination_ports": [
+                    "*"
+                  ],
+                  "protocols": [
+                    "TCP"
+                  ]
                 }
               ]
             }
@@ -352,6 +368,13 @@
       },
       {
         "stepId": "main"
+      },
+      {
+        "stepId": "9c4dc64b-8fbf-4e77-a7f6-48fb33423504",
+        "stepTitle": "Upgrade to ensure workspace is aware of address space removal",
+        "resourceType": "workspace",
+        "resourceAction": "upgrade",
+        "properties": []
       }
     ]
   }