[SB] fix to capture properties when using batches (#42598)

* fix to capture properties when using batches

* return none for partition key

* clean up comments

* fix some spacing issues

* enable partitions for queues and add tests

* disable other tests

* fix sleep time and get async cred

* fix test

* reset tests

* reset tests for mgmt

* vendor in to EH

* check partition key

* fix pylint

* update changelog

Author: kashifkhan

sdk/eventhub/azure-eventhub/azure/eventhub/_pyamqp/utils.py

@@ -4,6 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 import datetime
+from typing import TYPE_CHECKING
 from base64 import b64encode
 from hashlib import sha256
 from hmac import HMAC
@@ -12,6 +13,10 @@
 from datetime import timezone
 from .types import TYPE, VALUE, AMQPTypes
 from ._encode import encode_payload
+from .message import Properties
+
+if TYPE_CHECKING:
+    from .message import Message
 
 TZ_UTC: timezone = timezone.utc
 # Number of seconds between the Unix epoch (1/1/1970) and year 1 CE.
@@ -75,6 +80,13 @@ def add_batch(batch, message):
     encode_payload(output, message)
     batch[5].append(output)
 
+def set_message_properties(message, properties: list):
+    if not message[3]:
+        message[3] = Properties(*properties)
+
+def set_message_annotations(message, annotations: dict):
+    if not message[2]:
+        message[2] = annotations
 
 def encode_str(data, encoding="utf-8"):
     try:

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -8,6 +8,8 @@
 
 ### Bugs Fixed
 
+- Fixed a bug where batched messages couldn't be sent to a queue that had session & partitions enabled. ([#42598](https://github.com/Azure/azure-sdk-for-python/pull/42598))
+
 ### Other Changes
 
 ## 7.14.2 (2025-04-09)

sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py

@@ -14,6 +14,7 @@
 
 from .._pyamqp._message_backcompat import LegacyMessage, LegacyBatchMessage
 from .._pyamqp.message import Message as pyamqp_Message
+from .._pyamqp.utils import set_message_properties, set_message_annotations
 from .._transport._pyamqp_transport import PyamqpTransport
 
 from .constants import (
@@ -309,6 +310,7 @@ def partition_key(self) -> Optional[str]:
 
         :rtype: str or None
         """
+        opt_p_key = None
         try:
             opt_p_key = self._raw_amqp_message.annotations.get(_X_OPT_PARTITION_KEY)  # type: ignore
             if opt_p_key is not None:
@@ -675,6 +677,16 @@ def _add(self, add_message: Union[ServiceBusMessage, Mapping[str, Any], AmqpAnno
         self._count += 1
         self._messages.append(outgoing_sb_message)
 
+        if self._count == 1:  # Populate properties on the batch envelope from the first message
+            if outgoing_sb_message.message_id or outgoing_sb_message.session_id:
+                properties: List[Optional[str]] = [None] * 13
+                properties[0] = outgoing_sb_message.message_id
+                properties[10] = outgoing_sb_message.session_id
+                set_message_properties(self._message, properties)
+
+            if outgoing_sb_message.partition_key:
+                set_message_annotations(self._message, {_X_OPT_PARTITION_KEY: outgoing_sb_message.partition_key})
+
     @property
     def message(self) -> Union["BatchMessage", LegacyBatchMessage]:
         """DEPRECATED: Get the underlying uamqp.BatchMessage or LegacyBatchMessage.

sdk/servicebus/azure-servicebus/azure/servicebus/_pyamqp/utils.py

@@ -4,6 +4,7 @@
 # license information.
 # --------------------------------------------------------------------------
 import datetime
+from typing import TYPE_CHECKING
 from base64 import b64encode
 from hashlib import sha256
 from hmac import HMAC
@@ -12,6 +13,10 @@
 from datetime import timezone
 from .types import TYPE, VALUE, AMQPTypes
 from ._encode import encode_payload
+from .message import Properties
+
+if TYPE_CHECKING:
+    from .message import Message
 
 TZ_UTC: timezone = timezone.utc
 # Number of seconds between the Unix epoch (1/1/1970) and year 1 CE.
@@ -74,6 +79,13 @@ def add_batch(batch, message):
     encode_payload(output, message)
     batch[5].append(output)
 
+def set_message_properties(message, properties: list):
+    if not message[3]:
+        message[3] = Properties(*properties)
+
+def set_message_annotations(message, annotations: dict):
+    if not message[2]:
+        message[2] = annotations
 
 def encode_str(data, encoding="utf-8"):
     try:

sdk/servicebus/azure-servicebus/tests/async_tests/test_sessions_async.py

@@ -1533,3 +1533,38 @@ async def test_async_next_available_session_timeout_value(
                 async for msg in receiver:
                     pass
             assert time.time() - start_time2 > 65  # Default service timeout value is 65 seconds
+    
+    @pytest.mark.asyncio
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", requires_session=True, enable_partitioning=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasserAsync()
+    async def test_async_session_partition_batch(self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs):
+
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential(is_async=True)
+        messages = [
+            ServiceBusMessage("Message 1", session_id="mySessionId", message_id=uuid.uuid4(), partition_key="mySessionId"),
+            ServiceBusMessage("Message 2", session_id="mySessionId", message_id=uuid.uuid4(), partition_key="mySessionId")
+        ]
+        async with ServiceBusClient(
+            fully_qualified_namespace=fully_qualified_namespace,
+            credential=credential,
+            logging_enable=False,
+            uamqp_transport=uamqp_transport,
+        ) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                await sender.send_messages(messages)
+
+            received_messages = []
+            async with sb_client.get_queue_receiver(servicebus_queue.name, session_id="mySessionId", max_wait_time=10) as receiver:
+                async for message in receiver:
+                    received_messages.append(message)
+            assert len(received_messages) == 2
+            assert all(msg.session_id == "mySessionId" for msg in received_messages)
+            assert all(msg.partition_key == "mySessionId" for msg in received_messages)

sdk/servicebus/azure-servicebus/tests/servicebus_preparer.py

@@ -427,6 +427,7 @@ def __init__(
         requires_duplicate_detection=False,
         dead_lettering_on_message_expiration=False,
         requires_session=False,
+        enable_partitioning=False,
         lock_duration="PT30S",
         parameter_name=SERVICEBUS_QUEUE_PARAM,
         resource_group_parameter_name=RESOURCE_GROUP_PARAM,
@@ -452,12 +453,14 @@ def __init__(
             dead_lettering_on_message_expiration,
             requires_session,
             lock_duration,
+            enable_partitioning,
         )
 
         # Queue parameters
         self.requires_duplicate_detection = requires_duplicate_detection
         self.dead_lettering_on_message_expiration = dead_lettering_on_message_expiration
         self.requires_session = requires_session
+        self.enable_partitioning = enable_partitioning
         self.lock_duration = lock_duration
         if random_name_enabled:
             self.resource_moniker = self.name_prefix + "sbqueue"
@@ -481,6 +484,7 @@ def create_resource(self, name, **kwargs):
                             requires_duplicate_detection=self.requires_duplicate_detection,
                             dead_lettering_on_message_expiration=self.dead_lettering_on_message_expiration,
                             requires_session=self.requires_session,
+                            enable_partitioning=self.enable_partitioning,
                         ),
                     )
                     break

sdk/servicebus/azure-servicebus/tests/test_sessions.py

@@ -1676,3 +1676,37 @@ def test_next_available_session_timeout_value(
                     pass
 
             assert time.time() - start_time2 > 65  # Default service operation timeout is 65 seconds
+    
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedServiceBusResourceGroupPreparer(name_prefix="servicebustest")
+    @CachedServiceBusNamespacePreparer(name_prefix="servicebustest")
+    @ServiceBusQueuePreparer(name_prefix="servicebustest", requires_session=True, enable_partitioning=True)
+    @pytest.mark.parametrize("uamqp_transport", uamqp_transport_params, ids=uamqp_transport_ids)
+    @ArgPasser()
+    def test_session_partition_batch(self, uamqp_transport, *, servicebus_namespace=None, servicebus_queue=None, **kwargs):
+
+        fully_qualified_namespace = f"{servicebus_namespace.name}{SERVICEBUS_ENDPOINT_SUFFIX}"
+        credential = get_credential()
+        messages = [
+            ServiceBusMessage("Message 1", session_id="mySessionId", message_id=uuid.uuid4(), partition_key="mySessionId"),
+            ServiceBusMessage("Message 2", session_id="mySessionId", message_id=uuid.uuid4(), partition_key="mySessionId")
+        ]
+        with ServiceBusClient(
+            fully_qualified_namespace=fully_qualified_namespace,
+            credential=credential,
+            logging_enable=False,
+            uamqp_transport=uamqp_transport,
+        ) as sb_client:
+
+            with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                sender.send_messages(messages)
+            
+            received_messages = []
+            with sb_client.get_queue_receiver(servicebus_queue.name, session_id="mySessionId", max_wait_time=10) as receiver:
+                for message in receiver:
+                    received_messages.append(message)
+                
+            assert len(received_messages) == 2
+            assert all(msg.session_id == "mySessionId" for msg in received_messages)
+            assert all(msg.partition_key == "mySessionId" for msg in received_messages)