Service bus std

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),

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)

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))