fix bug in session id being empty string and partition key fails to parse empty string (Azure#21160)

Author: yunhaoling

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -8,6 +8,9 @@
 
 ### Bugs Fixed
 
+- Fixed bug that `ServiceBusReceiver` can not connect to sessionful entity with session id being empty string.
+- Fixed bug that `ServiceBusMessage.partition_key` can not parse empty string properly.
+
 ### Other Changes
 
 ## 7.3.4 (2021-10-06)

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

@@ -293,9 +293,15 @@ def partition_key(self):
         """
         p_key = None
         try:
-            p_key = self._raw_amqp_message.annotations.get(  # type: ignore
-                _X_OPT_PARTITION_KEY  # type: ignore
-            ) or self._raw_amqp_message.annotations.get(ANNOTATION_SYMBOL_PARTITION_KEY)  # type: ignore
+            # opt_p_key is used on the incoming message
+            opt_p_key = self._raw_amqp_message.annotations.get(_X_OPT_PARTITION_KEY)  # type: ignore
+            if opt_p_key is not None:
+                p_key = opt_p_key
+            # symbol_p_key is used on the outgoing message
+            symbol_p_key = self._raw_amqp_message.annotations.get(ANNOTATION_SYMBOL_PARTITION_KEY)  # type: ignore
+            if symbol_p_key is not None:
+                p_key = symbol_p_key
+
             return p_key.decode("UTF-8")  # type: ignore
         except (AttributeError, UnicodeDecodeError):
             return p_key
@@ -310,7 +316,7 @@ def partition_key(self, value):
                 )
             )
 
-        if value and self.session_id and value != self.session_id:
+        if value and self.session_id is not None and value != self.session_id:
             raise ValueError(
                 "partition_key:{} cannot be set to a different value than session_id:{}".format(
                     value, self.session_id

sdk/servicebus/azure-servicebus/azure/servicebus/_servicebus_receiver.py

@@ -156,7 +156,7 @@ def __init__(self, fully_qualified_namespace, credential, **kwargs):
 
         self._populate_attributes(**kwargs)
         self._session = (
-            ServiceBusSession(self._session_id, self) if self._session_id else None
+            None if self._session_id is None else ServiceBusSession(self._session_id, self)
         )
 
     def __iter__(self):

sdk/servicebus/azure-servicebus/azure/servicebus/aio/_servicebus_receiver_async.py

@@ -160,7 +160,7 @@ def __init__(
 
         self._populate_attributes(**kwargs)
         self._session = (
-            ServiceBusSession(self._session_id, self) if self._session_id else None
+            None if self._session_id is None else ServiceBusSession(self._session_id, self)
         )
 
     # Python 3.5 does not allow for yielding from a coroutine, so instead of the try-finally functional wrapper

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

@@ -78,6 +78,27 @@ async def test_async_session_by_session_client_conn_str_receive_handler_peeklock
 
             assert count == 3
 
+            session_id = ""
+            async with sb_client.get_queue_sender(servicebus_queue.name) as sender:
+                for i in range(3):
+                    message = ServiceBusMessage("Handler message no. {}".format(i), session_id=session_id)
+                    await sender.send_messages(message)
+
+            with pytest.raises(ServiceBusError):
+                await sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)._open_with_retry()
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, session_id=session_id, max_wait_time=5)
+            count = 0
+            async for message in receiver:
+                print_message(_logger, message)
+                assert message.session_id == session_id
+                count += 1
+                await receiver.complete_message(message)
+
+            await receiver.close()
+
+            assert count == 3
+
     @pytest.mark.liveTest
     @pytest.mark.live_test_only
     @CachedResourceGroupPreparer(name_prefix='servicebustest')

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

@@ -113,6 +113,63 @@ def test_session_by_session_client_conn_str_receive_handler_peeklock(self, servi
                 assert received_cnt_dic['0'] == 2 and received_cnt_dic['1'] == 2 and received_cnt_dic['2'] == 2
                 assert count == 6
 
+            session_id = ""
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name, session_id=session_id, max_wait_time=5)
+
+            with sender, receiver:
+                for i in range(3):
+                    message = ServiceBusMessage("Handler message no. {}".format(i))
+
+                    message.partition_key = 'pkey'
+
+                    message.session_id = session_id
+                    message.partition_key = session_id
+                    message.application_properties = {'key': 'value'}
+                    message.subject = 'label'
+                    message.content_type = 'application/text'
+                    message.correlation_id = 'cid'
+                    message.message_id = str(i)
+                    message.to = 'to'
+                    message.reply_to = 'reply_to'
+                    message.reply_to_session_id = 'reply_to_session_id'
+
+                    with pytest.raises(ValueError):
+                        message.partition_key = 'pkey'
+
+                    sender.send_messages(message)
+
+                with pytest.raises(ServiceBusError):
+                    receiver = sb_client.get_queue_receiver(servicebus_queue.name, max_wait_time=5)._open_with_retry()
+
+                count = 0
+                received_cnt_dic = {}
+                for message in receiver:
+                    print_message(_logger, message)
+                    assert message.delivery_count == 0
+                    assert message.application_properties
+                    assert message.application_properties[b'key'] == b'value'
+                    assert message.subject == 'label'
+                    assert message.content_type == 'application/text'
+                    assert message.correlation_id == 'cid'
+                    assert message.partition_key == session_id
+                    assert message.to == 'to'
+                    assert message.reply_to == 'reply_to'
+                    assert message.sequence_number
+                    assert message.enqueued_time_utc
+                    assert message.session_id == session_id
+                    assert message.reply_to_session_id == 'reply_to_session_id'
+                    count += 1
+                    receiver.complete_message(message)
+                    if message.message_id not in received_cnt_dic:
+                        received_cnt_dic[message.message_id] = 1
+                        sender.send_messages(message)
+                    else:
+                        received_cnt_dic[message.message_id] += 1
+
+                assert received_cnt_dic['0'] == 2 and received_cnt_dic['1'] == 2 and received_cnt_dic['2'] == 2
+                assert count == 6
+
     @pytest.mark.liveTest
     @pytest.mark.live_test_only
     @CachedResourceGroupPreparer(name_prefix='servicebustest')