Consolidate airlock storage architecture and implement metadata-based management

.gitignore

@@ -214,3 +214,4 @@ validation.txt
 
 /index.html
 .DS_Store
+*_old.tf

CHANGELOG.md

@@ -2,6 +2,7 @@
 ## (Unreleased)
 
 ENHANCEMENTS:
+* Add per-workspace `airlock_version` property (1=legacy, 2=consolidated) for backwards-compatible airlock storage migration. Add core-level `enable_legacy_airlock` toggle. Remove `USE_METADATA_STAGE_MANAGEMENT` environment variable. ([#4853](https://github.com/microsoft/AzureTRE/pull/4853), [#4358](https://github.com/microsoft/AzureTRE/issues/4358))
 * Specify default_outbound_access_enabled = false setting for all subnets ([#4757](https://github.com/microsoft/AzureTRE/pull/4757))
 * Pin all GitHub Actions workflow steps to full commit SHAs to prevent supply chain attacks plus update to latest releases ([#4886](https://github.com/microsoft/AzureTRE/pull/4886))
 

airlock_processor/BlobCreatedTrigger/__init__.py

@@ -11,6 +11,17 @@
 from shared_code.blob_operations import get_blob_info_from_topic_and_subject, get_blob_client_from_blob_info
 
 
+# Mapping from v2 container metadata stage to (completed_step, new_status)
+V2_STAGE_COMPLETION_MAP = {
+    constants.STAGE_IMPORT_APPROVED: (constants.STAGE_APPROVAL_INPROGRESS, constants.STAGE_APPROVED),
+    constants.STAGE_IMPORT_REJECTED: (constants.STAGE_REJECTION_INPROGRESS, constants.STAGE_REJECTED),
+    constants.STAGE_IMPORT_BLOCKED: (constants.STAGE_BLOCKING_INPROGRESS, constants.STAGE_BLOCKED_BY_SCAN),
+    constants.STAGE_EXPORT_APPROVED: (constants.STAGE_APPROVAL_INPROGRESS, constants.STAGE_APPROVED),
+    constants.STAGE_EXPORT_REJECTED: (constants.STAGE_REJECTION_INPROGRESS, constants.STAGE_REJECTED),
+    constants.STAGE_EXPORT_BLOCKED: (constants.STAGE_BLOCKING_INPROGRESS, constants.STAGE_BLOCKED_BY_SCAN),
+}
+
+
 def main(msg: func.ServiceBusMessage,
          stepResultEvent: func.Out[func.EventGridOutputEvent],
          dataDeletionEvent: func.Out[func.EventGridOutputEvent]):
@@ -23,6 +34,12 @@ def main(msg: func.ServiceBusMessage,
     topic = json_body["topic"]
     request_id = re.search(r'/blobServices/default/containers/(.*?)/blobs', json_body["subject"]).group(1)
 
+    # Check if this event is from a v2 consolidated storage account
+    if constants.STORAGE_ACCOUNT_NAME_AIRLOCK_CORE in topic or constants.STORAGE_ACCOUNT_NAME_AIRLOCK_WORKSPACE_GLOBAL in topic:
+        _handle_v2_blob_created(json_body, topic, request_id, stepResultEvent, dataDeletionEvent)
+        return
+
+    # Legacy v1 handling below
     # message originated from in-progress blob creation
     if constants.STORAGE_ACCOUNT_NAME_IMPORT_INPROGRESS in topic or constants.STORAGE_ACCOUNT_NAME_EXPORT_INPROGRESS in topic:
         try:
@@ -55,6 +72,9 @@ def main(msg: func.ServiceBusMessage,
     elif constants.STORAGE_ACCOUNT_NAME_IMPORT_BLOCKED in topic or constants.STORAGE_ACCOUNT_NAME_EXPORT_BLOCKED in topic:
         completed_step = constants.STAGE_BLOCKING_INPROGRESS
         new_status = constants.STAGE_BLOCKED_BY_SCAN
+    else:
+        logging.warning(f"Unknown storage account in topic: {topic}")
+        return
 
     # reply with a step completed event
     stepResultEvent.set(
@@ -88,3 +108,44 @@ def send_delete_event(dataDeletionEvent: func.Out[func.EventGridOutputEvent], js
             data_version=constants.DATA_DELETION_EVENT_DATA_VERSION
         )
     )
+
+
+def _handle_v2_blob_created(json_body, topic, request_id, stepResultEvent, dataDeletionEvent):
+    """Handle BlobCreated events from v2 consolidated storage accounts.
+
+    In v2, cross-account copies (e.g., import approval: core → workspace-global)
+    fire BlobCreated events. Container metadata determines the stage and appropriate
+    step result, matching the v1 pattern where BlobCreatedTrigger signals copy completion.
+    """
+    storage_account_name, _, _ = get_blob_info_from_topic_and_subject(
+        topic=json_body["topic"], subject=json_body["subject"])
+
+    from shared_code.blob_operations_metadata import get_container_metadata
+    try:
+        metadata = get_container_metadata(storage_account_name, request_id)
+    except Exception:
+        logging.warning(f"Could not read container metadata for request {request_id} on {storage_account_name}, skipping")
+        return
+
+    stage = metadata.get('stage', '')
+    logging.info(f"V2 BlobCreated for request {request_id}: stage={stage}, account={storage_account_name}")
+
+    if stage in V2_STAGE_COMPLETION_MAP:
+        completed_step, new_status = V2_STAGE_COMPLETION_MAP[stage]
+        logging.info(f"V2 copy completed for request {request_id}: {completed_step} -> {new_status}")
+
+        stepResultEvent.set(
+            func.EventGridOutputEvent(
+                id=str(uuid.uuid4()),
+                data={"completed_step": completed_step, "new_status": new_status, "request_id": request_id},
+                subject=request_id,
+                event_type="Airlock.StepResult",
+                event_time=datetime.datetime.now(datetime.UTC),
+                data_version=constants.STEP_RESULT_EVENT_DATA_VERSION))
+
+        # Send delete event for the source container (same as v1)
+        send_delete_event(dataDeletionEvent, json_body, request_id)
+    else:
+        # Non-terminal stages (e.g., import-external from user upload, export-internal)
+        # are not copy completions — ignore them
+        logging.info(f"V2 BlobCreated for non-terminal stage '{stage}' on request {request_id}, no action needed")

airlock_processor/StatusChangedQueueTrigger/__init__.py

@@ -9,7 +9,7 @@
 
 from exceptions import NoFilesInRequestException, TooManyFilesInRequestException
 
-from shared_code import blob_operations, constants
+from shared_code import blob_operations, constants, airlock_storage_helper, parsers
 from pydantic import BaseModel, parse_obj_as
 
 
@@ -19,6 +19,8 @@ class RequestProperties(BaseModel):
     previous_status: Optional[str]
     type: str
     workspace_id: str
+    review_workspace_id: Optional[str] = None
+    airlock_version: int = 1
 
 
 class ContainersCopyMetadata:
@@ -31,6 +33,8 @@ def __init__(self, source_account_name: str, dest_account_name: str):
 
 
 def main(msg: func.ServiceBusMessage, stepResultEvent: func.Out[func.EventGridOutputEvent], dataDeletionEvent: func.Out[func.EventGridOutputEvent]):
+    request_properties = None
+    request_files = None
     try:
         request_properties = extract_properties(msg)
         request_files = get_request_files(request_properties) if request_properties.new_status == constants.STAGE_SUBMITTED else None
@@ -53,13 +57,25 @@ def handle_status_changed(request_properties: RequestProperties, stepResultEvent
 
     logging.info('Processing request with id %s. new status is "%s", type is "%s"', req_id, new_status, request_type)
 
+    # Check if using metadata-based stage management (v2) or legacy per-stage accounts (v1)
+    use_metadata = request_properties.airlock_version >= 2
+
     if new_status == constants.STAGE_DRAFT:
-        account_name = get_storage_account(status=constants.STAGE_DRAFT, request_type=request_type, short_workspace_id=ws_id)
-        blob_operations.create_container(account_name, req_id)
+        if use_metadata:
+            from shared_code.blob_operations_metadata import create_container_with_metadata
+            account_name = airlock_storage_helper.get_storage_account_name_for_request(request_type, new_status, ws_id, airlock_version=request_properties.airlock_version)
+            stage = airlock_storage_helper.get_stage_from_status(request_type, new_status)
+            create_container_with_metadata(account_name, req_id, stage, workspace_id=ws_id, request_type=request_type)
+        else:
+            account_name = get_storage_account(status=constants.STAGE_DRAFT, request_type=request_type, short_workspace_id=ws_id)
+            blob_operations.create_container(account_name, req_id)
         return
 
     if new_status == constants.STAGE_CANCELLED:
-        storage_account_name = get_storage_account(previous_status, request_type, ws_id)
+        if use_metadata:
+            storage_account_name = airlock_storage_helper.get_storage_account_name_for_request(request_type, previous_status, ws_id, airlock_version=request_properties.airlock_version)
+        else:
+            storage_account_name = get_storage_account(previous_status, request_type, ws_id)
         container_to_delete_url = blob_operations.get_blob_url(account_name=storage_account_name, container_name=req_id)
         set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_properties, container_url=container_to_delete_url)
         return
@@ -68,11 +84,74 @@ def handle_status_changed(request_properties: RequestProperties, stepResultEvent
         set_output_event_to_report_request_files(stepResultEvent, request_properties, request_files)
 
     if (is_require_data_copy(new_status)):
-        logging.info('Request with id %s. requires data copy between storage accounts', req_id)
-        containers_metadata = get_source_dest_for_copy(new_status=new_status, previous_status=previous_status, request_type=request_type, short_workspace_id=ws_id)
-        blob_operations.create_container(containers_metadata.dest_account_name, req_id)
-        blob_operations.copy_data(containers_metadata.source_account_name,
-                                  containers_metadata.dest_account_name, req_id)
+        if use_metadata:
+            # Metadata mode: Update container stage instead of copying
+            from shared_code.blob_operations_metadata import update_container_stage, create_container_with_metadata
+
+            # For import submit, use review_workspace_id so data goes to review workspace storage
+            effective_ws_id = ws_id
+            if new_status == constants.STAGE_SUBMITTED and request_type.lower() == constants.IMPORT_TYPE and request_properties.review_workspace_id:
+                effective_ws_id = request_properties.review_workspace_id
+
+            # Get the storage account (might change from core to workspace or vice versa)
+            source_account = airlock_storage_helper.get_storage_account_name_for_request(request_type, previous_status, ws_id, airlock_version=request_properties.airlock_version)
+            dest_account = airlock_storage_helper.get_storage_account_name_for_request(request_type, new_status, effective_ws_id, airlock_version=request_properties.airlock_version)
+            new_stage = airlock_storage_helper.get_stage_from_status(request_type, new_status)
+
+            if source_account == dest_account:
+                # Same storage account - just update metadata
+                logging.info(f'Request {req_id}: Updating container stage to {new_stage} (no copy needed)')
+                update_container_stage(source_account, req_id, new_stage, changed_by='system')
+
+                # In v2, same-account transitions don't fire BlobCreated events.
+                # For SUBMITTED, v1 relies on BlobCreatedTrigger to handle the malware scanning gate
+                # (skip to in_review when scanning is disabled). We handle this inline for v2.
+                if new_status == constants.STAGE_SUBMITTED:
+                    try:
+                        enable_malware_scanning = parsers.parse_bool(os.environ["ENABLE_MALWARE_SCANNING"])
+                    except KeyError:
+                        logging.error("environment variable 'ENABLE_MALWARE_SCANNING' does not exist. Cannot continue.")
+                        raise
+                    if not enable_malware_scanning:
+                        logging.info(f'Request {req_id}: Malware scanning disabled, skipping to in_review')
+                        stepResultEvent.set(
+                            func.EventGridOutputEvent(
+                                id=str(uuid.uuid4()),
+                                data={"completed_step": constants.STAGE_SUBMITTED, "new_status": constants.STAGE_IN_REVIEW, "request_id": req_id},
+                                subject=req_id,
+                                event_type="Airlock.StepResult",
+                                event_time=datetime.datetime.now(datetime.UTC),
+                                data_version=constants.STEP_RESULT_EVENT_DATA_VERSION))
+                    else:
+                        logging.info(f'Request {req_id}: Malware scanning enabled, waiting for scan result')
+                elif new_status in [constants.STAGE_REJECTION_INPROGRESS, constants.STAGE_BLOCKING_INPROGRESS]:
+                    # Terminal transitions: emit StepResult immediately since no BlobCreated event will fire
+                    final_status = constants.STAGE_REJECTED if new_status == constants.STAGE_REJECTION_INPROGRESS else constants.STAGE_BLOCKED_BY_SCAN
+                    logging.info(f'Request {req_id}: Emitting StepResult for terminal transition {new_status} -> {final_status}')
+                    stepResultEvent.set(
+                        func.EventGridOutputEvent(
+                            id=str(uuid.uuid4()),
+                            data={"completed_step": new_status, "new_status": final_status, "request_id": req_id},
+                            subject=req_id,
+                            event_type="Airlock.StepResult",
+                            event_time=datetime.datetime.now(datetime.UTC),
+                            data_version=constants.STEP_RESULT_EVENT_DATA_VERSION))
+            else:
+                # Different storage account (e.g., core → workspace on import approval,
+                # workspace → core on export approval) - need to copy.
+                # BlobCreatedTrigger will fire when the copy completes and emit the StepResult,
+                # matching the v1 async pattern for large data transfers.
+                logging.info(f'Request {req_id}: Copying from {source_account} to {dest_account}')
+                create_container_with_metadata(dest_account, req_id, new_stage, workspace_id=effective_ws_id, request_type=request_type)
+                blob_operations.copy_data(source_account, dest_account, req_id)
+        else:
+            # Legacy mode: Copy data between storage accounts
+            logging.info('Request with id %s. requires data copy between storage accounts', req_id)
+            review_ws_id = request_properties.review_workspace_id
+            containers_metadata = get_source_dest_for_copy(new_status=new_status, previous_status=previous_status, request_type=request_type, short_workspace_id=ws_id, review_workspace_id=review_ws_id)
+            blob_operations.create_container(containers_metadata.dest_account_name, req_id)
+            blob_operations.copy_data(containers_metadata.source_account_name,
+                                      containers_metadata.dest_account_name, req_id)
         return
 
     # Other statuses which do not require data copy are dismissed as we don't need to do anything...
@@ -102,7 +181,7 @@ def is_require_data_copy(new_status: str):
     return False
 
 
-def get_source_dest_for_copy(new_status: str, previous_status: str, request_type: str, short_workspace_id: str) -> ContainersCopyMetadata:
+def get_source_dest_for_copy(new_status: str, previous_status: str, request_type: str, short_workspace_id: str, review_workspace_id: str = None) -> ContainersCopyMetadata:
     # sanity
     if is_require_data_copy(new_status) is False:
         raise Exception("Given new status is not supported")
@@ -115,7 +194,7 @@ def get_source_dest_for_copy(new_status: str, previous_status: str, request_type
         raise Exception(msg)
 
     source_account_name = get_storage_account(previous_status, request_type, short_workspace_id)
-    dest_account_name = get_storage_account_destination_for_copy(new_status, request_type, short_workspace_id)
+    dest_account_name = get_storage_account_destination_for_copy(new_status, request_type, short_workspace_id, review_workspace_id=review_workspace_id)
     return ContainersCopyMetadata(source_account_name, dest_account_name)
 
 
@@ -151,12 +230,14 @@ def get_storage_account(status: str, request_type: str, short_workspace_id: str)
     raise Exception(error_message)
 
 
-def get_storage_account_destination_for_copy(new_status: str, request_type: str, short_workspace_id: str) -> str:
+def get_storage_account_destination_for_copy(new_status: str, request_type: str, short_workspace_id: str, review_workspace_id: str = None) -> str:
     tre_id = _get_tre_id()
 
     if request_type == constants.IMPORT_TYPE:
         if new_status == constants.STAGE_SUBMITTED:
-            return constants.STORAGE_ACCOUNT_NAME_IMPORT_INPROGRESS + tre_id
+            # Import submit: copy to review workspace storage, or tre_id for legacy compatibility
+            dest_id = review_workspace_id if review_workspace_id else tre_id
+            return constants.STORAGE_ACCOUNT_NAME_IMPORT_INPROGRESS + dest_id
         elif new_status == constants.STAGE_APPROVAL_INPROGRESS:
             return constants.STORAGE_ACCOUNT_NAME_IMPORT_APPROVED + short_workspace_id
         elif new_status == constants.STAGE_REJECTION_INPROGRESS:
@@ -218,7 +299,13 @@ def set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_pr
 
 
 def get_request_files(request_properties: RequestProperties):
-    storage_account_name = get_storage_account(request_properties.previous_status, request_properties.type, request_properties.workspace_id)
+    use_metadata = request_properties.airlock_version >= 2
+    if use_metadata:
+        storage_account_name = airlock_storage_helper.get_storage_account_name_for_request(
+            request_properties.type, request_properties.previous_status, request_properties.workspace_id,
+            airlock_version=request_properties.airlock_version)
+    else:
+        storage_account_name = get_storage_account(request_properties.previous_status, request_properties.type, request_properties.workspace_id)
     return blob_operations.get_request_files(account_name=storage_account_name, request_id=request_properties.request_id)
 
 

airlock_processor/_version.py

@@ -1 +1 @@
-__version__ = "0.8.9"
+__version__ = "0.8.13"