Merge changes from base branch (andystaples/add-batch-actions)

Copilot · Copilot · commit 7357abe9def3 · 2026-02-12T21:56:21.000Z
diff --git a/durabletask/client.py b/durabletask/client.py
@@ -57,16 +57,16 @@ 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
 
 
+@dataclass
 class CleanEntityStorageResult:
-    def __init__(self, empty_entities_removed: int, orphaned_locks_released: int):
-        self.empty_entities_removed = empty_entities_removed
-        self.orphaned_locks_released = orphaned_locks_released
+    empty_entities_removed: int
+    orphaned_locks_released: int
 
 
 class OrchestrationFailedError(Exception):
@@ -153,7 +153,7 @@ def schedule_new_orchestration(self, orchestrator: Union[task.Orchestrator[TInpu
         req = pb.CreateInstanceRequest(
             name=name,
             instanceId=instance_id if instance_id else uuid.uuid4().hex,
-            input=helpers.get_string_value(shared.to_json(input)),
+            input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
             scheduledStartTimestamp=helpers.new_timestamp(start_at) if start_at else None,
             version=helpers.get_string_value(version if version else self.default_version),
             orchestrationIdReusePolicy=reuse_id_policy,
@@ -189,31 +189,42 @@ 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
-        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
+
+        self._logger.info(f"Querying orchestration instances with filters - "
+                          f"created_time_from={created_time_from}, "
+                          f"created_time_to={created_time_to}, "
+                          f"runtime_status={[str(status) for status in runtime_status] if runtime_status else None}, "
+                          f"max_instance_count={max_instance_count}, "
+                          f"fetch_inputs_and_outputs={fetch_inputs_and_outputs}, "
+                          f"continuation_token={_continuation_token.value if _continuation_token else None}")
+
+        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
 
@@ -264,7 +275,8 @@ def raise_orchestration_event(self, instance_id: str, event_name: str, *,
         req = pb.RaiseEventRequest(
             instanceId=instance_id,
             name=event_name,
-            input=helpers.get_string_value(shared.to_json(data)))
+            input=helpers.get_string_value(shared.to_json(data) if data is not None else None)
+        )
 
         self._logger.info(f"Raising event '{event_name}' for instance '{instance_id}'.")
         self._stub.RaiseEvent(req)
@@ -274,7 +286,7 @@ def terminate_orchestration(self, instance_id: str, *,
                                 recursive: bool = True):
         req = pb.TerminateRequest(
             instanceId=instance_id,
-            output=helpers.get_string_value(shared.to_json(output)),
+            output=helpers.get_string_value(shared.to_json(output) if output is not None else None),
             recursive=recursive)
 
         self._logger.info(f"Terminating instance '{instance_id}'.")
@@ -301,6 +313,11 @@ def purge_orchestrations_by(self,
                                 created_time_to: Optional[datetime] = None,
                                 runtime_status: Optional[List[OrchestrationStatus]] = None,
                                 recursive: bool = False) -> PurgeInstancesResult:
+        self._logger.info("Purging orchestrations by filter: "
+                          f"created_time_from={created_time_from}, "
+                          f"created_time_to={created_time_to}, "
+                          f"runtime_status={[str(status) for status in runtime_status] if runtime_status else None}, "
+                          f"recursive={recursive}")
         resp: pb.PurgeInstancesResponse = self._stub.PurgeInstances(pb.PurgeInstancesRequest(
             instanceId=None,
             purgeInstanceFilter=pb.PurgeInstanceFilter(
@@ -319,7 +336,7 @@ def signal_entity(self,
         req = pb.SignalEntityRequest(
             instanceId=str(entity_instance_id),
             name=operation_name,
-            input=helpers.get_string_value(shared.to_json(input)),
+            input=helpers.get_string_value(shared.to_json(input) if input is not None else None),
             requestId=str(uuid.uuid4()),
             scheduledTime=None,
             parentTraceContext=None,
@@ -362,31 +379,38 @@ def get_entities_by(self,
                         page_size: Optional[int] = None,
                         _continuation_token: Optional[pb2.StringValue] = None
                         ) -> List[EntityMetadata]:
-        self._logger.info("Getting entities")
-        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
+        self._logger.info(f"Retrieving entities by filter: "
+                          f"instance_id_starts_with={instance_id_starts_with}, "
+                          f"last_modified_from={last_modified_from}, "
+                          f"last_modified_to={last_modified_to}, "
+                          f"include_state={include_state}, "
+                          f"include_transient={include_transient}, "
+                          f"page_size={page_size}")
+
+        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
 
     def clean_entity_storage(self,
@@ -395,23 +419,27 @@ def clean_entity_storage(self,
                              _continuation_token: Optional[pb2.StringValue] = None
                              ) -> CleanEntityStorageResult:
         self._logger.info("Cleaning entity storage")
-        req = pb.CleanEntityStorageRequest(
-            removeEmptyEntities=remove_empty_entities,
-            releaseOrphanedLocks=release_orphaned_locks,
-            continuationToken=_continuation_token
-        )
-        resp: pb.CleanEntityStorageResponse = self._stub.CleanEntityStorage(req)
-        empty_entities_removed = resp.emptyEntitiesRemoved
-        orphaned_locks_released = resp.orphanedLocksReleased
-
-        if resp.continuationToken and resp.continuationToken.value != "0":
-            self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, cleaning next page...")
-            next_result = self.clean_entity_storage(
-                remove_empty_entities=remove_empty_entities,
-                release_orphaned_locks=release_orphaned_locks,
-                _continuation_token=resp.continuationToken
+        
+        empty_entities_removed = 0
+        orphaned_locks_released = 0
+        
+        while True:
+            req = pb.CleanEntityStorageRequest(
+                removeEmptyEntities=remove_empty_entities,
+                releaseOrphanedLocks=release_orphaned_locks,
+                continuationToken=_continuation_token
             )
-            empty_entities_removed += next_result.empty_entities_removed
-            orphaned_locks_released += next_result.orphaned_locks_released
+            resp: pb.CleanEntityStorageResponse = self._stub.CleanEntityStorage(req)
+            empty_entities_removed += resp.emptyEntitiesRemoved
+            orphaned_locks_released += resp.orphanedLocksReleased
+
+            if resp.continuationToken and resp.continuationToken.value and resp.continuationToken.value != "0":
+                self._logger.info(f"Received continuation token with value {resp.continuationToken.value}, cleaning next page...")
+                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 CleanEntityStorageResult(empty_entities_removed, orphaned_locks_released)
diff --git a/durabletask/entities/entity_metadata.py b/durabletask/entities/entity_metadata.py
@@ -1,6 +1,5 @@
 from datetime import datetime, timezone
 from typing import Any, Optional, Type, TypeVar, Union, overload
-from warnings import deprecated
 from durabletask.entities.entity_instance_id import EntityInstanceId
 
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -43,7 +42,6 @@ def __init__(self,
         self.includes_state = includes_state
         self._state = state
 
-    @deprecated("This method is deprecated. Use 'from_entity_metadata' instead.")
     @staticmethod
     def from_entity_response(entity_response: pb.GetEntityResponse, includes_state: bool):
         return EntityMetadata.from_entity_metadata(entity_response.entity, includes_state)
diff --git a/tests/durabletask-azuremanaged/test_dts_batch_actions.py b/tests/durabletask-azuremanaged/test_dts_batch_actions.py
@@ -43,20 +43,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"'
diff --git a/tests/durabletask/test_batch_actions.py b/tests/durabletask/test_batch_actions.py
