PR feedback

Author: andystaples

durabletask/extensions/azure_blob_payloads/blob_payload_store.py

@@ -10,6 +10,7 @@
 import uuid
 from typing import Optional
 
+from azure.core.exceptions import ResourceExistsError
 from azure.storage.blob import BlobServiceClient
 from azure.storage.blob.aio import BlobServiceClient as AsyncBlobServiceClient
 
@@ -82,6 +83,30 @@ def __init__(self, options: BlobPayloadStoreOptions):
 
         self._ensure_container_created = False
 
+    # ------------------------------------------------------------------
+    # Lifecycle / resource management
+    # ------------------------------------------------------------------
+
+    def close(self) -> None:
+        """Close the underlying sync blob service client."""
+        self._blob_service_client.close()
+
+    async def close_async(self) -> None:
+        """Close the underlying async blob service client."""
+        await self._async_blob_service_client.close()
+
+    def __enter__(self) -> BlobPayloadStore:
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.close()
+
+    async def __aenter__(self) -> BlobPayloadStore:
+        return self
+
+    async def __aexit__(self, *args: object) -> None:
+        await self.close_async()
+
     @property
     def options(self) -> BlobPayloadStoreOptions:
         return self._options
@@ -185,8 +210,7 @@ def _ensure_container_sync(self) -> None:
         container_client = self._blob_service_client.get_container_client(self._container_name)
         try:
             container_client.create_container()
-        except Exception:
-            # Container may already exist — that is fine.
+        except ResourceExistsError:
             pass
         self._ensure_container_created = True
 
@@ -196,7 +220,6 @@ async def _ensure_container_async(self) -> None:
         container_client = self._async_blob_service_client.get_container_client(self._container_name)
         try:
             await container_client.create_container()
-        except Exception:
-            # Container may already exist — that is fine.
+        except ResourceExistsError:
             pass
         self._ensure_container_created = True

durabletask/payload/helpers.py

@@ -84,6 +84,12 @@ async def deexternalize_payloads_async(
 # Internal recursive walkers – sync
 # ------------------------------------------------------------------
 
+def _is_map_field(fd) -> bool:
+    """Return True if the field descriptor represents a protobuf map field."""
+    mt = fd.message_type
+    return mt is not None and fd.is_repeated and mt.GetOptions().map_entry
+
+
 def _walk_and_externalize(
     msg: proto_message.Message,
     store: PayloadStore,
@@ -95,7 +101,21 @@ def _walk_and_externalize(
         if fd.message_type is None:
             continue
 
-        if fd.is_repeated:
+        if _is_map_field(fd):
+            # Map fields: iterate values. ScalarMap values are not
+            # messages and will be skipped by the isinstance check.
+            for map_value in getattr(msg, fd.name).values():
+                if isinstance(map_value, proto_message.Message):
+                    if isinstance(map_value, wrappers_pb2.StringValue):
+                        _try_externalize_field(
+                            fd.name, map_value, store,
+                            threshold, max_bytes, instance_id,
+                        )
+                    else:
+                        _walk_and_externalize(
+                            map_value, store, threshold, max_bytes, instance_id
+                        )
+        elif fd.is_repeated:
             value = getattr(msg, fd.name)
             for item in value:
                 if isinstance(item, proto_message.Message):
@@ -162,7 +182,14 @@ def _walk_and_deexternalize(
         if fd.message_type is None:
             continue
 
-        if fd.is_repeated:
+        if _is_map_field(fd):
+            for map_value in getattr(msg, fd.name).values():
+                if isinstance(map_value, proto_message.Message):
+                    if isinstance(map_value, wrappers_pb2.StringValue):
+                        _try_deexternalize_field(map_value, store)
+                    else:
+                        _walk_and_deexternalize(map_value, store)
+        elif fd.is_repeated:
             value = getattr(msg, fd.name)
             for item in value:
                 if isinstance(item, proto_message.Message):
@@ -207,7 +234,19 @@ async def _walk_and_externalize_async(
         if fd.message_type is None:
             continue
 
-        if fd.is_repeated:
+        if _is_map_field(fd):
+            for map_value in getattr(msg, fd.name).values():
+                if isinstance(map_value, proto_message.Message):
+                    if isinstance(map_value, wrappers_pb2.StringValue):
+                        await _try_externalize_field_async(
+                            fd.name, map_value, store,
+                            threshold, max_bytes, instance_id,
+                        )
+                    else:
+                        await _walk_and_externalize_async(
+                            map_value, store, threshold, max_bytes, instance_id,
+                        )
+        elif fd.is_repeated:
             value = getattr(msg, fd.name)
             for item in value:
                 if isinstance(item, proto_message.Message):
@@ -273,7 +312,14 @@ async def _walk_and_deexternalize_async(
         if fd.message_type is None:
             continue
 
-        if fd.is_repeated:
+        if _is_map_field(fd):
+            for map_value in getattr(msg, fd.name).values():
+                if isinstance(map_value, proto_message.Message):
+                    if isinstance(map_value, wrappers_pb2.StringValue):
+                        await _try_deexternalize_field_async(map_value, store)
+                    else:
+                        await _walk_and_deexternalize_async(map_value, store)
+        elif fd.is_repeated:
             value = getattr(msg, fd.name)
             for item in value:
                 if isinstance(item, proto_message.Message):