diff --git a/airlock_processor/BlobCreatedTrigger/__init__.py b/airlock_processor/BlobCreatedTrigger/__init__.py index f119ad3eda..c8bda61b64 100644 --- a/airlock_processor/BlobCreatedTrigger/__init__.py +++ b/airlock_processor/BlobCreatedTrigger/__init__.py @@ -6,23 +6,28 @@ import os import azure.functions as func +from azure.storage.blob.aio import BlobServiceClient -from shared_code import constants, parsers -from shared_code.blob_operations import get_blob_info_from_topic_and_subject, get_blob_client_from_blob_info +from shared_code import constants, parsers, sb_helpers +from shared_code.blob_operations import get_blob_info_from_topic_and_subject, get_credential, get_account_url -def main(msg: func.ServiceBusMessage, - stepResultEvent: func.Out[func.EventGridOutputEvent], - dataDeletionEvent: func.Out[func.EventGridOutputEvent]): +async def main(msg: func.ServiceBusMessage, + stepResultEvent: func.Out[func.EventGridOutputEvent], + dataDeletionEvent: func.Out[func.EventGridOutputEvent]): logging.info("Python ServiceBus topic trigger processed message - A new blob was created!.") body = msg.get_body().decode('utf-8') - logging.info('Python ServiceBus queue trigger processed message: %s', body) + logging.info('Python ServiceBus topic trigger raw body: %s', body) + payload = await sb_helpers.receive_message_payload(body) + json_body = json.loads(payload) - json_body = json.loads(body) topic = json_body["topic"] request_id = re.search(r'/blobServices/default/containers/(.*?)/blobs', json_body["subject"]).group(1) + completed_step = None + new_status = None + # 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: @@ -35,7 +40,7 @@ def main(msg: func.ServiceBusMessage, # If malware scanning is enabled, the fact that the blob was created can be dismissed. # It will be consumed by the malware scanning service logging.info('Malware scanning is enabled. no action to perform.') - send_delete_event(dataDeletionEvent, json_body, request_id) + await send_delete_event(dataDeletionEvent, json_body, request_id) return else: logging.info('Malware scanning is disabled. Completing the submitted stage (moving to in_review).') @@ -57,31 +62,41 @@ def main(msg: func.ServiceBusMessage, new_status = constants.STAGE_BLOCKED_BY_SCAN # reply with a step completed event - 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)) + if completed_step and new_status: + data = {"completed_step": completed_step, "new_status": new_status, "request_id": request_id} + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) + stepResultEvent.set( + func.EventGridOutputEvent( + id=str(uuid.uuid4()), + data=offloaded_data, + 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(dataDeletionEvent, json_body, request_id) + await send_delete_event(dataDeletionEvent, json_body, request_id) -def send_delete_event(dataDeletionEvent: func.Out[func.EventGridOutputEvent], json_body, request_id): +async def send_delete_event(dataDeletionEvent: func.Out[func.EventGridOutputEvent], json_body, request_id): # check blob metadata to find the blob it was copied from - blob_client = get_blob_client_from_blob_info( - *get_blob_info_from_topic_and_subject(topic=json_body["topic"], subject=json_body["subject"])) - blob_metadata = blob_client.get_blob_properties()["metadata"] - copied_from = json.loads(blob_metadata["copied_from"]) - logging.info(f"copied from history: {copied_from}") + storage_account_name, container_name, blob_name = get_blob_info_from_topic_and_subject(topic=json_body["topic"], subject=json_body["subject"]) + + credential = await get_credential() + async with credential: + async with BlobServiceClient(account_url=get_account_url(storage_account_name), credential=credential) as blob_service_client: + blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name) + blob_properties = await blob_client.get_blob_properties() + blob_metadata = blob_properties["metadata"] + copied_from = json.loads(blob_metadata["copied_from"]) + logging.info(f"copied from history: {copied_from}") # signal that the container where we copied from can now be deleted + data = {"blob_to_delete": copied_from[-1]} # last container in copied_from is the one we just copied from + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) dataDeletionEvent.set( func.EventGridOutputEvent( id=str(uuid.uuid4()), - data={"blob_to_delete": copied_from[-1]}, # last container in copied_from is the one we just copied from + data=offloaded_data, subject=request_id, event_type="Airlock.DataDeletion", event_time=datetime.datetime.now(datetime.UTC), diff --git a/airlock_processor/DataDeletionTrigger/__init__.py b/airlock_processor/DataDeletionTrigger/__init__.py index 581094981d..bd383f744c 100644 --- a/airlock_processor/DataDeletionTrigger/__init__.py +++ b/airlock_processor/DataDeletionTrigger/__init__.py @@ -2,46 +2,49 @@ import json import azure.functions as func -from azure.storage.blob import BlobServiceClient +from azure.storage.blob.aio import BlobServiceClient -from shared_code import blob_operations +from shared_code import blob_operations, sb_helpers -def delete_blob_and_container_if_last_blob(blob_url: str): +async def delete_blob_and_container_if_last_blob(blob_url: str): storage_account_name, container_name, blob_name = blob_operations.get_blob_info_from_blob_url(blob_url=blob_url) - credential = blob_operations.get_credential() - blob_service_client = BlobServiceClient( - account_url=blob_operations.get_account_url(storage_account_name), - credential=credential) - container_client = blob_service_client.get_container_client(container_name) - - if not blob_name: - logging.info(f'No specific blob specified, deleting the entire container: {container_name}') - container_client.delete_container() - return - - # If it's the only blob in the container, we need to delete the container too - # Check how many blobs are in the container (note: this exhausts the generator) - blobs_num = sum(1 for _ in container_client.list_blobs()) - logging.info(f'Found {blobs_num} blobs in the container') - - # Deleting blob - logging.info(f'Deleting blob {blob_name}...') - blob_client = container_client.get_blob_client(blob_name) - blob_client.delete_blob() - - if blobs_num == 1: - # Need to delete the container too - logging.info(f'There was one blob in the container. Deleting container {container_name}...') - container_client.delete_container() - - -def main(msg: func.ServiceBusMessage): + credential = await blob_operations.get_credential() + async with BlobServiceClient( + account_url=blob_operations.get_account_url(storage_account_name), + credential=credential) as blob_service_client: + container_client = blob_service_client.get_container_client(container_name) + + if not blob_name: + logging.info(f'No specific blob specified, deleting the entire container: {container_name}') + await container_client.delete_container() + return + + # If it's the only blob in the container, we need to delete the container too + # Check how many blobs are in the container + blobs_num = 0 + async for _ in container_client.list_blobs(): + blobs_num += 1 + logging.info(f'Found {blobs_num} blobs in the container') + + # Deleting blob + logging.info(f'Deleting blob {blob_name}...') + blob_client = container_client.get_blob_client(blob_name) + await blob_client.delete_blob() + + if blobs_num == 1: + # Need to delete the container too + logging.info(f'There was one blob in the container. Deleting container {container_name}...') + await container_client.delete_container() + + +async def main(msg: func.ServiceBusMessage): body = msg.get_body().decode('utf-8') - logging.info(f'Python ServiceBus queue trigger processed message: {body}') - json_body = json.loads(body) + logging.info(f'Python ServiceBus queue trigger raw body: {body}') + payload = await sb_helpers.receive_message_payload(body) + json_body = json.loads(payload) blob_url = json_body["data"]["blob_to_delete"] logging.info(f'Blob to delete is {blob_url}') - delete_blob_and_container_if_last_blob(blob_url) + await delete_blob_and_container_if_last_blob(blob_url) diff --git a/airlock_processor/ScanResultTrigger/__init__.py b/airlock_processor/ScanResultTrigger/__init__.py index 1e4ffa4274..0b3a8a57a1 100644 --- a/airlock_processor/ScanResultTrigger/__init__.py +++ b/airlock_processor/ScanResultTrigger/__init__.py @@ -5,15 +5,16 @@ import uuid import json import os -from shared_code import constants, blob_operations, parsers +from shared_code import constants, blob_operations, parsers, sb_helpers -def main(msg: func.ServiceBusMessage, +async def main(msg: func.ServiceBusMessage, outputEvent: func.Out[func.EventGridOutputEvent]): logging.info("Python ServiceBus queue trigger processed message - Malware scan result arrived!") body = msg.get_body().decode('utf-8') - logging.info(f'Python ServiceBus queue trigger processed message: {body}') + logging.info(f'Python ServiceBus queue trigger raw body: {body}') + payload = await sb_helpers.receive_message_payload(body) status_message = None try: @@ -31,7 +32,7 @@ def main(msg: func.ServiceBusMessage, raise Exception(error_msg) try: - json_body = json.loads(body) + json_body = json.loads(payload) blob_uri = json_body["data"]["blobUri"] verdict = json_body["data"]["scanResultType"] except KeyError as e: @@ -53,11 +54,13 @@ def main(msg: func.ServiceBusMessage, status_message = verdict # Send the event to indicate this step is done (and to request a new status change) + data = {"completed_step": completed_step, "new_status": new_status, "request_id": request_id, "status_message": status_message} + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) outputEvent.set( func.EventGridOutputEvent( id=str(uuid.uuid4()), - data={"completed_step": completed_step, "new_status": new_status, "request_id": request_id, "status_message": status_message}, + data=offloaded_data, subject=request_id, event_type="Airlock.StepResult", event_time=datetime.datetime.now(datetime.UTC), - data_version=constants.STEP_RESULT_EVENT_DATA_VERSION)) + data_version=constants.STEP_RESULT_EVENT_DATA_VERSION)) \ No newline at end of file diff --git a/airlock_processor/StatusChangedQueueTrigger/__init__.py b/airlock_processor/StatusChangedQueueTrigger/__init__.py index db64d72a48..683f7238e2 100644 --- a/airlock_processor/StatusChangedQueueTrigger/__init__.py +++ b/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, sb_helpers from pydantic import BaseModel, parse_obj_as @@ -30,21 +30,21 @@ def __init__(self, source_account_name: str, dest_account_name: str): self.dest_account_name = dest_account_name -def main(msg: func.ServiceBusMessage, stepResultEvent: func.Out[func.EventGridOutputEvent], dataDeletionEvent: func.Out[func.EventGridOutputEvent]): +async def main(msg: func.ServiceBusMessage, stepResultEvent: func.Out[func.EventGridOutputEvent], dataDeletionEvent: func.Out[func.EventGridOutputEvent]): try: - request_properties = extract_properties(msg) - request_files = get_request_files(request_properties) if request_properties.new_status == constants.STAGE_SUBMITTED else None - handle_status_changed(request_properties, stepResultEvent, dataDeletionEvent, request_files) + request_properties = await extract_properties(msg) + request_files = await get_request_files(request_properties) if request_properties.new_status == constants.STAGE_SUBMITTED else None + await handle_status_changed(request_properties, stepResultEvent, dataDeletionEvent, request_files) except NoFilesInRequestException: - set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.NO_FILES_IN_REQUEST_MESSAGE, request_files=request_files) + await set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.NO_FILES_IN_REQUEST_MESSAGE, request_files=request_files) except TooManyFilesInRequestException: - set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.TOO_MANY_FILES_IN_REQUEST_MESSAGE, request_files=request_files) + await set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.TOO_MANY_FILES_IN_REQUEST_MESSAGE, request_files=request_files) except Exception: - set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.UNKNOWN_REASON_MESSAGE, request_files=request_files) + await set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason=constants.UNKNOWN_REASON_MESSAGE, request_files=request_files) -def handle_status_changed(request_properties: RequestProperties, stepResultEvent: func.Out[func.EventGridOutputEvent], dataDeletionEvent: func.Out[func.EventGridOutputEvent], request_files): +async def handle_status_changed(request_properties: RequestProperties, stepResultEvent: func.Out[func.EventGridOutputEvent], dataDeletionEvent: func.Out[func.EventGridOutputEvent], request_files): new_status = request_properties.new_status previous_status = request_properties.previous_status req_id = request_properties.request_id @@ -55,34 +55,35 @@ def handle_status_changed(request_properties: RequestProperties, stepResultEvent 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) + await 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) 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) + await set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_properties, container_url=container_to_delete_url) return if new_status == constants.STAGE_SUBMITTED: - set_output_event_to_report_request_files(stepResultEvent, request_properties, request_files) + await 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) + await blob_operations.create_container(containers_metadata.dest_account_name, req_id) + await 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... -def extract_properties(msg: func.ServiceBusMessage) -> RequestProperties: +async def extract_properties(msg: func.ServiceBusMessage) -> RequestProperties: try: body = msg.get_body().decode('utf-8') - logging.debug('Python ServiceBus queue trigger processed message: %s', body) - json_body = json.loads(body) + logging.debug('Python ServiceBus queue trigger raw body: %s', body) + payload = await sb_helpers.receive_message_payload(body) + json_body = json.loads(payload) result = parse_obj_as(RequestProperties, json_body["data"]) if not result: raise Exception("Failed parsing request properties") @@ -179,36 +180,42 @@ def get_storage_account_destination_for_copy(new_status: str, request_type: str, raise Exception(error_message) -def set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason, request_files): +async def set_output_event_to_report_failure(stepResultEvent, request_properties, failure_reason, request_files): logging.exception(f"Failed processing Airlock request with ID: '{request_properties.request_id}', changing request status to '{constants.STAGE_FAILED}'.") + data = {"completed_step": request_properties.new_status, "new_status": constants.STAGE_FAILED, "request_id": request_properties.request_id, "request_files": request_files, "status_message": failure_reason} + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) stepResultEvent.set( func.EventGridOutputEvent( id=str(uuid.uuid4()), - data={"completed_step": request_properties.new_status, "new_status": constants.STAGE_FAILED, "request_id": request_properties.request_id, "request_files": request_files, "status_message": failure_reason}, + data=offloaded_data, subject=request_properties.request_id, event_type="Airlock.StepResult", event_time=datetime.datetime.now(datetime.UTC), data_version=constants.STEP_RESULT_EVENT_DATA_VERSION)) -def set_output_event_to_report_request_files(stepResultEvent, request_properties, request_files): +async def set_output_event_to_report_request_files(stepResultEvent, request_properties, request_files): logging.info(f'Sending file enumeration result for request with ID: {request_properties.request_id} result: {request_files}') + data = {"completed_step": request_properties.new_status, "request_id": request_properties.request_id, "request_files": request_files} + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) stepResultEvent.set( func.EventGridOutputEvent( id=str(uuid.uuid4()), - data={"completed_step": request_properties.new_status, "request_id": request_properties.request_id, "request_files": request_files}, + data=offloaded_data, subject=request_properties.request_id, event_type="Airlock.StepResult", event_time=datetime.datetime.now(datetime.UTC), data_version=constants.STEP_RESULT_EVENT_DATA_VERSION)) -def set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_properties, container_url): +async def set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_properties, container_url): logging.info(f'Sending container deletion event for request ID: {request_properties.request_id}. container URL: {container_url}') + data = {"blob_to_delete": container_url} + offloaded_data = await sb_helpers.wrap_payload_for_offloading(data) dataDeletionEvent.set( func.EventGridOutputEvent( id=str(uuid.uuid4()), - data={"blob_to_delete": container_url}, + data=offloaded_data, subject=request_properties.request_id, event_type="Airlock.DataDeletion", event_time=datetime.datetime.now(datetime.UTC), @@ -217,9 +224,9 @@ def set_output_event_to_trigger_container_deletion(dataDeletionEvent, request_pr ) -def get_request_files(request_properties: RequestProperties): +async def get_request_files(request_properties: RequestProperties): 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) + return await blob_operations.get_request_files(account_name=storage_account_name, request_id=request_properties.request_id) def _get_tre_id(): diff --git a/airlock_processor/shared_code/blob_operations.py b/airlock_processor/shared_code/blob_operations.py index 211b1925aa..8438911144 100644 --- a/airlock_processor/shared_code/blob_operations.py +++ b/airlock_processor/shared_code/blob_operations.py @@ -6,8 +6,9 @@ from typing import Tuple from azure.core.exceptions import ResourceExistsError -from azure.identity import DefaultAzureCredential -from azure.storage.blob import ContainerSasPermissions, generate_container_sas, BlobServiceClient +from azure.identity.aio import DefaultAzureCredential +from azure.storage.blob.aio import BlobServiceClient +from azure.storage.blob import ContainerSasPermissions, generate_container_sas from exceptions import NoFilesInRequestException, TooManyFilesInRequestException @@ -16,94 +17,97 @@ def get_account_url(account_name: str) -> str: return f"https://{account_name}.blob.{get_storage_endpoint_suffix()}/" -def get_blob_client_from_blob_info(storage_account_name: str, container_name: str, blob_name: str): - source_blob_service_client = BlobServiceClient(account_url=get_account_url(storage_account_name), - credential=get_credential()) - source_container_client = source_blob_service_client.get_container_client(container_name) - return source_container_client.get_blob_client(blob_name) - - -def create_container(account_name: str, request_id: str): +async def create_container(account_name: str, request_id: str): try: container_name = request_id - blob_service_client = BlobServiceClient(account_url=get_account_url(account_name), - credential=get_credential()) - blob_service_client.create_container(container_name) - logging.info(f'Container created for request id: {request_id}.') + credential = await get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url=get_account_url(account_name), + credential=credential) + async with blob_service_client: + await blob_service_client.create_container(container_name) + logging.info(f'Container created for request id: {request_id}.') except ResourceExistsError: logging.info(f'Did not create a new container. Container already exists for request id: {request_id}.') -def get_request_files(account_name: str, request_id: str) -> list: +async def get_request_files(account_name: str, request_id: str) -> list: files = [] - blob_service_client = BlobServiceClient(account_url=get_account_url(account_name), credential=get_credential()) - container_client = blob_service_client.get_container_client(container=request_id) + credential = await get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url=get_account_url(account_name), credential=credential) + async with blob_service_client: + container_client = blob_service_client.get_container_client(container=request_id) - for blob in container_client.list_blobs(): - files.append({"name": blob.name, "size": blob.size}) + async for blob in container_client.list_blobs(): + files.append({"name": blob.name, "size": blob.size}) return files -def copy_data(source_account_name: str, destination_account_name: str, request_id: str): - credential = get_credential() - container_name = request_id - - source_blob_service_client = BlobServiceClient(account_url=get_account_url(source_account_name), - credential=credential) - source_container_client = source_blob_service_client.get_container_client(container_name) - - # Check that we are copying exactly one blob - found_blobs = 0 - blob_name = "" - for blob in source_container_client.list_blobs(): - blob_name = blob.name - if found_blobs > 0: - msg = "Request with id {} contains more than 1 file. flow aborted.".format(request_id) - logging.error(msg) - raise TooManyFilesInRequestException(msg) - found_blobs += 1 - - if found_blobs == 0: - msg = "Request with id {} did not contain any files. flow aborted.".format(request_id) - logging.error(msg) - raise NoFilesInRequestException(msg) - - # token geneation with expiry of 1 hour. since its not shared, we can leave it to expire (no need to track/delete) - # Remove sas token if not needed: https://github.com/microsoft/AzureTRE/issues/2034 - start = datetime.now(UTC) - timedelta(minutes=15) - expiry = datetime.now(UTC) + timedelta(hours=1) - udk = source_blob_service_client.get_user_delegation_key(key_start_time=start, key_expiry_time=expiry) - - sas_token = generate_container_sas(container_name=container_name, - account_name=source_account_name, - user_delegation_key=udk, - permission=ContainerSasPermissions(read=True), - start=start, - expiry=expiry) - - source_blob = source_container_client.get_blob_client(blob_name) - source_url = f'{source_blob.url}?{sas_token}' - - # Set metadata to include the blob url that it is copied from - metadata = source_blob.get_blob_properties()["metadata"] - copied_from = json.loads(metadata["copied_from"]) if "copied_from" in metadata else [] - metadata["copied_from"] = json.dumps(copied_from + [source_blob.url]) - - # Copy files - dest_blob_service_client = BlobServiceClient(account_url=get_account_url(destination_account_name), - credential=credential) - copied_blob = dest_blob_service_client.get_blob_client(container_name, source_blob.blob_name) - copy = copied_blob.start_copy_from_url(source_url, metadata=metadata) - - try: - logging.info("Copy operation returned 'copy_id': '%s', 'copy_status': '%s'", copy["copy_id"], - copy["copy_status"]) - except KeyError as e: - logging.error(f"Failed getting operation id and status {e}") - +async def copy_data(source_account_name: str, destination_account_name: str, request_id: str): + credential = await get_credential() + async with credential: + container_name = request_id -def get_credential() -> DefaultAzureCredential: + source_blob_service_client = BlobServiceClient(account_url=get_account_url(source_account_name), + credential=credential) + async with source_blob_service_client: + source_container_client = source_blob_service_client.get_container_client(container_name) + + # Check that we are copying exactly one blob + found_blobs = 0 + blob_name = "" + async for blob in source_container_client.list_blobs(): + blob_name = blob.name + if found_blobs > 0: + msg = "Request with id {} contains more than 1 file. flow aborted.".format(request_id) + logging.error(msg) + raise TooManyFilesInRequestException(msg) + found_blobs += 1 + + if found_blobs == 0: + msg = "Request with id {} did not contain any files. flow aborted.".format(request_id) + logging.error(msg) + raise NoFilesInRequestException(msg) + + # token geneation with expiry of 1 hour. since its not shared, we can leave it to expire (no need to track/delete) + # Remove sas token if not needed: https://github.com/microsoft/AzureTRE/issues/2034 + start = datetime.now(UTC) - timedelta(minutes=15) + expiry = datetime.now(UTC) + timedelta(hours=1) + udk = await source_blob_service_client.get_user_delegation_key(key_start_time=start, key_expiry_time=expiry) + + sas_token = generate_container_sas(container_name=container_name, + account_name=source_account_name, + user_delegation_key=udk, + permission=ContainerSasPermissions(read=True), + start=start, + expiry=expiry) + + source_blob = source_container_client.get_blob_client(blob_name) + source_url = f'{source_blob.url}?{sas_token}' + + # Set metadata to include the blob url that it is copied from + props = await source_blob.get_blob_properties() + metadata = props["metadata"] + copied_from = json.loads(metadata["copied_from"]) if "copied_from" in metadata else [] + metadata["copied_from"] = json.dumps(copied_from + [source_blob.url]) + + # Copy files + dest_blob_service_client = BlobServiceClient(account_url=get_account_url(destination_account_name), + credential=credential) + async with dest_blob_service_client: + copied_blob = dest_blob_service_client.get_blob_client(container_name, source_blob.blob_name) + copy = await copied_blob.start_copy_from_url(source_url, metadata=metadata) + + try: + logging.info("Copy operation returned 'copy_id': '%s', 'copy_status': '%s'", copy["copy_id"], + copy["copy_status"]) + except KeyError as e: + logging.error(f"Failed getting operation id and status {e}") + + +async def get_credential() -> DefaultAzureCredential: managed_identity = os.environ.get("MANAGED_IDENTITY_CLIENT_ID") if managed_identity: logging.info("using the Airlock processor's managed identity to get credentials.") diff --git a/airlock_processor/shared_code/constants.py b/airlock_processor/shared_code/constants.py index 277312d1cb..a90d21922e 100644 --- a/airlock_processor/shared_code/constants.py +++ b/airlock_processor/shared_code/constants.py @@ -1,3 +1,11 @@ +import os + +# SB Offloading +SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME = os.environ.get("SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME", "") +SERVICE_BUS_MESSAGES_CONTAINER_NAME = os.environ.get("SERVICE_BUS_MESSAGES_CONTAINER_NAME", "sb-messages") +STORAGE_ENDPOINT_SUFFIX = os.environ.get("STORAGE_ENDPOINT_SUFFIX", "core.windows.net") +SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD = int(os.environ.get("SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD", "200000")) + # RG CORE_RESOURCE_GROUP_NAME = "rg-{}" WORKSPACE_RESOURCE_GROUP_NAME = "rg-{}-ws-{}" diff --git a/airlock_processor/shared_code/sb_helpers.py b/airlock_processor/shared_code/sb_helpers.py new file mode 100644 index 0000000000..ccc1e1dc20 --- /dev/null +++ b/airlock_processor/shared_code/sb_helpers.py @@ -0,0 +1,65 @@ +import json +import uuid +import os +import logging +from azure.storage.blob.aio import BlobServiceClient +from azure.identity.aio import DefaultAzureCredential +from shared_code import constants + + +async def _get_credential(): + msi_id = os.environ.get("MANAGED_IDENTITY_CLIENT_ID") + return DefaultAzureCredential(managed_identity_client_id=msi_id) if msi_id else DefaultAzureCredential() + + +async def receive_message_payload(msg_body: str) -> str: + try: + body_json = json.loads(msg_body) + if "claim_check" in body_json: + blob_path = body_json["claim_check"] + logging.info(f"Message has claim check: {blob_path}. Downloading from blob storage.") + return await _download_from_blob(blob_path) + except json.JSONDecodeError: + pass + + return msg_body + + +async def _download_from_blob(blob_path: str) -> str: + account_url = f"https://{constants.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME}.blob.{constants.STORAGE_ENDPOINT_SUFFIX}" + container_name, blob_name = blob_path.split("/", 1) + credential = await _get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name) + download_stream = await blob_client.download_blob() + content = await download_stream.readall() + return content.decode("utf-8") + + +async def wrap_payload_for_offloading(payload: dict) -> dict: + """ + Checks if the payload exceeds the threshold and offloads to blob storage if it does. + Returns the original payload or a reference to the offloaded blob. + """ + payload_str = json.dumps(payload) + if constants.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME and len(payload_str) > constants.SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD: + logging.info(f"Payload size {len(payload_str)} exceeds threshold. Offloading to blob storage.") + blob_url = await _offload_to_blob(payload_str) + return { + "claim_check": blob_url + } + return payload + + +async def _offload_to_blob(content: str) -> str: + account_url = f"https://{constants.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME}.blob.{constants.STORAGE_ENDPOINT_SUFFIX}" + credential = await _get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_name = f"msg-{uuid.uuid4()}.json" + blob_client = blob_service_client.get_blob_client(container=constants.SERVICE_BUS_MESSAGES_CONTAINER_NAME, blob=blob_name) + await blob_client.upload_blob(content) + return f"{constants.SERVICE_BUS_MESSAGES_CONTAINER_NAME}/{blob_name}" diff --git a/airlock_processor/tests/shared_code/test_blob_operations.py b/airlock_processor/tests/shared_code/test_blob_operations.py index c3a921f0b7..b4461800c9 100644 --- a/airlock_processor/tests/shared_code/test_blob_operations.py +++ b/airlock_processor/tests/shared_code/test_blob_operations.py @@ -1,9 +1,9 @@ from collections import namedtuple import json import pytest -from mock import MagicMock, patch +from mock import AsyncMock, MagicMock, patch -from shared_code.blob_operations import get_blob_info_from_topic_and_subject, get_blob_info_from_blob_url, copy_data, get_blob_url, get_storage_endpoint_suffix +import shared_code.blob_operations as blob_ops from exceptions import TooManyFilesInRequestException, NoFilesInRequestException @@ -17,78 +17,113 @@ def test_get_blob_info_from_topic_and_subject(self): topic = "/subscriptions/SUB_ID/resourceGroups/RG_NAME/providers/Microsoft.Storage/storageAccounts/ST_ACC_NAME" subject = "/blobServices/default/containers/c144728c-3c69-4a58-afec-48c2ec8bfd45/blobs/BLOB" - storage_account_name, container_name, blob_name = get_blob_info_from_topic_and_subject(topic=topic, subject=subject) + storage_account_name, container_name, blob_name = blob_ops.get_blob_info_from_topic_and_subject(topic=topic, subject=subject) assert storage_account_name == "ST_ACC_NAME" assert container_name == "c144728c-3c69-4a58-afec-48c2ec8bfd45" assert blob_name == "BLOB" def test_get_blob_info_from_url(self): - url = f"https://stalimextest.blob.{get_storage_endpoint_suffix()}/c144728c-3c69-4a58-afec-48c2ec8bfd45/test_dataset.txt" + url = f"https://stalimextest.blob.{blob_ops.get_storage_endpoint_suffix()}/c144728c-3c69-4a58-afec-48c2ec8bfd45/test_dataset.txt" - storage_account_name, container_name, blob_name = get_blob_info_from_blob_url(blob_url=url) + storage_account_name, container_name, blob_name = blob_ops.get_blob_info_from_blob_url(blob_url=url) assert storage_account_name == "stalimextest" assert container_name == "c144728c-3c69-4a58-afec-48c2ec8bfd45" assert blob_name == "test_dataset.txt" - @patch("shared_code.blob_operations.BlobServiceClient") - def test_copy_data_fails_if_too_many_blobs_to_copy(self, mock_blob_service_client): - mock_blob_service_client().get_container_client().list_blobs = MagicMock(return_value=[get_test_blob()("a"), get_test_blob()("b")]) + @pytest.mark.asyncio + async def test_copy_data_fails_if_too_many_blobs_to_copy(self): + with patch("shared_code.blob_operations.get_credential", new_callable=AsyncMock) as mock_get_credential, \ + patch("shared_code.blob_operations.BlobServiceClient") as mock_blob_service_client: - with pytest.raises(TooManyFilesInRequestException): - copy_data("source_acc", "dest_acc", "req_id") + mock_get_credential.return_value.__aenter__.return_value = MagicMock() - @patch("shared_code.blob_operations.BlobServiceClient") - def test_copy_data_fails_if_no_blobs_to_copy(self, mock_blob_service_client): - mock_blob_service_client().get_container_client().list_blobs = MagicMock(return_value=[]) + mock_client_instance = mock_blob_service_client.return_value + mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance) - with pytest.raises(NoFilesInRequestException): - copy_data("source_acc", "dest_acc", "req_id") + mock_container_client = MagicMock() + mock_client_instance.get_container_client.return_value = mock_container_client - @patch("shared_code.blob_operations.BlobServiceClient") - @patch("shared_code.blob_operations.generate_container_sas", return_value="sas") - def test_copy_data_adds_copied_from_metadata(self, _, mock_blob_service_client): - source_url = f"http://storageacct.blob.{get_storage_endpoint_suffix()}/container/blob" + mock_list_blobs = AsyncMock() + mock_list_blobs.__aiter__.return_value = [get_test_blob()("a"), get_test_blob()("b")] + mock_container_client.list_blobs.return_value = mock_list_blobs - # Check for two scenarios: when there's no copied_from history in metadata, and when there is some - for source_metadata, dest_metadata in [ - ({"a": "b"}, {"a": "b", "copied_from": json.dumps([source_url])}), - ({"a": "b", "copied_from": json.dumps(["old_url"])}, {"a": "b", "copied_from": json.dumps(["old_url", source_url])}) - ]: - source_blob_client_mock = MagicMock() - source_blob_client_mock.url = source_url - source_blob_client_mock.get_blob_properties = MagicMock(return_value={"metadata": source_metadata}) + with pytest.raises(TooManyFilesInRequestException): + await blob_ops.copy_data("source_acc", "dest_acc", "req_id") - dest_blob_client_mock = MagicMock() - dest_blob_client_mock.bla = "bla" - dest_blob_client_mock.start_copy_from_url = MagicMock(return_value={"copy_id": "123", "copy_status": "status"}) + @pytest.mark.asyncio + async def test_copy_data_fails_if_no_blobs_to_copy(self): + with patch("shared_code.blob_operations.get_credential", new_callable=AsyncMock) as mock_get_credential, \ + patch("shared_code.blob_operations.BlobServiceClient") as mock_blob_service_client: - # Set source blob mock - mock_blob_service_client().get_container_client().get_blob_client = MagicMock(return_value=source_blob_client_mock) - # Set dest blob mock - mock_blob_service_client().get_blob_client = MagicMock(return_value=dest_blob_client_mock) + mock_get_credential.return_value.__aenter__.return_value = MagicMock() - # Any additional mocks for the copy_data method to work - mock_blob_service_client().get_user_delegation_key = MagicMock(return_value="key") - mock_blob_service_client().get_container_client().list_blobs = MagicMock(return_value=[get_test_blob()("a")]) + mock_client_instance = mock_blob_service_client.return_value + mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance) - copy_data("source_acc", "dest_acc", "req_id") + mock_container_client = MagicMock() + mock_client_instance.get_container_client.return_value = mock_container_client - # Check that copied_from field was set correctly in the metadata - dest_blob_client_mock.start_copy_from_url.assert_called_with(f"{source_url}?sas", metadata=dest_metadata) + mock_list_blobs = AsyncMock() + mock_list_blobs.__aiter__.return_value = [] + mock_container_client.list_blobs.return_value = mock_list_blobs + + with pytest.raises(NoFilesInRequestException): + await blob_ops.copy_data("source_acc", "dest_acc", "req_id") + + @pytest.mark.asyncio + async def test_copy_data_adds_copied_from_metadata(self): + with patch("shared_code.blob_operations.get_credential", new_callable=AsyncMock) as mock_get_credential, \ + patch("shared_code.blob_operations.BlobServiceClient") as mock_blob_service_client, \ + patch("shared_code.blob_operations.generate_container_sas", return_value="sas"): + + mock_get_credential.return_value.__aenter__.return_value = MagicMock() + + source_url = f"https://storageacct.blob.{blob_ops.get_storage_endpoint_suffix()}/container/blob" + + # Check for two scenarios: when there's no copied_from history in metadata, and when there is some + for source_metadata, dest_metadata in [ + ({"a": "b"}, {"a": "b", "copied_from": json.dumps([source_url])}), + ({"a": "b", "copied_from": json.dumps(["old_url"])}, {"a": "b", "copied_from": json.dumps(["old_url", source_url])}) + ]: + mock_client_instance = mock_blob_service_client.return_value + mock_client_instance.__aenter__ = AsyncMock(return_value=mock_client_instance) + + source_blob_client_mock = MagicMock() + source_blob_client_mock.url = source_url + source_blob_client_mock.get_blob_properties = AsyncMock(return_value={"metadata": source_metadata}) + + dest_blob_client_mock = MagicMock() + dest_blob_client_mock.start_copy_from_url = AsyncMock(return_value={"copy_id": "123", "copy_status": "status"}) + + mock_container_client = MagicMock() + mock_client_instance.get_container_client.return_value = mock_container_client + mock_container_client.get_blob_client.return_value = source_blob_client_mock + + mock_list_blobs = AsyncMock() + mock_list_blobs.__aiter__.return_value = [get_test_blob()("a")] + mock_container_client.list_blobs.return_value = mock_list_blobs + + mock_client_instance.get_user_delegation_key = AsyncMock(return_value="key") + mock_client_instance.get_blob_client.return_value = dest_blob_client_mock + + await blob_ops.copy_data("source_acc", "dest_acc", "req_id") + + # Check that copied_from field was set correctly in the metadata + dest_blob_client_mock.start_copy_from_url.assert_called_with(f"{source_url}?sas", metadata=dest_metadata) def test_get_blob_url_should_return_blob_url(self): account_name = "account" container_name = "container" blob_name = "blob" - blob_url = get_blob_url(account_name, container_name, blob_name) - assert blob_url == f"https://{account_name}.blob.{get_storage_endpoint_suffix()}/{container_name}/{blob_name}" + blob_url = blob_ops.get_blob_url(account_name, container_name, blob_name) + assert blob_url == f"https://{account_name}.blob.{blob_ops.get_storage_endpoint_suffix()}/{container_name}/{blob_name}" def test_get_blob_url_without_blob_name_should_return_container_url(self): account_name = "account" container_name = "container" - blob_url = get_blob_url(account_name, container_name) - assert blob_url == f"https://{account_name}.blob.{get_storage_endpoint_suffix()}/{container_name}/" + blob_url = blob_ops.get_blob_url(account_name, container_name) + assert blob_url == f"https://{account_name}.blob.{blob_ops.get_storage_endpoint_suffix()}/{container_name}/" diff --git a/airlock_processor/tests/test_data_deletion_trigger.py b/airlock_processor/tests/test_data_deletion_trigger.py index 65634f09f7..20505dd526 100644 --- a/airlock_processor/tests/test_data_deletion_trigger.py +++ b/airlock_processor/tests/test_data_deletion_trigger.py @@ -1,33 +1,68 @@ -from mock import patch, MagicMock +from mock import patch, MagicMock, AsyncMock +import pytest from DataDeletionTrigger import delete_blob_and_container_if_last_blob from shared_code.blob_operations import get_storage_endpoint_suffix +@pytest.mark.asyncio class TestDataDeletionTrigger(): + @patch("DataDeletionTrigger.blob_operations.get_credential") @patch("DataDeletionTrigger.BlobServiceClient") - def test_delete_blob_and_container_if_last_blob_deletes_container(self, mock_blob_service_client): + async def test_delete_blob_and_container_if_last_blob_deletes_container(self, mock_blob_service_client, _): blob_url = f"https://stalimextest.blob.{get_storage_endpoint_suffix()}/c144728c-3c69-4a58-afec-48c2ec8bfd45/test_dataset.txt" - mock_blob_service_client().get_container_client().list_blobs = MagicMock(return_value=["blob"]) + mock_service_instance = MagicMock() + mock_blob_service_client.return_value.__aenter__.return_value = mock_service_instance - delete_blob_and_container_if_last_blob(blob_url) + mock_container_client = MagicMock() + mock_service_instance.get_container_client.return_value = mock_container_client - mock_blob_service_client().get_container_client().delete_container.assert_called_once() + mock_list_blobs = AsyncMock() + mock_list_blobs.__aiter__.return_value = ["blob"] + mock_container_client.list_blobs.return_value = mock_list_blobs + mock_container_client.delete_container = AsyncMock() + mock_container_client.get_blob_client.return_value.delete_blob = AsyncMock() + + await delete_blob_and_container_if_last_blob(blob_url) + + mock_container_client.delete_container.assert_called_once() + + @patch("DataDeletionTrigger.blob_operations.get_credential") @patch("DataDeletionTrigger.BlobServiceClient") - def test_delete_blob_and_container_if_last_blob_doesnt_delete_container(self, mock_blob_service_client): + async def test_delete_blob_and_container_if_last_blob_doesnt_delete_container(self, mock_blob_service_client, _): blob_url = f"https://stalimextest.blob.{get_storage_endpoint_suffix()}/c144728c-3c69-4a58-afec-48c2ec8bfd45/test_dataset.txt" - mock_blob_service_client().get_container_client().list_blobs = MagicMock(return_value=["blob1", "blob2"]) + mock_service_instance = MagicMock() + mock_blob_service_client.return_value.__aenter__.return_value = mock_service_instance + + mock_container_client = MagicMock() + mock_service_instance.get_container_client.return_value = mock_container_client + + mock_list_blobs = AsyncMock() + mock_list_blobs.__aiter__.return_value = ["blob1", "blob2"] + mock_container_client.list_blobs.return_value = mock_list_blobs - delete_blob_and_container_if_last_blob(blob_url) + mock_container_client.delete_container = AsyncMock() + mock_container_client.get_blob_client.return_value.delete_blob = AsyncMock() - mock_blob_service_client().get_container_client().delete_container.assert_not_called() + await delete_blob_and_container_if_last_blob(blob_url) + mock_container_client.delete_container.assert_not_called() + + @patch("DataDeletionTrigger.blob_operations.get_credential") @patch("DataDeletionTrigger.BlobServiceClient") - def test_delete_blob_and_container_if_last_blob_deletes_container_if_no_blob_specified(self, mock_blob_service_client): + async def test_delete_blob_and_container_if_last_blob_deletes_container_if_no_blob_specified(self, mock_blob_service_client, _): blob_url = f"https://stalimextest.blob.{get_storage_endpoint_suffix()}/c144728c-3c69-4a58-afec-48c2ec8bfd45/" - delete_blob_and_container_if_last_blob(blob_url) - mock_blob_service_client().get_container_client().delete_container.assert_called_once() + + mock_service_instance = MagicMock() + mock_blob_service_client.return_value.__aenter__.return_value = mock_service_instance + + mock_container_client = MagicMock() + mock_service_instance.get_container_client.return_value = mock_container_client + mock_container_client.delete_container = AsyncMock() + + await delete_blob_and_container_if_last_blob(blob_url) + mock_container_client.delete_container.assert_called_once() diff --git a/airlock_processor/tests/test_status_change_queue_trigger.py b/airlock_processor/tests/test_status_change_queue_trigger.py index 4ce518c09d..6e7cf46190 100644 --- a/airlock_processor/tests/test_status_change_queue_trigger.py +++ b/airlock_processor/tests/test_status_change_queue_trigger.py @@ -9,38 +9,44 @@ from shared_code import constants +@pytest.mark.asyncio class TestPropertiesExtraction(): - def test_extract_prop_valid_body_return_all_values(self): + async def test_extract_prop_valid_body_return_all_values(self): message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"456\" ,\"previous_status\":\"789\" , \"type\":\"101112\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - req_prop = extract_properties(message) + req_prop = await extract_properties(message) assert req_prop.request_id == "123" assert req_prop.new_status == "456" assert req_prop.previous_status == "789" assert req_prop.type == "101112" assert req_prop.workspace_id == "ws1" - def test_extract_prop_missing_arg_throws(self): + async def test_extract_prop_missing_arg_throws(self): message_body = "{ \"data\": { \"status\":\"456\" , \"type\":\"789\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - pytest.raises(ValidationError, extract_properties, message) + with pytest.raises(ValidationError): + await extract_properties(message) message_body = "{ \"data\": { \"request_id\":\"123\", \"type\":\"789\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - pytest.raises(ValidationError, extract_properties, message) + with pytest.raises(ValidationError): + await extract_properties(message) message_body = "{ \"data\": { \"request_id\":\"123\",\"status\":\"456\" , \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - pytest.raises(ValidationError, extract_properties, message) + with pytest.raises(ValidationError): + await extract_properties(message) message_body = "{ \"data\": { \"request_id\":\"123\",\"status\":\"456\" , \"type\":\"789\" }}" message = _mock_service_bus_message(body=message_body) - pytest.raises(ValidationError, extract_properties, message) + with pytest.raises(ValidationError): + await extract_properties(message) - def test_extract_prop_invalid_json_throws(self): + async def test_extract_prop_invalid_json_throws(self): message_body = "Hi" message = _mock_service_bus_message(body=message_body) - pytest.raises(JSONDecodeError, extract_properties, message) + with pytest.raises(JSONDecodeError): + await extract_properties(message) class TestDataCopyProperties(): @@ -60,62 +66,69 @@ def test_only_specific_status_are_triggering_copy(self): assert is_require_data_copy("blocking_in_progress") def test_wrong_status_raises_when_getting_storage_account_properties(self): - pytest.raises(Exception, get_source_dest_for_copy, "Miaow", "import") + with pytest.raises(Exception): + get_source_dest_for_copy("Miaow", "import", "type", "ws1") def test_wrong_type_raises_when_getting_storage_account_properties(self): - pytest.raises(Exception, get_source_dest_for_copy, "accepted", "somethingelse") + with pytest.raises(Exception): + get_source_dest_for_copy("accepted", "somethingelse", "type", "ws1") +@pytest.mark.asyncio class TestFileEnumeration(): @patch("StatusChangedQueueTrigger.set_output_event_to_report_request_files") @patch("StatusChangedQueueTrigger.get_request_files") @patch("StatusChangedQueueTrigger.is_require_data_copy", return_value=False) @patch.dict(os.environ, {"TRE_ID": "tre-id"}, clear=True) - def test_get_request_files_should_be_called_on_submit_stage(self, _, mock_get_request_files, mock_set_output_event_to_report_request_files): + async def test_get_request_files_should_be_called_on_submit_stage(self, _, mock_get_request_files, mock_set_output_event_to_report_request_files): message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"submitted\" ,\"previous_status\":\"draft\" , \"type\":\"export\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) + mock_get_request_files.return_value = [] + await main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) assert mock_get_request_files.called assert mock_set_output_event_to_report_request_files.called @patch("StatusChangedQueueTrigger.set_output_event_to_report_failure") @patch("StatusChangedQueueTrigger.get_request_files") @patch("StatusChangedQueueTrigger.handle_status_changed") - def test_get_request_files_should_not_be_called_if_new_status_is_not_submit(self, _, mock_get_request_files, mock_set_output_event_to_report_failure): + async def test_get_request_files_should_not_be_called_if_new_status_is_not_submit(self, _, mock_get_request_files, mock_set_output_event_to_report_failure): message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"fake-status\" ,\"previous_status\":\"None\" , \"type\":\"export\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) + await main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) assert not mock_get_request_files.called assert not mock_set_output_event_to_report_failure.called @patch("StatusChangedQueueTrigger.set_output_event_to_report_failure") @patch("StatusChangedQueueTrigger.get_request_files") @patch("StatusChangedQueueTrigger.handle_status_changed", side_effect=Exception) - def test_get_request_files_should_be_called_when_failing_during_submit_stage(self, _, mock_get_request_files, mock_set_output_event_to_report_failure): + async def test_get_request_files_should_be_called_when_failing_during_submit_stage(self, _, mock_get_request_files, mock_set_output_event_to_report_failure): message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"submitted\" ,\"previous_status\":\"draft\" , \"type\":\"export\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) + mock_get_request_files.return_value = [] + await main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) assert mock_get_request_files.called assert mock_set_output_event_to_report_failure.called @patch("StatusChangedQueueTrigger.blob_operations.get_request_files") @patch.dict(os.environ, {"TRE_ID": "tre-id"}, clear=True) - def test_get_request_files_called_with_correct_storage_account(self, mock_get_request_files): + async def test_get_request_files_called_with_correct_storage_account(self, mock_get_request_files): source_storage_account_for_submitted_stage = constants.STORAGE_ACCOUNT_NAME_EXPORT_INTERNAL + 'ws1' message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"submitted\" ,\"previous_status\":\"draft\" , \"type\":\"export\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - request_properties = extract_properties(message) - get_request_files(request_properties) + request_properties = await extract_properties(message) + mock_get_request_files.return_value = [] + await get_request_files(request_properties) mock_get_request_files.assert_called_with(account_name=source_storage_account_for_submitted_stage, request_id=request_properties.request_id) +@pytest.mark.asyncio class TestFilesDeletion(): @patch("StatusChangedQueueTrigger.set_output_event_to_trigger_container_deletion") @patch.dict(os.environ, {"TRE_ID": "tre-id"}, clear=True) - def test_delete_request_files_should_be_called_on_cancel_stage(self, mock_set_output_event_to_trigger_container_deletion): + async def test_delete_request_files_should_be_called_on_cancel_stage(self, mock_set_output_event_to_trigger_container_deletion): message_body = "{ \"data\": { \"request_id\":\"123\",\"new_status\":\"cancelled\" ,\"previous_status\":\"draft\" , \"type\":\"export\", \"workspace_id\":\"ws1\" }}" message = _mock_service_bus_message(body=message_body) - main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) + await main(msg=message, stepResultEvent=MagicMock(), dataDeletionEvent=MagicMock()) assert mock_set_output_event_to_trigger_container_deletion.called diff --git a/api_app/.env.sample b/api_app/.env.sample index 0cf370e342..6437d86b1e 100644 --- a/api_app/.env.sample +++ b/api_app/.env.sample @@ -44,6 +44,9 @@ SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE=__CHANGE_ME__ SERVICE_BUS_RESOURCE_REQUEST_QUEUE=workspacequeue SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE=deploymentstatus SERVICE_BUS_STEP_RESULT_QUEUE=airlock-step-result +SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME=__CHANGE_ME__ +SERVICE_BUS_MESSAGES_CONTAINER_NAME=sb-messages +SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD=200000 # Event grid configuration # ------------------------- diff --git a/api_app/_version.py b/api_app/_version.py index 6623c5202f..7c4a9591e1 100644 --- a/api_app/_version.py +++ b/api_app/_version.py @@ -1 +1 @@ -__version__ = "0.25.14" +__version__ = "0.26.0" diff --git a/api_app/core/config.py b/api_app/core/config.py index d2f1cf1fa4..b0ce648c58 100644 --- a/api_app/core/config.py +++ b/api_app/core/config.py @@ -41,6 +41,9 @@ SERVICE_BUS_RESOURCE_REQUEST_QUEUE: str = config("SERVICE_BUS_RESOURCE_REQUEST_QUEUE", default="") SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE: str = config("SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE", default="") SERVICE_BUS_STEP_RESULT_QUEUE: str = config("SERVICE_BUS_STEP_RESULT_QUEUE", default="") +SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME: str = config("SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME", default="") +SERVICE_BUS_MESSAGES_CONTAINER_NAME: str = config("SERVICE_BUS_MESSAGES_CONTAINER_NAME", default="sb-messages") +SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD: int = config("SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD", cast=int, default=200000) # 200KB # Event grid configuration EVENT_GRID_STATUS_CHANGED_TOPIC_ENDPOINT: str = config("EVENT_GRID_STATUS_CHANGED_TOPIC_ENDPOINT", default="") diff --git a/api_app/event_grid/helpers.py b/api_app/event_grid/helpers.py index bcad3e65e1..615afb98ab 100644 --- a/api_app/event_grid/helpers.py +++ b/api_app/event_grid/helpers.py @@ -1,9 +1,22 @@ +import json from azure.eventgrid import EventGridEvent from azure.eventgrid.aio import EventGridPublisherClient -from core import credentials +from core import credentials, config +from service_bus.helpers import _offload_to_blob +from services.logging import logger async def publish_event(event: EventGridEvent, topic_endpoint: str): + # Claim check pattern + if config.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME: + data_str = json.dumps(event.data) + if len(data_str) > config.SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD: + logger.info(f"Event data size {len(data_str)} exceeds threshold. Offloading to blob storage.") + blob_url = await _offload_to_blob(data_str) + event.data = { + "claim_check": blob_url + } + async with credentials.get_credential_async_context() as credential: client = EventGridPublisherClient(topic_endpoint, credential) async with client: diff --git a/api_app/service_bus/airlock_request_status_update.py b/api_app/service_bus/airlock_request_status_update.py index a643404a86..529ad4cd43 100644 --- a/api_app/service_bus/airlock_request_status_update.py +++ b/api_app/service_bus/airlock_request_status_update.py @@ -1,3 +1,5 @@ +from service_bus.helpers import receive_message_payload +from service_bus.service_bus_consumer import ServiceBusConsumer import asyncio import json import time @@ -18,10 +20,10 @@ from resources import strings -class AirlockStatusUpdater(): +class AirlockStatusUpdater(ServiceBusConsumer): def __init__(self): - pass + super().__init__("airlock_status_updater") async def init_repos(self): self.airlock_request_repo = await AirlockRequestRepository.create() @@ -36,9 +38,13 @@ async def receive_messages(self): try: current_time = time.time() polling_count += 1 + + # Update heartbeat for supervisor monitoring + self.update_heartbeat() + # Log a heartbeat message every 60 seconds to show the service is still working if current_time - last_heartbeat_time >= 60: - logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue {polling_count} times in the last minute") + logger.info(f"{config.SERVICE_BUS_STEP_RESULT_QUEUE} queue polled {polling_count} times in the last minute") last_heartbeat_time = current_time polling_count = 0 @@ -47,37 +53,48 @@ async def receive_messages(self): receiver = service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_STEP_RESULT_QUEUE) logger.debug(f"Looking for new messages on {config.SERVICE_BUS_STEP_RESULT_QUEUE} queue...") async with receiver: - received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=1) - for msg in received_msgs: - async with AutoLockRenewer() as renewer: - renewer.register(receiver, msg, max_lock_renewal_duration=60) - complete_message = await self.process_message(msg) - if complete_message: - await receiver.complete_message(msg) - else: - # could have been any kind of transient issue, we'll abandon back to the queue, and retry - await receiver.abandon_message(msg) + while True: + # Update heartbeat inside the loop + self.update_heartbeat() + + received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=60) + if not received_msgs: + break + + for msg in received_msgs: + async with AutoLockRenewer() as renewer: + renewer.register(receiver, msg, max_lock_renewal_duration=60) + complete_message = await self.process_message(msg) + if complete_message: + await receiver.complete_message(msg) + else: + # could have been any kind of transient issue, we'll abandon back to the queue, and retry + await receiver.abandon_message(msg) await asyncio.sleep(10) except OperationTimeoutError: # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available logger.debug("No sessions for this process. Will look again...") + await asyncio.sleep(10) - except ServiceBusConnectionError: + except ServiceBusConnectionError as e: # Occasionally there will be a transient / network-level error in connecting to SB. - logger.info("Unknown Service Bus connection error. Will retry...") + logger.warning(f"Service Bus connection error (will retry): {e}") + await asyncio.sleep(10) except Exception as e: # Catch all other exceptions, log them via .exception to get the stack trace, and reconnect - logger.exception(f"Unknown exception. Will retry - {e}") + logger.exception(f"Unexpected error in message processing: {type(e).__name__}: {e}") + await asyncio.sleep(10) async def process_message(self, msg): with tracer.start_as_current_span("process_message") as current_span: complete_message = False try: - message = parse_obj_as(StepResultStatusUpdateMessage, json.loads(str(msg))) + payload = await receive_message_payload(msg) + message = parse_obj_as(StepResultStatusUpdateMessage, json.loads(payload)) current_span.set_attribute("step_id", message.id) current_span.set_attribute("event_type", message.eventType) diff --git a/api_app/service_bus/deployment_status_updater.py b/api_app/service_bus/deployment_status_updater.py index 41670464c7..4327dd6643 100644 --- a/api_app/service_bus/deployment_status_updater.py +++ b/api_app/service_bus/deployment_status_updater.py @@ -1,7 +1,8 @@ -import asyncio import json import uuid import time +import asyncio +from typing import Dict, List, Any from pydantic import ValidationError, parse_obj_as @@ -10,7 +11,7 @@ from db.repositories.resources_history import ResourceHistoryRepository from models.domain.request_action import RequestAction from db.repositories.resource_templates import ResourceTemplateRepository -from service_bus.helpers import send_deployment_message, update_resource_for_step +from service_bus.helpers import send_deployment_message, update_resource_for_step, receive_message_payload from azure.servicebus import NEXT_AVAILABLE_SESSION from azure.servicebus.exceptions import OperationTimeoutError, ServiceBusConnectionError from azure.servicebus.aio import ServiceBusClient, AutoLockRenewer @@ -21,11 +22,12 @@ from models.domain.operation import DeploymentStatusUpdateMessage, Operation, OperationStep, Status from resources import strings from services.logging import logger, tracer +from service_bus.service_bus_consumer import ServiceBusConsumer -class DeploymentStatusUpdater(): +class DeploymentStatusUpdater(ServiceBusConsumer): def __init__(self): - pass + super().__init__("deployment_status_updater") async def init_repos(self): self.operations_repo = await OperationRepository.create() @@ -33,9 +35,6 @@ async def init_repos(self): self.resource_template_repo = await ResourceTemplateRepository.create() self.resource_history_repo = await ResourceHistoryRepository.create() - def run(self, *args, **kwargs): - asyncio.run(self.receive_messages()) - async def receive_messages(self): with tracer.start_as_current_span("deployment_status_receive_messages"): last_heartbeat_time = 0 @@ -45,9 +44,12 @@ async def receive_messages(self): try: current_time = time.time() polling_count += 1 + + # Update heartbeat for supervisor monitoring + self.update_heartbeat() # Log a heartbeat message every 60 seconds to show the service is still working if current_time - last_heartbeat_time >= 60: - logger.info(f"Queue reader heartbeat: Polled {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue {polling_count} times in the last minute") + logger.info(f"{config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue polled {polling_count} times in the last minute") last_heartbeat_time = current_time polling_count = 0 @@ -55,39 +57,53 @@ async def receive_messages(self): service_bus_client = ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) logger.debug(f"Looking for new messages on {config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE} queue...") - # max_wait_time=1 -> don't hold the session open after processing of the message has finished - async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=1, session_id=NEXT_AVAILABLE_SESSION) as receiver: + # max_wait_time=60 -> wait up to 60 seconds for a session to become available + async with service_bus_client.get_queue_receiver(queue_name=config.SERVICE_BUS_DEPLOYMENT_STATUS_UPDATE_QUEUE, max_wait_time=60, session_id=NEXT_AVAILABLE_SESSION) as receiver: logger.info(f"Got a session containing messages: {receiver.session.session_id}") async with AutoLockRenewer() as renewer: renewer.register(receiver, receiver.session, max_lock_renewal_duration=60) - async for msg in receiver: - complete_message = await self.process_message(msg) - if complete_message: - await receiver.complete_message(msg) - else: - # could have been any kind of transient issue, we'll abandon back to the queue, and retry - await receiver.abandon_message(msg) + + while True: + # Update heartbeat inside the session loop too + self.update_heartbeat() + + received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=1) + if not received_msgs: + break + + for msg in received_msgs: + complete_message = await self.process_message(msg) + if complete_message: + await receiver.complete_message(msg) + else: + # could have been any kind of transient issue, we'll abandon back to the queue, and retry + await receiver.abandon_message(msg) + logger.info(f"Closing session: {receiver.session.session_id}") except OperationTimeoutError: # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available logger.debug("No sessions for this process. Will look again...") + await asyncio.sleep(10) - except ServiceBusConnectionError: + except ServiceBusConnectionError as e: # Occasionally there will be a transient / network-level error in connecting to SB. - logger.info("Unknown Service Bus connection error. Will retry...") + logger.warning(f"Service Bus connection error (will retry): {e}") + await asyncio.sleep(10) except Exception as e: # Catch all other exceptions, log them via .exception to get the stack trace, and reconnect - logger.exception(f"Unknown exception. Will retry - {e}") + logger.exception(f"Unexpected error in message processing: {type(e).__name__}: {e}") + await asyncio.sleep(10) - async def process_message(self, msg): + async def process_message(self, msg) -> bool: complete_message = False message = "" with tracer.start_as_current_span("process_message") as current_span: try: - message = parse_obj_as(DeploymentStatusUpdateMessage, json.loads(str(msg))) + payload = await receive_message_payload(msg) + message = parse_obj_as(DeploymentStatusUpdateMessage, json.loads(payload)) current_span.set_attribute("step_id", message.stepId) current_span.set_attribute("operation_id", message.operationId) @@ -98,6 +114,7 @@ async def process_message(self, msg): logger.info(f"Update status in DB for {message.operationId} - {message.status}") except (json.JSONDecodeError, ValidationError): logger.exception(f"{strings.DEPLOYMENT_STATUS_MESSAGE_FORMAT_INCORRECT}: {msg.correlation_id}") + complete_message = True except Exception: logger.exception(f"Exception processing message: {msg.correlation_id}") @@ -115,6 +132,11 @@ async def update_status_in_database(self, message: DeploymentStatusUpdateMessage try: # update the op operation = await self.operations_repo.get_operation_by_id(str(message.operationId)) + + # Add null safety for operation steps + if not operation.steps: + raise ValueError(f"Operation {message.operationId} has no steps") + step_to_update = None is_last_step = False @@ -128,7 +150,7 @@ async def update_status_in_database(self, message: DeploymentStatusUpdateMessage is_last_step = True if step_to_update is None: - raise f"Error finding step {message.stepId} in operation {message.operationId}" + raise ValueError(f"Step {message.stepId} not found in operation {message.operationId}") # update the step status step_to_update.status = message.status @@ -159,7 +181,8 @@ async def update_status_in_database(self, message: DeploymentStatusUpdateMessage # more steps in the op to do? if is_last_step is False: - assert current_step_index < (len(operation.steps) - 1) + if current_step_index >= len(operation.steps) - 1: + raise ValueError(f"Step index {current_step_index} is the last step in operation (has {len(operation.steps)} steps), but more steps were expected") next_step = operation.steps[current_step_index + 1] # catch any errors in updating the resource - maybe Cosmos / schema invalid etc, and report them back to the op @@ -255,7 +278,7 @@ def get_failure_status_for_action(self, action: RequestAction): return status - def create_updated_resource_document(self, resource: dict, message: DeploymentStatusUpdateMessage): + def create_updated_resource_document(self, resource: Dict[str, Any], message: DeploymentStatusUpdateMessage) -> Dict[str, Any]: """ Merge the outputs with the resource document to persist """ @@ -268,7 +291,7 @@ def create_updated_resource_document(self, resource: dict, message: DeploymentSt return resource - def convert_outputs_to_dict(self, outputs_list: [Output]): + def convert_outputs_to_dict(self, outputs_list: List[Output]) -> Dict[str, Any]: """ Convert a list of Porter outputs to a dictionary """ diff --git a/api_app/service_bus/helpers.py b/api_app/service_bus/helpers.py index 56ff47a724..af118381bf 100644 --- a/api_app/service_bus/helpers.py +++ b/api_app/service_bus/helpers.py @@ -1,5 +1,8 @@ +import json +import uuid from azure.servicebus import ServiceBusMessage from azure.servicebus.aio import ServiceBusClient +from azure.storage.blob.aio import BlobServiceClient from pydantic import parse_obj_as from resources import strings from db.repositories.resources_history import ResourceHistoryRepository @@ -25,6 +28,23 @@ async def _send_message(message: ServiceBusMessage, queue: str): :param queue: The Service Bus queue to send the message to. :type queue: str """ + # Claim check pattern + if config.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME: + body_str = b"".join(message.body).decode("utf-8") if isinstance(message.body, list) else str(message.body) + if len(body_str) > config.SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD: + logger.info(f"Message size {len(body_str)} exceeds threshold {config.SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD}. Offloading to blob storage.") + blob_url = await _offload_to_blob(body_str) + offload_message = { + "claim_check": blob_url + } + # We must recreate the ServiceBusMessage because body is read-only + message = ServiceBusMessage( + body=json.dumps(offload_message), + correlation_id=message.correlation_id, + session_id=message.session_id, + application_properties=message.application_properties + ) + async with credentials.get_credential_async_context() as credential: service_bus_client = ServiceBusClient(config.SERVICE_BUS_FULLY_QUALIFIED_NAMESPACE, credential) @@ -35,6 +55,43 @@ async def _send_message(message: ServiceBusMessage, queue: str): await sender.send_messages(message) +async def _offload_to_blob(content: str) -> str: + account_url = f"https://{config.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME}.blob.{config.STORAGE_ENDPOINT_SUFFIX}" + async with credentials.get_credential_async_context() as credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_name = f"msg-{uuid.uuid4()}.json" + blob_client = blob_service_client.get_blob_client(container=config.SERVICE_BUS_MESSAGES_CONTAINER_NAME, blob=blob_name) + await blob_client.upload_blob(content) + return f"{config.SERVICE_BUS_MESSAGES_CONTAINER_NAME}/{blob_name}" + + +async def receive_message_payload(msg: ServiceBusMessage) -> str: + body_str = b"".join(msg.body).decode("utf-8") + try: + body_json = json.loads(body_str) + if "claim_check" in body_json: + blob_path = body_json["claim_check"] + logger.info(f"Message has claim check: {blob_path}. Downloading from blob storage.") + return await _download_from_blob(blob_path) + except json.JSONDecodeError: + pass + + return body_str + + +async def _download_from_blob(blob_path: str) -> str: + account_url = f"https://{config.SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME}.blob.{config.STORAGE_ENDPOINT_SUFFIX}" + container_name, blob_name = blob_path.split("/", 1) + async with credentials.get_credential_async_context() as credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name) + download_stream = await blob_client.download_blob() + content = await download_stream.readall() + return content.decode("utf-8") + + async def send_deployment_message(content, correlation_id, session_id, action): resource_request_message = ServiceBusMessage(body=content, correlation_id=correlation_id, session_id=session_id) logger.info(f"Sending resource request message with correlation ID {resource_request_message.correlation_id}, action: {action}") diff --git a/api_app/service_bus/service_bus_consumer.py b/api_app/service_bus/service_bus_consumer.py new file mode 100644 index 0000000000..34b6163451 --- /dev/null +++ b/api_app/service_bus/service_bus_consumer.py @@ -0,0 +1,95 @@ +import asyncio +import time + +from services.logging import logger + +# Configuration constants for monitoring intervals +HEARTBEAT_CHECK_INTERVAL_SECONDS = 60 +HEARTBEAT_STALENESS_THRESHOLD_SECONDS = 300 +RESTART_DELAY_SECONDS = 5 +MAX_RESTART_DELAY_SECONDS = 300 +SUPERVISOR_ERROR_DELAY_SECONDS = 30 + + +class ServiceBusConsumer: + + def __init__(self, consumer_name: str): + self.service_name = consumer_name.replace('_', ' ').title() + self._last_heartbeat: float = time.monotonic() + self._restart_delay: float = RESTART_DELAY_SECONDS + logger.info(f"Initializing {self.service_name}") + + def update_heartbeat(self): + self._last_heartbeat = time.monotonic() + + def check_heartbeat(self, max_age_seconds: int = HEARTBEAT_STALENESS_THRESHOLD_SECONDS) -> bool: + age = time.monotonic() - self._last_heartbeat + if age > max_age_seconds: + logger.warning(f"{self.service_name} heartbeat is {age:.1f}s old (threshold: {max_age_seconds}s)") + return False + return True + + async def _receive_messages_loop(self): + """Run receive_messages() in a loop with exponential backoff on failure.""" + while True: + try: + start_time = time.monotonic() + logger.info(f"Starting {self.service_name} receive_messages loop...") + await self.receive_messages() + logger.warning(f"{self.service_name} receive_messages() returned unexpectedly") + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception(f"{self.service_name} receive_messages failed: {e}") + + # Reset backoff if the consumer ran long enough to be considered healthy + elapsed = time.monotonic() - start_time + if elapsed > self._restart_delay: + self._restart_delay = RESTART_DELAY_SECONDS + + logger.info(f"{self.service_name} restarting in {self._restart_delay:.0f}s...") + await asyncio.sleep(self._restart_delay) + self._restart_delay = min(self._restart_delay * 2, MAX_RESTART_DELAY_SECONDS) + + async def supervisor_with_heartbeat_check(self): + task = None + try: + while True: + try: + if task is None or task.done(): + if task and task.done(): + try: + await task + except Exception as e: + logger.exception(f"{self.service_name} task failed unexpectedly: {e}") + await asyncio.sleep(RESTART_DELAY_SECONDS) + + logger.info(f"Starting {self.service_name} task...") + task = asyncio.create_task(self._receive_messages_loop()) + self.update_heartbeat() + + await asyncio.sleep(HEARTBEAT_CHECK_INTERVAL_SECONDS) + + if not self.check_heartbeat(): + logger.warning(f"{self.service_name} heartbeat stale, restarting...") + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + task = None + self._restart_delay = RESTART_DELAY_SECONDS + except Exception as e: + logger.exception(f"{self.service_name} supervisor error: {e}") + await asyncio.sleep(SUPERVISOR_ERROR_DELAY_SECONDS) + finally: + if task and not task.done(): + logger.info(f"Cleaning up {self.service_name} task...") + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def receive_messages(self): + raise NotImplementedError("Subclasses must implement receive_messages()") diff --git a/api_app/tests_ma/test_service_bus/test_airlock_request_status_update.py b/api_app/tests_ma/test_service_bus/test_airlock_request_status_update.py index 6404ba122f..f82960299d 100644 --- a/api_app/tests_ma/test_service_bus/test_airlock_request_status_update.py +++ b/api_app/tests_ma/test_service_bus/test_airlock_request_status_update.py @@ -97,11 +97,12 @@ def sample_airlock_request(status=AirlockRequestStatus.Submitted): class ServiceBusReceivedMessageMock: def __init__(self, message: dict): - self.message = json.dumps(message) + self._message = json.dumps(message) + self.body = [self._message.encode("utf-8")] self.correlation_id = "test_correlation_id" def __str__(self): - return self.message + return self._message @patch("event_grid.helpers.EventGridPublisherClient") diff --git a/api_app/tests_ma/test_service_bus/test_deployment_status_update.py b/api_app/tests_ma/test_service_bus/test_deployment_status_update.py index db80c5b1f7..6e4dd69a9d 100644 --- a/api_app/tests_ma/test_service_bus/test_deployment_status_update.py +++ b/api_app/tests_ma/test_service_bus/test_deployment_status_update.py @@ -71,12 +71,13 @@ class ServiceBusReceivedMessageMock: def __init__(self, message: dict): - self.message = json.dumps(message) + self._message = json.dumps(message) + self.body = [self._message.encode("utf-8")] self.correlation_id = "test_correlation_id" self.session_id = "test_session_id" def __str__(self): - return self.message + return self._message def create_sample_workspace_object(workspace_id): @@ -124,8 +125,8 @@ async def test_receiving_bad_json_logs_error(logging_mock, payload): status_updater = DeploymentStatusUpdater() complete_message = await status_updater.process_message(service_bus_received_message_mock) - # bad message data will fail. we don't mark complete=true since we want the message in the DLQ - assert complete_message is False + # bad message data will fail. we mark complete=true since we want the message removed from the queue + assert complete_message is True # check we logged the error error_message = logging_mock.call_args.args[0] diff --git a/config.sample.yaml b/config.sample.yaml index 95fedf2835..a618710590 100644 --- a/config.sample.yaml +++ b/config.sample.yaml @@ -51,6 +51,7 @@ tre: # firewall_force_tunnel_ip: __CHANGE_ME__ firewall_sku: Standard app_gateway_sku: Standard_v2 + #service_bus_sku: Premium deploy_bastion: true # See https://learn.microsoft.com/en-us/azure/bastion/bastion-overview#sku # Set to Basic if wish to connect to VMs in workspaces. diff --git a/config_schema.json b/config_schema.json index abfaf97217..0b855b6487 100644 --- a/config_schema.json +++ b/config_schema.json @@ -133,6 +133,14 @@ "description": "VMSS SKU for the Resource Processor scale set (e.g., Standard_B2s).", "type": "string" }, + "service_bus_sku": { + "description": "SKU of the Service Bus namespace. Must be Standard or Premium as sessions are required.", + "type": "string", + "enum": [ + "Standard", + "Premium" + ] + }, "user_management_enabled": { "description": "When true, TreAdmins can assign/deassign users to workspaces via the UI.", "type": "boolean" diff --git a/core/terraform/airlock/airlock_processor.tf b/core/terraform/airlock/airlock_processor.tf index 981e81bb13..f13df67c1b 100644 --- a/core/terraform/airlock/airlock_processor.tf +++ b/core/terraform/airlock/airlock_processor.tf @@ -95,6 +95,7 @@ resource "azurerm_linux_function_app" "airlock_function_app" { "TRE_ID" = var.tre_id "WEBSITE_CONTENTOVERVNET" = 1 "STORAGE_ENDPOINT_SUFFIX" = module.terraform_azurerm_environment_configuration.storage_suffix + "SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME" = var.service_bus_messages_storage_account_name "AzureWebJobsStorage__clientId" = azurerm_user_assigned_identity.airlock_id.client_id "AzureWebJobsStorage__credential" = "managedidentity" diff --git a/core/terraform/airlock/outputs.tf b/core/terraform/airlock/outputs.tf index 5a71e75030..d227520a8e 100644 --- a/core/terraform/airlock/outputs.tf +++ b/core/terraform/airlock/outputs.tf @@ -21,3 +21,7 @@ output "event_grid_airlock_notification_topic_resource_id" { output "airlock_malware_scan_result_topic_name" { value = local.scan_result_topic_name } + +output "airlock_id_principal_id" { + value = azurerm_user_assigned_identity.airlock_id.principal_id +} diff --git a/core/terraform/airlock/variables.tf b/core/terraform/airlock/variables.tf index 69888118d0..07b729c50b 100644 --- a/core/terraform/airlock/variables.tf +++ b/core/terraform/airlock/variables.tf @@ -107,3 +107,8 @@ variable "encryption_key_versionless_id" { type = string description = "Versionless ID of the encryption key in the key vault" } + +variable "service_bus_messages_storage_account_name" { + type = string + description = "Name of the storage account for Service Bus messages offloading" +} diff --git a/core/terraform/api-webapp.tf b/core/terraform/api-webapp.tf index 47afeb83cb..89e97e2a71 100644 --- a/core/terraform/api-webapp.tf +++ b/core/terraform/api-webapp.tf @@ -62,6 +62,7 @@ resource "azurerm_linux_web_app" "api" { RESOURCE_MANAGER_ENDPOINT = module.terraform_azurerm_environment_configuration.resource_manager_endpoint MICROSOFT_GRAPH_URL = module.terraform_azurerm_environment_configuration.microsoft_graph_endpoint STORAGE_ENDPOINT_SUFFIX = module.terraform_azurerm_environment_configuration.storage_suffix + SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME = azurerm_storage_account.stg.name ENABLE_AIRLOCK_EMAIL_CHECK = var.enable_airlock_email_check LOGGING_LEVEL = var.logging_level OTEL_RESOURCE_ATTRIBUTES = "service.name=api,service.version=${local.version}" diff --git a/core/terraform/firewall/rules.tf b/core/terraform/firewall/rules.tf index cba038487d..b918cf88b3 100644 --- a/core/terraform/firewall/rules.tf +++ b/core/terraform/firewall/rules.tf @@ -25,6 +25,27 @@ resource "azurerm_firewall_policy_rule_collection_group" "core" { ] source_ip_groups = [var.resource_processor_ip_group_id] } + + dynamic "rule" { + for_each = var.service_bus_sku == "Standard" ? [1] : [] + content { + name = "servicebus-standard-egress" + protocols = [ + "TCP" + ] + destination_addresses = [ + "ServiceBus" + ] + destination_ports = [ + "443", + "5671" + ] + source_ip_groups = [ + var.resource_processor_ip_group_id, + var.web_app_ip_group_id + ] + } + } } network_rule_collection { diff --git a/core/terraform/firewall/variables.tf b/core/terraform/firewall/variables.tf index 458bd0bf1e..0053a12933 100644 --- a/core/terraform/firewall/variables.tf +++ b/core/terraform/firewall/variables.tf @@ -64,3 +64,8 @@ variable "shared_services_ip_group_id" { type = string description = "Shared Services IP Group" } + +variable "service_bus_sku" { + type = string + description = "SKU of the Service Bus namespace (Standard, Premium)" +} diff --git a/core/terraform/main.tf b/core/terraform/main.tf index b2f9a6f225..f5019e1412 100644 --- a/core/terraform/main.tf +++ b/core/terraform/main.tf @@ -96,6 +96,8 @@ module "network" { core_address_space = var.core_address_space arm_environment = var.arm_environment firewall_force_tunnel_ip = var.firewall_force_tunnel_ip + service_bus_sku = var.service_bus_sku + } module "firewall" { @@ -114,6 +116,7 @@ module "firewall" { shared_services_ip_group_id = module.network.shared_services_ip_group_id web_app_ip_group_id = module.network.web_app_ip_group_id airlock_processor_ip_group_id = module.network.airlock_processor_ip_group_id + service_bus_sku = var.service_bus_sku } module "appgateway" { @@ -144,29 +147,30 @@ module "appgateway" { } module "airlock_resources" { - source = "./airlock" - tre_id = var.tre_id - location = var.location - resource_group_name = azurerm_resource_group.core.name - airlock_storage_subnet_id = module.network.airlock_storage_subnet_id - airlock_events_subnet_id = module.network.airlock_events_subnet_id - docker_registry_server = local.docker_registry_server - acr_id = data.azurerm_container_registry.acr.id - api_principal_id = azurerm_user_assigned_identity.id.principal_id - airlock_app_service_plan_sku = var.core_app_service_plan_sku - airlock_processor_subnet_id = module.network.airlock_processor_subnet_id - airlock_servicebus = azurerm_servicebus_namespace.sb - airlock_servicebus_fqdn = azurerm_servicebus_namespace.sb.endpoint - applicationinsights_connection_string = module.azure_monitor.app_insights_connection_string - enable_malware_scanning = var.enable_airlock_malware_scanning - arm_environment = var.arm_environment - tre_core_tags = local.tre_core_tags - log_analytics_workspace_id = module.azure_monitor.log_analytics_workspace_id - blob_core_dns_zone_id = module.network.blob_core_dns_zone_id - file_core_dns_zone_id = module.network.file_core_dns_zone_id - queue_core_dns_zone_id = module.network.queue_core_dns_zone_id - table_core_dns_zone_id = module.network.table_core_dns_zone_id - eventgrid_private_dns_zone_id = module.network.eventgrid_private_dns_zone_id + source = "./airlock" + tre_id = var.tre_id + location = var.location + resource_group_name = azurerm_resource_group.core.name + airlock_storage_subnet_id = module.network.airlock_storage_subnet_id + airlock_events_subnet_id = module.network.airlock_events_subnet_id + docker_registry_server = local.docker_registry_server + acr_id = data.azurerm_container_registry.acr.id + api_principal_id = azurerm_user_assigned_identity.id.principal_id + airlock_app_service_plan_sku = var.core_app_service_plan_sku + airlock_processor_subnet_id = module.network.airlock_processor_subnet_id + airlock_servicebus = azurerm_servicebus_namespace.sb + airlock_servicebus_fqdn = azurerm_servicebus_namespace.sb.endpoint + applicationinsights_connection_string = module.azure_monitor.app_insights_connection_string + enable_malware_scanning = var.enable_airlock_malware_scanning + arm_environment = var.arm_environment + tre_core_tags = local.tre_core_tags + log_analytics_workspace_id = module.azure_monitor.log_analytics_workspace_id + blob_core_dns_zone_id = module.network.blob_core_dns_zone_id + file_core_dns_zone_id = module.network.file_core_dns_zone_id + queue_core_dns_zone_id = module.network.queue_core_dns_zone_id + table_core_dns_zone_id = module.network.table_core_dns_zone_id + eventgrid_private_dns_zone_id = module.network.eventgrid_private_dns_zone_id + service_bus_messages_storage_account_name = azurerm_storage_account.stg.name enable_local_debugging = var.enable_local_debugging myip = local.myip @@ -219,6 +223,8 @@ module "resource_processor_vmss_porter" { enable_airlock_malware_scanning = var.enable_airlock_malware_scanning airlock_malware_scan_result_topic_name = module.airlock_resources.airlock_malware_scan_result_topic_name firewall_policy_id = module.firewall.firewall_policy_id + service_bus_messages_storage_account_name = azurerm_storage_account.stg.name + storage_endpoint_suffix = module.terraform_azurerm_environment_configuration.storage_suffix depends_on = [ module.network, diff --git a/core/terraform/modules_move_definitions.tf b/core/terraform/modules_move_definitions.tf index e0ffe47e3e..e2ca4c7f48 100644 --- a/core/terraform/modules_move_definitions.tf +++ b/core/terraform/modules_move_definitions.tf @@ -137,10 +137,6 @@ moved { from = module.servicebus.azurerm_private_endpoint.sbpe to = azurerm_private_endpoint.sbpe } -moved { - from = module.servicebus.azurerm_servicebus_namespace_network_rule_set.servicebus_network_rule_set - to = azurerm_servicebus_namespace_network_rule_set.servicebus_network_rule_set -} # Keyvault moved { diff --git a/core/terraform/network/network.tf b/core/terraform/network/network.tf index d7d3a513f1..73fbd51ec5 100644 --- a/core/terraform/network/network.tf +++ b/core/terraform/network/network.tf @@ -42,6 +42,8 @@ resource "azurerm_virtual_network" "core" { actions = ["Microsoft.Network/virtualNetworks/subnets/action"] } } + + service_endpoints = ["Microsoft.ServiceBus"] } subnet { @@ -50,6 +52,7 @@ resource "azurerm_virtual_network" "core" { private_endpoint_network_policies = "Disabled" security_group = azurerm_network_security_group.default_rules.id route_table_id = azurerm_route_table.rt.id + service_endpoints = var.service_bus_sku == "Standard" ? ["Microsoft.ServiceBus"] : [] } subnet { @@ -58,6 +61,7 @@ resource "azurerm_virtual_network" "core" { private_endpoint_network_policies = "Disabled" security_group = azurerm_network_security_group.default_rules.id route_table_id = azurerm_route_table.rt.id + service_endpoints = ["Microsoft.ServiceBus"] } subnet { @@ -76,7 +80,7 @@ resource "azurerm_virtual_network" "core" { } } - service_endpoints = ["Microsoft.Storage"] + service_endpoints = ["Microsoft.Storage", "Microsoft.ServiceBus"] } subnet { @@ -84,7 +88,6 @@ resource "azurerm_virtual_network" "core" { address_prefixes = [local.airlock_notifications_subnet_address_prefix] private_endpoint_network_policies = "Disabled" security_group = azurerm_network_security_group.default_rules.id - delegation { name = "delegation" diff --git a/core/terraform/network/variables.tf b/core/terraform/network/variables.tf index 965f7bfc15..07a6cb7be3 100644 --- a/core/terraform/network/variables.tf +++ b/core/terraform/network/variables.tf @@ -16,3 +16,6 @@ variable "arm_environment" { variable "firewall_force_tunnel_ip" { type = string } +variable "service_bus_sku" { + type = string +} diff --git a/core/terraform/resource_processor/vmss_porter/cloud-config.yaml b/core/terraform/resource_processor/vmss_porter/cloud-config.yaml index 9c3c331423..abf1dba5c2 100644 --- a/core/terraform/resource_processor/vmss_porter/cloud-config.yaml +++ b/core/terraform/resource_processor/vmss_porter/cloud-config.yaml @@ -57,6 +57,9 @@ write_files: OTEL_RESOURCE_ATTRIBUTES=service.name=resource_processor,service.version=${resource_processor_vmss_porter_image_tag} OTEL_EXPERIMENTAL_RESOURCE_DETECTORS=azure_vm LOGGING_LEVEL=${logging_level} + SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME=${service_bus_messages_storage_account_name} + STORAGE_ENDPOINT_SUFFIX=${storage_endpoint_suffix} + SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD=200000 ${rp_bundle_values} - path: /etc/cron.hourly/docker-prune # An hourly cron job to have docker free disk space. Running this frquently diff --git a/core/terraform/resource_processor/vmss_porter/locals.tf b/core/terraform/resource_processor/vmss_porter/locals.tf index 7c0ec06402..e19f94283c 100644 --- a/core/terraform/resource_processor/vmss_porter/locals.tf +++ b/core/terraform/resource_processor/vmss_porter/locals.tf @@ -47,6 +47,8 @@ locals { aad_authority_url = module.terraform_azurerm_environment_configuration.active_directory_endpoint microsoft_graph_fqdn = regex("(?:(?P[^:/?#]+):)?(?://(?P[^/?#:]*))?", module.terraform_azurerm_environment_configuration.microsoft_graph_endpoint).fqdn logging_level = var.logging_level + service_bus_messages_storage_account_name = var.service_bus_messages_storage_account_name + storage_endpoint_suffix = var.storage_endpoint_suffix rp_bundle_values = local.rp_bundle_values_formatted }) } diff --git a/core/terraform/resource_processor/vmss_porter/outputs.tf b/core/terraform/resource_processor/vmss_porter/outputs.tf new file mode 100644 index 0000000000..d5b21887b7 --- /dev/null +++ b/core/terraform/resource_processor/vmss_porter/outputs.tf @@ -0,0 +1,3 @@ +output "vmss_msi_principal_id" { + value = azurerm_user_assigned_identity.vmss_msi.principal_id +} diff --git a/core/terraform/resource_processor/vmss_porter/variables.tf b/core/terraform/resource_processor/vmss_porter/variables.tf index 342c27b207..6e4cf19181 100644 --- a/core/terraform/resource_processor/vmss_porter/variables.tf +++ b/core/terraform/resource_processor/vmss_porter/variables.tf @@ -124,3 +124,13 @@ variable "firewall_policy_id" { type = string description = "ID of the firewall policy to use for the resource processor" } + +variable "service_bus_messages_storage_account_name" { + type = string + description = "Name of the storage account for Service Bus messages offloading" +} + +variable "storage_endpoint_suffix" { + type = string + description = "Storage endpoint suffix" +} diff --git a/core/terraform/servicebus.tf b/core/terraform/servicebus.tf index 7fbfb99af8..d9c81ea8a7 100644 --- a/core/terraform/servicebus.tf +++ b/core/terraform/servicebus.tf @@ -2,34 +2,12 @@ resource "azurerm_servicebus_namespace" "sb" { name = "sb-${var.tre_id}" location = azurerm_resource_group.core.location resource_group_name = azurerm_resource_group.core.name - sku = "Premium" - premium_messaging_partitions = "1" - capacity = "1" + sku = var.service_bus_sku + premium_messaging_partitions = var.service_bus_sku == "Premium" ? 1 : null + capacity = var.service_bus_sku == "Premium" ? 1 : 0 local_auth_enabled = false tags = local.tre_core_tags - # Block public access - # See https://docs.microsoft.com/azure/service-bus-messaging/service-bus-service-endpoints - network_rule_set { - ip_rules = var.enable_local_debugging ? [local.myip] : null - - # Allows the Eventgrid to access the SB - trusted_services_allowed = true - - # We must enable the Airlock events subnet to access the SB, as the Eventgrid topics can't send messages over PE - # https://docs.microsoft.com/en-us/azure/event-grid/consume-private-endpoints - default_action = "Deny" - public_network_access_enabled = true - network_rules { - subnet_id = module.network.airlock_events_subnet_id - ignore_missing_vnet_service_endpoint = false - } - network_rules { - subnet_id = module.network.airlock_notification_subnet_id - ignore_missing_vnet_service_endpoint = false - } - } - dynamic "customer_managed_key" { for_each = var.enable_cmk_encryption ? [1] : [] content { @@ -47,6 +25,27 @@ resource "azurerm_servicebus_namespace" "sb" { } } + dynamic "network_rule_set" { + for_each = var.service_bus_sku == "Premium" ? [1] : [] + content { + default_action = "Deny" + public_network_access_enabled = true + trusted_services_allowed = true + + ip_rules = var.enable_local_debugging ? [local.myip] : [] + + # Always include Airlock subnets + network_rules { + subnet_id = module.network.airlock_events_subnet_id + ignore_missing_vnet_service_endpoint = false + } + network_rules { + subnet_id = module.network.airlock_notification_subnet_id + ignore_missing_vnet_service_endpoint = false + } + } + } + lifecycle { ignore_changes = [tags] } } @@ -64,7 +63,8 @@ resource "azurerm_servicebus_queue" "service_bus_deployment_status_update_queue" # The returned payload might be large, especially for errors. # Cosmos is the final destination of the messages where 2048 is the limit. - max_message_size_in_kilobytes = 2048 # default=1024 + # Standard SKU supports up to 256 KB. Premium supports up to 100 MB. + max_message_size_in_kilobytes = var.service_bus_sku == "Premium" ? 2048 : null partitioning_enabled = false requires_session = true @@ -88,6 +88,7 @@ resource "azurerm_private_dns_zone_virtual_network_link" "servicebuslink" { } resource "azurerm_private_endpoint" "sbpe" { + count = var.service_bus_sku == "Premium" ? 1 : 0 name = "pe-${azurerm_servicebus_namespace.sb.name}" location = azurerm_resource_group.core.location resource_group_name = azurerm_resource_group.core.name diff --git a/core/terraform/storage.tf b/core/terraform/storage.tf index 1994d84b5c..746eb12368 100644 --- a/core/terraform/storage.tf +++ b/core/terraform/storage.tf @@ -35,6 +35,31 @@ resource "azurerm_storage_account" "stg" { lifecycle { ignore_changes = [infrastructure_encryption_enabled, tags] } } +resource "azurerm_storage_container" "sb_messages" { + name = "sb-messages" + storage_account_id = azurerm_storage_account.stg.id + container_access_type = "private" +} + +resource "azurerm_role_assignment" "api_stg_blob_contributor" { + scope = azurerm_storage_account.stg.id + role_definition_name = "Storage Blob Data Contributor" + principal_id = azurerm_user_assigned_identity.id.principal_id +} + +resource "azurerm_role_assignment" "rp_stg_blob_contributor" { + count = var.resource_processor_type == "vmss_porter" ? 1 : 0 + scope = azurerm_storage_account.stg.id + role_definition_name = "Storage Blob Data Contributor" + principal_id = module.resource_processor_vmss_porter[0].vmss_msi_principal_id +} + +resource "azurerm_role_assignment" "airlock_stg_blob_contributor" { + scope = azurerm_storage_account.stg.id + role_definition_name = "Storage Blob Data Contributor" + principal_id = module.airlock_resources.airlock_id_principal_id +} + resource "azurerm_private_endpoint" "blobpe" { name = "pe-blob-${var.tre_id}" location = azurerm_resource_group.core.location diff --git a/core/terraform/variables.tf b/core/terraform/variables.tf index e813751745..09a09ad48a 100644 --- a/core/terraform/variables.tf +++ b/core/terraform/variables.tf @@ -300,3 +300,13 @@ variable "private_agent_subnet_id" { type = string default = "" } + +variable "service_bus_sku" { + type = string + default = "Premium" + description = "The SKU of the Service Bus namespace. Must be Standard or Premium as sessions are required." + validation { + condition = contains(["Standard", "Premium"], var.service_bus_sku) + error_message = "Invalid service_bus_sku value" + } +} diff --git a/core/version.txt b/core/version.txt index 91819ca382..fd86b3ee91 100644 --- a/core/version.txt +++ b/core/version.txt @@ -1 +1 @@ -__version__ = "0.16.15" +__version__ = "0.17.0" diff --git a/docs/tre-admins/environment-variables.md b/docs/tre-admins/environment-variables.md index ac33efe640..e01ea3f92b 100644 --- a/docs/tre-admins/environment-variables.md +++ b/docs/tre-admins/environment-variables.md @@ -42,6 +42,7 @@ | `WORKSPACE_APP_SERVICE_PLAN_SKU` | Optional. The SKU used for AppService plan used in E2E tests unless otherwise specified. Default value is `P1v2`. | | `RESOURCE_PROCESSOR_NUMBER_PROCESSES_PER_INSTANCE` | Optional. The number of processes to instantiate when the Resource Processor starts. Equates to the number of parallel deployment operations possible in your TRE. Defaults to `5`. | | `FIREWALL_SKU` | Optional. The SKU of the Azure Firewall instance. Default value is `Standard`. Allowed values [`Basic`, `Standard`, `Premium`]. See [Azure Firewall SKU feature comparison](https://learn.microsoft.com/en-us/azure/firewall/choose-firewall-sku). | +| `SERVICE_BUS_SKU` | Optional. The SKU of the Service Bus namespace. Default value is `Premium`. Allowed values [`Standard`, `Premium`]. Standard SKU is recommended for dev/test environments to reduce costs. | | `APP_GATEWAY_SKU` | Optional. The SKU of the Application Gateway. Default value is `Standard_v2`. Allowed values [`Standard_v2`, `WAF_v2`] | | `DEPLOY_BASTION` | Optional. If set to `true`, an Azure Bastion instance will be deployed. Default value is `true`. | | `BASTION_SKU` | Optional. The SKU of the Azure Bastion instance. Default value is `Basic`. Allowed values [`Developer`, `Standard`, `Basic`, `Premium`]. See [Azure Bastion SKU feature comparison](https://learn.microsoft.com/en-us/azure/bastion/bastion-overview#sku). | diff --git a/resource_processor/_version.py b/resource_processor/_version.py index 26c36ca79a..9e78220f94 100644 --- a/resource_processor/_version.py +++ b/resource_processor/_version.py @@ -1 +1 @@ -__version__ = "0.13.3" +__version__ = "0.14.0" diff --git a/resource_processor/shared/config.py b/resource_processor/shared/config.py index e937d24552..6ae0a8f65c 100644 --- a/resource_processor/shared/config.py +++ b/resource_processor/shared/config.py @@ -24,6 +24,10 @@ def get_config() -> dict: config["azure_environment"] = os.environ.get("AZURE_ENVIRONMENT", "AzureCloud") config["aad_authority_url"] = os.environ.get("AAD_AUTHORITY_URL", "https://login.microsoftonline.com") config["microsoft_graph_fqdn"] = os.environ.get("MICROSOFT_GRAPH_FQDN", "graph.microsoft.com") + config["service_bus_messages_storage_account_name"] = os.environ.get("SERVICE_BUS_MESSAGES_STORAGE_ACCOUNT_NAME", "") + config["service_bus_messages_container_name"] = os.environ.get("SERVICE_BUS_MESSAGES_CONTAINER_NAME", "sb-messages") + config["storage_endpoint_suffix"] = os.environ.get("STORAGE_ENDPOINT_SUFFIX", "core.windows.net") + config["service_bus_message_offload_threshold"] = int(os.environ.get("SERVICE_BUS_MESSAGE_OFFLOAD_THRESHOLD", "200000")) try: config["number_processes_int"] = int(config["number_processes"]) diff --git a/resource_processor/shared/sb_helpers.py b/resource_processor/shared/sb_helpers.py new file mode 100644 index 0000000000..8c17b2c36f --- /dev/null +++ b/resource_processor/shared/sb_helpers.py @@ -0,0 +1,86 @@ +import json +import uuid +from azure.servicebus import ServiceBusMessage +from azure.servicebus.aio import ServiceBusClient +from azure.storage.blob.aio import BlobServiceClient +from azure.identity.aio import DefaultAzureCredential +from shared.logging import logger + + +def _get_credential(): + # In VMSS, VMSS_MSI_ID is used. If not set, it will use the system assigned identity or ambient credentials. + from shared.config import get_config + config = get_config() + msi_id = config.get("vmss_msi_id") + return DefaultAzureCredential(managed_identity_client_id=msi_id) if msi_id else DefaultAzureCredential() + + +async def send_message(message: ServiceBusMessage, queue: str, config: dict): + """ + Sends the given message to the given queue in the Service Bus. + """ + # Claim check pattern + if config.get("service_bus_messages_storage_account_name"): + body_str = b"".join(message.body).decode("utf-8") if isinstance(message.body, list) else str(message.body) + if len(body_str) > config.get("service_bus_message_offload_threshold", 200000): + logger.info(f"Message size {len(body_str)} exceeds threshold. Offloading to blob storage.") + blob_url = await _offload_to_blob(body_str, config) + offload_message = { + "claim_check": blob_url + } + # We must recreate the ServiceBusMessage because body is read-only + message = ServiceBusMessage( + body=json.dumps(offload_message), + correlation_id=message.correlation_id, + session_id=message.session_id, + application_properties=message.application_properties + ) + + credential = _get_credential() + async with credential: + service_bus_client = ServiceBusClient(config["service_bus_namespace"], credential) + + async with service_bus_client: + sender = service_bus_client.get_queue_sender(queue_name=queue) + + async with sender: + await sender.send_messages(message) + + +async def _offload_to_blob(content: str, config: dict) -> str: + account_url = f"https://{config['service_bus_messages_storage_account_name']}.blob.{config['storage_endpoint_suffix']}" + credential = _get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_name = f"msg-{uuid.uuid4()}.json" + blob_client = blob_service_client.get_blob_client(container=config["service_bus_messages_container_name"], blob=blob_name) + await blob_client.upload_blob(content) + return f"{config['service_bus_messages_container_name']}/{blob_name}" + + +async def receive_message_payload(msg: ServiceBusMessage, config: dict) -> str: + body_str = b"".join(msg.body).decode("utf-8") + try: + body_json = json.loads(body_str) + if "claim_check" in body_json: + blob_path = body_json["claim_check"] + logger.info(f"Message has claim check: {blob_path}. Downloading from blob storage.") + return await _download_from_blob(blob_path, config) + except json.JSONDecodeError: + pass + + return body_str + + +async def _download_from_blob(blob_path: str, config: dict) -> str: + account_url = f"https://{config['service_bus_messages_storage_account_name']}.blob.{config['storage_endpoint_suffix']}" + container_name, blob_name = blob_path.split("/", 1) + credential = _get_credential() + async with credential: + blob_service_client = BlobServiceClient(account_url, credential=credential) + async with blob_service_client: + blob_client = blob_service_client.get_blob_client(container=container_name, blob=blob_name) + download_stream = await blob_client.download_blob() + content = await download_stream.readall() + return content.decode("utf-8") diff --git a/resource_processor/vmss_porter/requirements.txt b/resource_processor/vmss_porter/requirements.txt index b984362eb0..c735de0582 100644 --- a/resource_processor/vmss_porter/requirements.txt +++ b/resource_processor/vmss_porter/requirements.txt @@ -3,4 +3,5 @@ azure-cli-core==2.68.0 azure-identity==1.25.1 azure-monitor-opentelemetry==1.6.4 azure-servicebus==7.14.3 +azure-storage-blob==12.27.1 opentelemetry-instrumentation-logging==0.49b2 diff --git a/resource_processor/vmss_porter/runner.py b/resource_processor/vmss_porter/runner.py index 120ececba0..043fae90a1 100644 --- a/resource_processor/vmss_porter/runner.py +++ b/resource_processor/vmss_porter/runner.py @@ -7,6 +7,7 @@ from helpers.commands import azure_acr_login_command, azure_login_command, build_porter_command, build_porter_command_for_outputs, apply_porter_credentials_sets_command, run_command_helper from shared.config import get_config from helpers.httpserver import start_server +from resource_processor.shared import sb_helpers from shared.logging import initialize_logging, logger, tracer from shared.config import VERSION @@ -58,53 +59,71 @@ async def receive_message(service_bus_client, config: dict, keep_running=lambda: polling_count = 0 logger.debug("Looking for new session...") - # max_wait_time=1 -> don't hold the session open after processing of the message has finished - async with service_bus_client.get_queue_receiver(queue_name=q_name, max_wait_time=1, session_id=NEXT_AVAILABLE_SESSION) as receiver: + # max_wait_time=60 -> wait up to 60 seconds for a session to become available + async with service_bus_client.get_queue_receiver(queue_name=q_name, max_wait_time=60, session_id=NEXT_AVAILABLE_SESSION) as receiver: logger.info(f"Got a session containing messages: {receiver.session.session_id}") async with AutoLockRenewer() as renewer: # allow a session to be auto lock renewed for up to an hour - if it's processing a message renewer.register(receiver, receiver.session, max_lock_renewal_duration=3600) - async for msg in receiver: - result = True - message = "" - - try: - message = json.loads(str(msg)) - except (json.JSONDecodeError) as e: - logger.error(f"Received bad service bus resource request message: {e}") - - with tracer.start_as_current_span("receive_message") as current_span: - current_span.set_attribute("resource_id", message["id"]) - current_span.set_attribute("action", message["action"]) - current_span.set_attribute("step_id", message["stepId"]) - current_span.set_attribute("operation_id", message["operationId"]) - logger.info(f"Message received for resource_id={message['id']}, operation_id={message['operationId']}, step_id={message['stepId']}") - - result = await invoke_porter_action(message, service_bus_client, config) - - if result: - logger.info(f"Resource request for {message} is complete") - else: - logger.error('Message processing failed!') - - logger.info(f"Message for resource_id={message['id']}, operation_id={message['operationId']} processed as {result} and marked complete.") - await receiver.complete_message(msg) - - logger.info(f"Closing session: {receiver.session.session_id}") + while True: + polling_count += 1 + # Update heartbeat inside the session loop + if time.time() - last_heartbeat_time >= 60: + logger.info(f"Queue reader heartbeat: Polled for messages {polling_count} times in the last minute") + last_heartbeat_time = time.time() + polling_count = 0 + + received_msgs = await receiver.receive_messages(max_message_count=10, max_wait_time=1) + if not received_msgs: + break + + for msg in received_msgs: + result = True + message = "" + + try: + payload = await sb_helpers.receive_message_payload(msg, config) + message = json.loads(payload) + except (json.JSONDecodeError) as e: + logger.error(f"Received bad service bus resource request message: {e}") + await receiver.complete_message(msg) + continue + + with tracer.start_as_current_span("receive_message") as current_span: + current_span.set_attribute("resource_id", message["id"]) + current_span.set_attribute("action", message["action"]) + current_span.set_attribute("step_id", message["stepId"]) + current_span.set_attribute("operation_id", message["operationId"]) + logger.info(f"Message received for resource_id={message['id']}, operation_id={message['operationId']}, step_id={message['stepId']}") + + result = await invoke_porter_action(message, service_bus_client, config) + + if result: + logger.info(f"Resource request for {message} is complete") + else: + logger.error('Message processing failed!') + + logger.info(f"Message for resource_id={message['id']}, operation_id={message['operationId']} processed as {result} and marked complete.") + await receiver.complete_message(msg) + + logger.info(f"Closing session: {receiver.session.session_id}") except OperationTimeoutError: # Timeout occurred whilst connecting to a session - this is expected and indicates no non-empty sessions are available logger.debug("No sessions for this process. Will look again...") + await asyncio.sleep(10) except ServiceBusConnectionError: # Occasionally there will be a transient / network-level error in connecting to SB. logger.info("Unknown Service Bus connection error. Will retry...") + await asyncio.sleep(10) except Exception: # Catch all other exceptions, log them via .exception to get the stack trace, sleep, and reconnect logger.exception("Unknown exception. Will retry...") + await asyncio.sleep(10) async def run_porter(command_parts_list: list, config: dict): @@ -169,11 +188,10 @@ async def invoke_porter_action(msg_body: dict, sb_client: ServiceBusClient, conf installation_id = msg_body["id"] action = msg_body["action"] logger.info(f"{action} action starting for {installation_id}...") - sb_sender = sb_client.get_queue_sender(queue_name=config["deployment_status_queue"]) # post an update message to set the status to an 'in progress' one resource_request_message = service_bus_message_generator(msg_body, statuses.in_progress_status_string_for[action], "Job starting") - await sb_sender.send_messages(ServiceBusMessage(body=resource_request_message, correlation_id=msg_body["id"], session_id=msg_body["operationId"])) + await sb_helpers.send_message(ServiceBusMessage(body=resource_request_message, correlation_id=msg_body["id"], session_id=msg_body["operationId"]), config["deployment_status_queue"], config) logger.info(f'Sent status message for {installation_id} - {statuses.in_progress_status_string_for[action]} - Job starting') # Build and run porter command (flagging if its a built-in action or custom so we can adapt porter command appropriately) @@ -237,7 +255,7 @@ async def invoke_porter_action(msg_body: dict, sb_client: ServiceBusClient, conf resource_request_message = service_bus_message_generator(msg_body, status_for_sb_message, status_message, outputs) - await sb_sender.send_messages(ServiceBusMessage(body=resource_request_message, correlation_id=msg_body["id"], session_id=msg_body["operationId"])) + await sb_helpers.send_message(ServiceBusMessage(body=resource_request_message, correlation_id=msg_body["id"], session_id=msg_body["operationId"]), config["deployment_status_queue"], config) logger.info(f"Sent status message for {installation_id}: {status_for_sb_message}") # return true as want to continue processing the message @@ -276,10 +294,15 @@ async def get_porter_outputs(msg_body: dict, config: dict): async def runner(process_number: int, config: dict): - with tracer.start_as_current_span(process_number): - async with default_credentials(config["vmss_msi_id"]) as credential: - service_bus_client = ServiceBusClient(config["service_bus_namespace"], credential) - await receive_message(service_bus_client, config) + with tracer.start_as_current_span(str(process_number)): + while True: + try: + async with default_credentials(config["vmss_msi_id"]) as credential: + async with ServiceBusClient(config["service_bus_namespace"], credential) as service_bus_client: + await receive_message(service_bus_client, config) + except Exception: + logger.exception("Exception in runner loop") + await asyncio.sleep(10) async def check_runners(processes: list, httpserver: Process, keep_running=lambda: True): @@ -318,7 +341,7 @@ async def check_runners(processes: list, httpserver: Process, keep_running=lambd logger.info(f"Starting {num} processes...") for i in range(num): logger.info(f"Starting process {str(i)}") - process = Process(target=lambda: asyncio.run(runner(i, config))) + process = Process(target=lambda i=i: asyncio.run(runner(i, config))) processes.append(process) process.start()