Feedback

Author: andystaples

durabletask/client.py

@@ -57,10 +57,10 @@ def raise_if_failed(self):
                 self.failure_details)
 
 
+@dataclass
 class PurgeInstancesResult:
-    def __init__(self, deleted_instance_count: int, is_complete: bool):
-        self.deleted_instance_count = deleted_instance_count
-        self.is_complete = is_complete
+    deleted_instance_count: int
+    is_complete: bool
 
 
 class OrchestrationFailedError(Exception):
@@ -183,7 +183,7 @@ def get_orchestration_state_by(self,
                                    _continuation_token: Optional[pb2.StringValue] = None
                                    ) -> List[OrchestrationState]:
         if max_instance_count is None:
-            # DTS backend does not behave well with max_instance_count = None, so we set to max 32-bit signed value
+            # Some backends do not behave well with max_instance_count = None, so we set to max 32-bit signed value
             max_instance_count = (1 << 31) - 1
 
         self._logger.info(f"Querying orchestration instances with filters - "
@@ -194,29 +194,31 @@ def get_orchestration_state_by(self,
                           f"fetch_inputs_and_outputs={fetch_inputs_and_outputs}, "
                           f"continuation_token={_continuation_token.value if _continuation_token else None}")
 
-        req = pb.QueryInstancesRequest(
-            query=pb.InstanceQuery(
-                runtimeStatus=[status.value for status in runtime_status] if runtime_status else None,
-                createdTimeFrom=helpers.new_timestamp(created_time_from) if created_time_from else None,
-                createdTimeTo=helpers.new_timestamp(created_time_to) if created_time_to else None,
-                maxInstanceCount=max_instance_count,
-                fetchInputsAndOutputs=fetch_inputs_and_outputs,
-                continuationToken=_continuation_token
-            )
-        )
-        resp: pb.QueryInstancesResponse = self._stub.QueryInstances(req)
-        states = [parse_orchestration_state(res) for res in resp.orchestrationState]
-        # Check the value for continuationToken - none or "0" indicates that there are no more results.
-        if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
-            self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next list of instances...")
-            states += self.get_orchestration_state_by(
-                created_time_from,
-                created_time_to,
-                runtime_status,
-                max_instance_count,
-                fetch_inputs_and_outputs,
-                _continuation_token=resp.continuationToken
+        states = []
+
+        while True:
+            req = pb.QueryInstancesRequest(
+                query=pb.InstanceQuery(
+                    runtimeStatus=[status.value for status in runtime_status] if runtime_status else None,
+                    createdTimeFrom=helpers.new_timestamp(created_time_from) if created_time_from else None,
+                    createdTimeTo=helpers.new_timestamp(created_time_to) if created_time_to else None,
+                    maxInstanceCount=max_instance_count,
+                    fetchInputsAndOutputs=fetch_inputs_and_outputs,
+                    continuationToken=_continuation_token
+                )
             )
+            resp: pb.QueryInstancesResponse = self._stub.QueryInstances(req)
+            states += [parse_orchestration_state(res) for res in resp.orchestrationState]
+            # Check the value for continuationToken - none or "0" indicates that there are no more results.
+            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
+                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next list of instances...")
+                if _continuation_token and _continuation_token.value and _continuation_token.value == resp.continuationToken.value:
+                    self._logger.warning(f"Received the same continuation token value {resp.continuationToken.value} again, stopping to avoid infinite loop.")
+                    break
+                _continuation_token = resp.continuationToken
+            else:
+                break
+
         states = [state for state in states if state is not None]  # Filter out any None values
         return states
 
@@ -377,28 +379,29 @@ def get_entities_by(self,
                           f"include_state={include_state}, "
                           f"include_transient={include_transient}, "
                           f"page_size={page_size}")
-        query_request = pb.QueryEntitiesRequest(
-            query=pb.EntityQuery(
-                instanceIdStartsWith=helpers.get_string_value(instance_id_starts_with),
-                lastModifiedFrom=helpers.new_timestamp(last_modified_from) if last_modified_from else None,
-                lastModifiedTo=helpers.new_timestamp(last_modified_to) if last_modified_to else None,
-                includeState=include_state,
-                includeTransient=include_transient,
-                pageSize=helpers.get_int_value(page_size),
-                continuationToken=_continuation_token
-            )
-        )
-        resp: pb.QueryEntitiesResponse = self._stub.QueryEntities(query_request)
-        entities = [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
-        if resp.continuationToken and resp.continuationToken.value != "0":
-            self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next page of entities...")
-            entities += self.get_entities_by(
-                instance_id_starts_with=instance_id_starts_with,
-                last_modified_from=last_modified_from,
-                last_modified_to=last_modified_to,
-                include_state=include_state,
-                include_transient=include_transient,
-                page_size=page_size,
-                _continuation_token=resp.continuationToken
+
+        entities = []
+
+        while True:
+            query_request = pb.QueryEntitiesRequest(
+                query=pb.EntityQuery(
+                    instanceIdStartsWith=helpers.get_string_value(instance_id_starts_with),
+                    lastModifiedFrom=helpers.new_timestamp(last_modified_from) if last_modified_from else None,
+                    lastModifiedTo=helpers.new_timestamp(last_modified_to) if last_modified_to else None,
+                    includeState=include_state,
+                    includeTransient=include_transient,
+                    pageSize=helpers.get_int_value(page_size),
+                    continuationToken=_continuation_token
+                )
             )
+            resp: pb.QueryEntitiesResponse = self._stub.QueryEntities(query_request)
+            entities += [EntityMetadata.from_entity_metadata(entity, query_request.query.includeState) for entity in resp.entities]
+            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
+                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, fetching next page of entities...")
+                if _continuation_token and _continuation_token.value and _continuation_token.value == resp.continuationToken.value:
+                    self._logger.warning(f"Received the same continuation token value {resp.continuationToken.value} again, stopping to avoid infinite loop.")
+                    break
+                _continuation_token = resp.continuationToken
+            else:
+                break
         return entities

tests/durabletask-azuremanaged/test_dts_batch_actions.py

@@ -37,20 +37,18 @@ def test_get_all_orchestration_states():
     assert this_orch.instance_id == id
 
     assert all_orchestrations is not None
-    assert len(all_orchestrations) > 1
-    print(f"Received {len(all_orchestrations)} orchestrations")
-    assert len([o for o in all_orchestrations if o.instance_id == id]) == 1
-    orchestration_state = [o for o in all_orchestrations if o.instance_id == id][0]
+    matching_orchestrations = [o for o in all_orchestrations if o.instance_id == id]
+    assert len(matching_orchestrations) == 1
+    orchestration_state = matching_orchestrations[0]
     assert orchestration_state.runtime_status == client.OrchestrationStatus.COMPLETED
     assert orchestration_state.serialized_input is None
     assert orchestration_state.serialized_output is None
     assert orchestration_state.failure_details is None
 
     assert all_orchestrations_with_state is not None
-    assert len(all_orchestrations_with_state) > 1
-    print(f"Received {len(all_orchestrations_with_state)} orchestrations")
-    assert len([o for o in all_orchestrations_with_state if o.instance_id == id]) == 1
-    orchestration_state = [o for o in all_orchestrations_with_state if o.instance_id == id][0]
+    matching_orchestrations = [o for o in all_orchestrations_with_state if o.instance_id == id]
+    assert len(matching_orchestrations) == 1
+    orchestration_state = matching_orchestrations[0]
     assert orchestration_state.runtime_status == client.OrchestrationStatus.COMPLETED
     assert orchestration_state.serialized_input == '"Hello"'
     assert orchestration_state.serialized_output == '"Complete"'

tests/durabletask/test_batch_actions.py

@@ -1,6 +1,12 @@
 
+import pytest
 from durabletask import client, task, worker
 
+# NOTE: These tests assume a sidecar process is running. Example command:
+#       go install github.com/microsoft/durabletask-go@main
+#       durabletask-go --port 4001
+pytestmark = pytest.mark.e2e
+
 
 def empty_orchestrator(ctx: task.OrchestrationContext, _):
     return "Complete"
@@ -23,10 +29,9 @@ def test_get_all_orchestration_states():
     assert this_orch.instance_id == id
 
     assert all_orchestrations is not None
-    assert len(all_orchestrations) > 1
-    print(f"Received {len(all_orchestrations)} orchestrations")
-    assert len([o for o in all_orchestrations if o.instance_id == id]) == 1
-    orchestration_state = [o for o in all_orchestrations if o.instance_id == id][0]
+    matching_orchestrations = [o for o in all_orchestrations if o.instance_id == id]
+    assert len(matching_orchestrations) == 1
+    orchestration_state = matching_orchestrations[0]
     assert orchestration_state.runtime_status == client.OrchestrationStatus.COMPLETED
     assert orchestration_state.serialized_input == '"Hello"'
     assert orchestration_state.serialized_output == '"Complete"'