Add messsage state to SB received message (Azure#22837)

* initial commit

* working

* lint

* tests

* comments

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

Co-authored-by: Adam Ling (MSFT) <adam_ling@outlook.com>

* version

* typo

* int enum

* oops

Co-authored-by: Adam Ling (MSFT) <adam_ling@outlook.com>

Author: rakshith91

sdk/servicebus/azure-servicebus/CHANGELOG.md

@@ -1,9 +1,12 @@
 # Release History
 
-## 7.5.1 (Unreleased)
+## 7.6.0 (Unreleased)
 
 ### Features Added
 
+- Introduce `ServiceBusMessageState` enum that can assume the values of `active`, `scheduled` or `deferred`.
+- Add `ServiceBusMessageState` property in `ServiceBusReceivedMessage`.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/servicebus/azure-servicebus/azure/servicebus/__init__.py

@@ -21,6 +21,7 @@
 from ._common.constants import (
     ServiceBusReceiveMode,
     ServiceBusSubQueue,
+    ServiceBusMessageState,
     ServiceBusSessionFilter,
     NEXT_AVAILABLE_SESSION,
 )
@@ -35,6 +36,7 @@
 __all__ = [
     "ServiceBusMessage",
     "ServiceBusMessageBatch",
+    "ServiceBusMessageState",
     "ServiceBusReceivedMessage",
     "NEXT_AVAILABLE_SESSION",
     "ServiceBusSubQueue",

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

@@ -165,12 +165,16 @@
 MAX_DURATION_VALUE = 922337203685477
 # equivalent to .NET Date("9999-12-31T07:59:59.000Z").getTime() in ms
 MAX_ABSOLUTE_EXPIRY_TIME = 253402243199000
-
+MESSAGE_STATE_NAME = b"x-opt-message-state"
 
 class ServiceBusReceiveMode(str, Enum):
     PEEK_LOCK = "peeklock"
     RECEIVE_AND_DELETE = "receiveanddelete"
 
+class ServiceBusMessageState(int, Enum):
+    ACTIVE = 0
+    DEFERRED = 1
+    SCHEDULED = 2
 
 # To enable extensible string enums for the public facing parameter, and translate to the "real" uamqp constants.
 ServiceBusToAMQPReceiveModeMap = {

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

@@ -19,6 +19,7 @@
 from .constants import (
     _BATCH_MESSAGE_OVERHEAD_COST,
     ServiceBusReceiveMode,
+    ServiceBusMessageState,
     _X_OPT_ENQUEUED_TIME,
     _X_OPT_SEQUENCE_NUMBER,
     _X_OPT_ENQUEUE_SEQUENCE_NUMBER,
@@ -35,6 +36,7 @@
     MESSAGE_PROPERTY_MAX_LENGTH,
     MAX_ABSOLUTE_EXPIRY_TIME,
     MAX_DURATION_VALUE,
+    MESSAGE_STATE_NAME
 )
 from ..amqp import (
     AmqpAnnotatedMessage,
@@ -967,6 +969,23 @@ def dead_letter_source(self):
                 pass
         return None
 
+    @property
+    def message_state(self):
+        # type: () -> ServiceBusMessageState
+        """
+        Defaults to Active. Represents the message state of the message. Can be Active, Deferred.
+        or Scheduled.
+
+        :rtype: ~azure.servicebus.ServiceBusMessageState
+        """
+        try:
+            message_state = self._raw_amqp_message.annotations.get(MESSAGE_STATE_NAME)
+            if not message_state:
+                return ServiceBusMessageState.ACTIVE
+            return ServiceBusMessageState(message_state)
+        except AttributeError:
+            return ServiceBusMessageState.ACTIVE
+
     @property
     def delivery_count(self):
         # type: () -> Optional[int]

sdk/servicebus/azure-servicebus/azure/servicebus/_version.py

@@ -3,4 +3,4 @@
 # Licensed under the MIT License.
 # ------------------------------------
 
-VERSION = "7.5.1"
+VERSION = "7.6.0"

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

@@ -28,6 +28,7 @@
     TransportType,
     ServiceBusReceiveMode,
     ServiceBusSubQueue,
+    ServiceBusMessageState
 )
 from azure.servicebus.amqp import (
     AmqpMessageHeader,
@@ -2123,3 +2124,57 @@ async def test_queue_async_send_amqp_annotated_message(self, servicebus_namespac
                     assert recv_sequence_msg == 3
                     assert recv_value_msg == 3
                     assert normal_msg == 4
+
+
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='servicebustest')
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @ServiceBusQueuePreparer(name_prefix='servicebustest', dead_lettering_on_message_expiration=True)
+    async def test_message_state_scheduled_async(self, servicebus_namespace_connection_string, servicebus_queue, **kwargs):
+        async with ServiceBusClient.from_connection_string(
+            servicebus_namespace_connection_string) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                for i in range(10):
+                    message = ServiceBusMessage("message no. {}".format(i))
+                    scheduled_time_utc = datetime.utcnow() + timedelta(seconds=30)
+                    sequence_number = await sender.schedule_messages(message, scheduled_time_utc)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name)
+            async with receiver:
+                messages = await receiver.peek_messages()
+                for msg in messages:
+                    assert msg.message_state == ServiceBusMessageState.SCHEDULED
+
+
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='servicebustest')
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @ServiceBusQueuePreparer(name_prefix='servicebustest', dead_lettering_on_message_expiration=True)
+    async def test_message_state_deferred_async(self, servicebus_namespace_connection_string, servicebus_queue, **kwargs):
+        async with ServiceBusClient.from_connection_string(
+            servicebus_namespace_connection_string) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            async with sender:
+                for i in range(10):
+                    message = ServiceBusMessage("message no. {}".format(i))
+                    await sender.send_messages(message)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name)
+            deferred_messages = []
+            async with receiver:
+                received_msgs = await receiver.receive_messages()
+                for message in received_msgs:
+                    assert message.message_state == ServiceBusMessageState.ACTIVE
+                    deferred_messages.append(message.sequence_number)
+                    await receiver.defer_message(message)
+                if deferred_messages:
+                    received_deferred_msg = await receiver.receive_deferred_messages(
+                        sequence_numbers=deferred_messages
+                        )                
+                for message in received_deferred_msg:
+                    assert message.message_state == ServiceBusMessageState.DEFERRED

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

@@ -28,6 +28,7 @@
     ServiceBusReceivedMessage,
     ServiceBusReceiveMode,
     ServiceBusSubQueue,
+    ServiceBusMessageState
 )
 from azure.servicebus.amqp import (
     AmqpMessageHeader,
@@ -39,7 +40,8 @@
     _X_OPT_LOCK_TOKEN,
     _X_OPT_PARTITION_KEY,
     _X_OPT_VIA_PARTITION_KEY,
-    _X_OPT_SCHEDULED_ENQUEUE_TIME
+    _X_OPT_SCHEDULED_ENQUEUE_TIME,
+    ServiceBusMessageState
 )
 from azure.servicebus._common.utils import utc_now
 from azure.servicebus.management._models import DictMixin
@@ -2581,3 +2583,54 @@ def test_queue_send_amqp_annotated_message(self, servicebus_namespace_connection
                     assert recv_data_msg == 3
                     assert recv_value_msg == 3
                     assert normal_msg == 4
+
+
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='servicebustest')
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @ServiceBusQueuePreparer(name_prefix='servicebustest', dead_lettering_on_message_expiration=True)
+    def test_message_state_scheduled(self, servicebus_namespace_connection_string, servicebus_queue, **kwargs):
+        with ServiceBusClient.from_connection_string(
+            servicebus_namespace_connection_string) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            for i in range(10):
+                message = ServiceBusMessage("message no. {}".format(i))
+                scheduled_time_utc = datetime.utcnow() + timedelta(seconds=30)
+                sequence_number = sender.schedule_messages(message, scheduled_time_utc)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name)
+            with receiver:
+                for msg in receiver.peek_messages():
+                    assert msg.message_state == ServiceBusMessageState.SCHEDULED
+
+
+    @pytest.mark.liveTest
+    @pytest.mark.live_test_only
+    @CachedResourceGroupPreparer(name_prefix='servicebustest')
+    @CachedServiceBusNamespacePreparer(name_prefix='servicebustest')
+    @ServiceBusQueuePreparer(name_prefix='servicebustest', dead_lettering_on_message_expiration=True)
+    def test_message_state_deferred(self, servicebus_namespace_connection_string, servicebus_queue, **kwargs):
+        with ServiceBusClient.from_connection_string(
+            servicebus_namespace_connection_string) as sb_client:
+
+            sender = sb_client.get_queue_sender(servicebus_queue.name)
+            for i in range(10):
+                message = ServiceBusMessage("message no. {}".format(i))
+                sender.send_messages(message)
+
+            receiver = sb_client.get_queue_receiver(servicebus_queue.name)
+            deferred_messages = []
+            with receiver:
+                received_msgs = receiver.receive_messages()
+                for message in received_msgs:
+                    assert message.message_state == ServiceBusMessageState.ACTIVE
+                    deferred_messages.append(message.sequence_number)
+                    receiver.defer_message(message)
+                if deferred_messages:
+                    received_deferred_msg = receiver.receive_deferred_messages(
+                        sequence_numbers=deferred_messages
+                        )
+                for message in received_deferred_msg:
+                    assert message.message_state == ServiceBusMessageState.DEFERRED