[EH] Partition Last Received Event Updating Fix (Azure#31352)

* partition last event received updated

* add test

* sync test

* delete container at the end

* add dev req

* start trying to move to consumer_async

* start aligning uamqp/pyamqp

* update conftest

* update changes

* attempt at one thread

* revert to 2 thread

* update test

* update test 2

* remove commented code

* remove dev_req

* remove comment

* update test

* update to kind

* remove import

* remove import

* don't pin websocket

Author: l0lawrence

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_consumer_async.py

@@ -177,12 +177,12 @@ def _create_handler(
     async def _open_with_retry(self) -> None:
         await self._do_retryable_operation(self._open, operation_need_param=False)
 
-    # only used by _uamqp_transport_async
     def _next_message_in_buffer(self):
         # pylint:disable=protected-access
         message = self._message_buffer.popleft()
         event_data = EventData._from_message(message)
-        event_data._uamqp_message = message
+        if self._amqp_transport.KIND == "uamqp":
+            event_data._uamqp_message = message
         self._last_received_event = event_data
         return event_data
 

sdk/eventhub/azure-eventhub/azure/eventhub/aio/_transport/_pyamqp_transport_async.py

@@ -156,11 +156,10 @@ def create_receive_client(*, config, **kwargs):  # pylint:disable=unused-argumen
     async def _callback_task(consumer, batch, max_batch_size, max_wait_time):
         while consumer._callback_task_run: # pylint: disable=protected-access
             async with consumer._message_buffer_lock: # pylint: disable=protected-access
-                messages = [
-                    consumer._message_buffer.popleft() # pylint: disable=protected-access
+                events = [
+                    consumer._next_message_in_buffer() # pylint: disable=protected-access
                     for _ in range(min(max_batch_size, len(consumer._message_buffer))) # pylint: disable=protected-access
                 ]
-            events = [EventData._from_message(message) for message in messages] # pylint: disable=protected-access
             now_time = time.time()
             if len(events) > 0:
                 await consumer._on_event_received(events if batch else events[0]) # pylint: disable=protected-access
@@ -228,6 +227,7 @@ async def receive_messages_async(consumer, batch, max_batch_size, max_wait_time)
         # pylint:disable=protected-access
         consumer._callback_task_run = True
         consumer._last_callback_called_time = time.time()
+
         callback_task = asyncio.create_task(
             PyamqpTransportAsync._callback_task(consumer, batch, max_batch_size, max_wait_time)
         )

sdk/eventhub/azure-eventhub/dev_requirements.txt

@@ -4,5 +4,7 @@
 azure-mgmt-eventhub<=10.1.0
 azure-mgmt-resource==20.0.0
 aiohttp>=3.0,<4.0
-websocket-client==1.4.2
+websocket-client
 -e ../../../tools/azure-devtools
+azure-eventhub-checkpointstoreblob-aio
+azure-eventhub-checkpointstoreblob

sdk/eventhub/azure-eventhub/tests/conftest.py

@@ -18,6 +18,8 @@
 from azure.eventhub import EventHubProducerClient
 from azure.eventhub._pyamqp import ReceiveClient
 from azure.eventhub._pyamqp.authentication import SASTokenAuth
+from azure.eventhub.extensions.checkpointstoreblob import BlobCheckpointStore
+from azure.eventhub.extensions.checkpointstoreblobaio import BlobCheckpointStore as BlobCheckpointStoreAsync
 try:
     import uamqp
     uamqp_transport_params = [True, False]
@@ -54,6 +56,24 @@ def sleep(request):
 def uamqp_transport(request):
     return request.param
 
+@pytest.fixture(scope="session")    
+def storage_connection_str():
+    try:
+        return os.environ['AZURE_STORAGE_CONN_STR']
+    except KeyError:
+        pytest.skip('AZURE_STORAGE_CONN_STR undefined')
+        return
+
+@pytest.fixture()    
+def checkpoint_store(storage_connection_str):
+    checkpoint_store = BlobCheckpointStore.from_connection_string(storage_connection_str, "blobcontainer" + str(uuid.uuid4()))
+    return checkpoint_store
+
+@pytest.fixture()    
+def checkpoint_store_aio(storage_connection_str):
+    checkpoint_store = BlobCheckpointStoreAsync.from_connection_string(storage_connection_str, "blobcontainer" + str(uuid.uuid4()))
+    return checkpoint_store
+
 def get_logger(filename, level=logging.INFO):
     azure_logger = logging.getLogger("azure.eventhub")
     azure_logger.setLevel(level)
@@ -179,6 +199,17 @@ def live_eventhub(resource_group, eventhub_namespace):  # pylint: disable=redefi
         except:
             warnings.warn(UserWarning("eventhub teardown failed"))
 
+@pytest.fixture()
+def resource_mgmt_client():
+    try:
+        SUBSCRIPTION_ID = os.environ["AZURE_SUBSCRIPTION_ID"]
+    except KeyError:
+        pytest.skip('AZURE_SUBSCRIPTION_ID defined')
+        return
+    base_url = os.environ.get("EVENTHUB_RESOURCE_MANAGER_URL", "https://management.azure.com/")
+    credential_scopes = ["{}.default".format(base_url)]
+    resource_client = EventHubManagementClient(EnvironmentCredential(), SUBSCRIPTION_ID, base_url=base_url, credential_scopes=credential_scopes)
+    yield resource_client
 
 @pytest.fixture()
 def connection_str(live_eventhub):

sdk/eventhub/azure-eventhub/tests/livetest/asynctests/test_consumer_client_async.py

@@ -9,6 +9,68 @@
 from azure.eventhub._constants import ALL_PARTITIONS
 
 
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_receive_storage_checkpoint_async(connstr_senders, uamqp_transport, checkpoint_store_aio, live_eventhub, resource_mgmt_client):
+    connection_str, senders = connstr_senders
+
+    for i in range(10):
+        senders[0].send(EventData("Test EventData"))
+        senders[1].send(EventData("Test EventData"))
+
+    try:
+        await checkpoint_store_aio._container_client.create_container()
+    except:
+        pass
+
+    client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', checkpoint_store=checkpoint_store_aio, uamqp_transport=uamqp_transport)
+
+    sequence_numbers_0 = []
+    sequence_numbers_1 = []
+    async def on_event(partition_context, event):
+        await partition_context.update_checkpoint(event)
+        sequence_num = event.sequence_number
+        if partition_context.partition_id == "0":
+            if sequence_num in sequence_numbers_0:
+                assert False
+            sequence_numbers_0.append(sequence_num)
+        else:
+            if sequence_num in sequence_numbers_1:
+                assert False
+            sequence_numbers_1.append(sequence_num)
+
+    async with client:
+        task = asyncio.ensure_future(
+            client.receive(on_event, starting_position="-1"))
+        # Update the eventhub
+        eventhub = resource_mgmt_client.event_hubs.get(
+            live_eventhub["resource_group"],
+            live_eventhub["namespace"],
+            live_eventhub["event_hub"]
+        )
+        properties = eventhub.as_dict()
+        if properties["message_retention_in_days"] == 1:
+            properties["message_retention_in_days"] = 2
+        else:
+            properties["message_retention_in_days"] = 1
+        resource_mgmt_client.event_hubs.create_or_update(
+            live_eventhub["resource_group"],
+            live_eventhub["namespace"],
+            live_eventhub["event_hub"],
+            properties
+        )
+        await asyncio.sleep(10)
+
+ 
+    await task
+    assert len(sequence_numbers_0) == 10
+    assert len(sequence_numbers_1) == 10
+
+    try:
+        await checkpoint_store_aio._container_client.delete_container()
+    except:
+        pass
+
 @pytest.mark.liveTest
 @pytest.mark.asyncio
 async def test_receive_no_partition_async(connstr_senders, uamqp_transport):

sdk/eventhub/azure-eventhub/tests/livetest/synctests/test_consumer_client.py

@@ -10,6 +10,70 @@
 from azure.eventhub._eventprocessor.in_memory_checkpoint_store import InMemoryCheckpointStore
 from azure.eventhub._constants import ALL_PARTITIONS
 
+@pytest.mark.liveTest
+@pytest.mark.asyncio
+async def test_receive_storage_checkpoint(connstr_senders, uamqp_transport, checkpoint_store, live_eventhub, resource_mgmt_client):
+    connection_str, senders = connstr_senders
+
+    for i in range(10):
+        senders[0].send(EventData("Test EventData"))
+        senders[1].send(EventData("Test EventData"))
+
+    try:
+        checkpoint_store._container_client.create_container()
+    except:
+        pass
+
+    client = EventHubConsumerClient.from_connection_string(connection_str, consumer_group='$default', checkpoint_store=checkpoint_store, uamqp_transport=uamqp_transport)
+
+    sequence_numbers_0 = []
+    sequence_numbers_1 = []
+    def on_event(partition_context, event):
+        partition_context.update_checkpoint(event)
+        sequence_num = event.sequence_number
+        if partition_context.partition_id == "0":
+            if sequence_num in sequence_numbers_0:
+                assert False
+            sequence_numbers_0.append(sequence_num)
+        else:
+            if sequence_num in sequence_numbers_1:
+                assert False
+            sequence_numbers_1.append(sequence_num)
+
+    with client:
+        worker = threading.Thread(target=client.receive,
+                                  args=(on_event,),
+                                  kwargs={"starting_position": "-1"})
+        worker.start()
+
+        # Update the eventhub
+        eventhub = resource_mgmt_client.event_hubs.get(
+            live_eventhub["resource_group"],
+            live_eventhub["namespace"],
+            live_eventhub["event_hub"]
+        )
+        properties = eventhub.as_dict()
+        if properties["message_retention_in_days"] == 1:
+            properties["message_retention_in_days"] = 2
+        else:
+            properties["message_retention_in_days"] = 1
+        resource_mgmt_client.event_hubs.create_or_update(
+            live_eventhub["resource_group"],
+            live_eventhub["namespace"],
+            live_eventhub["event_hub"],
+            properties
+        )
+        
+        time.sleep(20)
+
+ 
+    assert len(sequence_numbers_0) == 10
+    assert len(sequence_numbers_1) == 10
+
+    try:
+        checkpoint_store._container_client.delete_container()
+    except:
+        pass
 
 @pytest.mark.liveTest
 def test_receive_no_partition(connstr_senders, uamqp_transport):