PR Feedback

Author: andystaples

durabletask/client.py

@@ -57,6 +57,27 @@ def raise_if_failed(self):
                 self.failure_details)
 
 
+@dataclass
+class OrchestrationQuery:
+    created_time_from: Optional[datetime] = None
+    created_time_to: Optional[datetime] = None
+    runtime_status: Optional[List[OrchestrationStatus]] = None
+    # Some backends don't respond well with max_instance_count = None, so we use the integer limit for non-paginated
+    # results instead.
+    max_instance_count: Optional[int] = (1 << 31) - 1
+    fetch_inputs_and_outputs: bool = False
+
+
+@dataclass
+class EntityQuery:
+    instance_id_starts_with: Optional[str] = None
+    last_modified_from: Optional[datetime] = None
+    last_modified_to: Optional[datetime] = None
+    include_state: bool = True
+    include_transient: bool = False
+    page_size: Optional[int] = None
+
+
 @dataclass
 class PurgeInstancesResult:
     deleted_instance_count: int
@@ -170,46 +191,24 @@ def get_orchestration_state(self, instance_id: str, *, fetch_payloads: bool = Tr
         return new_orchestration_state(req.instanceId, res)
 
     def get_all_orchestration_states(self,
-                                     max_instance_count: Optional[int] = None,
-                                     fetch_inputs_and_outputs: bool = False) -> List[OrchestrationState]:
-        return self.get_orchestration_state_by(
-            created_time_from=None,
-            created_time_to=None,
-            runtime_status=None,
-            max_instance_count=max_instance_count,
-            fetch_inputs_and_outputs=fetch_inputs_and_outputs
-        )
+                                     orchestration_query: Optional[OrchestrationQuery] = None
+                                     ) -> List[OrchestrationState]:
+        if orchestration_query is None:
+            orchestration_query = OrchestrationQuery()
+        _continuation_token = None
 
-    def get_orchestration_state_by(self,
-                                   created_time_from: Optional[datetime] = None,
-                                   created_time_to: Optional[datetime] = None,
-                                   runtime_status: Optional[List[OrchestrationStatus]] = None,
-                                   max_instance_count: Optional[int] = None,
-                                   fetch_inputs_and_outputs: bool = False,
-                                   _continuation_token: Optional[pb2.StringValue] = None
-                                   ) -> List[OrchestrationState]:
-        if max_instance_count is None:
-            # 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 - "
-                          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}")
+        self._logger.info(f"Querying orchestration instances with query: {orchestration_query}")
 
         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,
+                    runtimeStatus=[status.value for status in orchestration_query.runtime_status] if orchestration_query.runtime_status else None,
+                    createdTimeFrom=helpers.new_timestamp(orchestration_query.created_time_from) if orchestration_query.created_time_from else None,
+                    createdTimeTo=helpers.new_timestamp(orchestration_query.created_time_to) if orchestration_query.created_time_to else None,
+                    maxInstanceCount=orchestration_query.max_instance_count,
+                    fetchInputsAndOutputs=orchestration_query.fetch_inputs_and_outputs,
                     continuationToken=_continuation_token
                 )
             )
@@ -318,7 +317,6 @@ def purge_orchestrations_by(self,
                           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(
                 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,
@@ -357,46 +355,24 @@ def get_entity(self,
         return EntityMetadata.from_entity_metadata(res.entity, include_state)
 
     def get_all_entities(self,
-                         include_state: bool = True,
-                         include_transient: bool = False,
-                         page_size: Optional[int] = None) -> List[EntityMetadata]:
-        return self.get_entities_by(
-            instance_id_starts_with=None,
-            last_modified_from=None,
-            last_modified_to=None,
-            include_state=include_state,
-            include_transient=include_transient,
-            page_size=page_size
-        )
+                         entity_query: Optional[EntityQuery] = None) -> List[EntityMetadata]:
+        if entity_query is None:
+            entity_query = EntityQuery()
+        _continuation_token = None
 
-    def get_entities_by(self,
-                        instance_id_starts_with: Optional[str] = None,
-                        last_modified_from: Optional[datetime] = None,
-                        last_modified_to: Optional[datetime] = None,
-                        include_state: bool = True,
-                        include_transient: bool = False,
-                        page_size: Optional[int] = None,
-                        _continuation_token: Optional[pb2.StringValue] = None
-                        ) -> List[EntityMetadata]:
-        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}")
+        self._logger.info(f"Retrieving entities by filter: {entity_query}")
 
         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),
+                    instanceIdStartsWith=helpers.get_string_value(entity_query.instance_id_starts_with),
+                    lastModifiedFrom=helpers.new_timestamp(entity_query.last_modified_from) if entity_query.last_modified_from else None,
+                    lastModifiedTo=helpers.new_timestamp(entity_query.last_modified_to) if entity_query.last_modified_to else None,
+                    includeState=entity_query.include_state,
+                    includeTransient=entity_query.include_transient,
+                    pageSize=helpers.get_int_value(entity_query.page_size),
                     continuationToken=_continuation_token
                 )
             )
@@ -414,13 +390,13 @@ def get_entities_by(self,
 
     def clean_entity_storage(self,
                              remove_empty_entities: bool = True,
-                             release_orphaned_locks: bool = True,
-                             _continuation_token: Optional[pb2.StringValue] = None
+                             release_orphaned_locks: bool = True
                              ) -> CleanEntityStorageResult:
         self._logger.info("Cleaning entity storage")
 
         empty_entities_removed = 0
         orphaned_locks_released = 0
+        _continuation_token = None
 
         while True:
             req = pb.CleanEntityStorageRequest(

tests/durabletask-azuremanaged/test_dts_batch_actions.py

@@ -1,4 +1,5 @@
 
+import asyncio
 import logging
 import os
 import time
@@ -37,7 +38,9 @@ def test_get_all_orchestration_states():
         c.wait_for_orchestration_completion(id, timeout=30)
 
         all_orchestrations = c.get_all_orchestration_states()
-        all_orchestrations_with_state = c.get_all_orchestration_states(fetch_inputs_and_outputs=True)
+        query = client.OrchestrationQuery()
+        query.fetch_inputs_and_outputs = True
+        all_orchestrations_with_state = c.get_all_orchestration_states(query)
         this_orch = c.get_orchestration_state(id)
 
     assert this_orch is not None
@@ -84,16 +87,16 @@ def test_get_orchestration_state_by_status():
             pass  # Expected failure
 
         # Query by completed status
-        completed_orchestrations = c.get_orchestration_state_by(
-            runtime_status=[client.OrchestrationStatus.COMPLETED],
-            fetch_inputs_and_outputs=True
-        )
+        query = client.OrchestrationQuery()
+        query.runtime_status = [client.OrchestrationStatus.COMPLETED]
+        query.fetch_inputs_and_outputs = True
+        completed_orchestrations = c.get_all_orchestration_states(query)
 
         # Query by failed status
-        failed_orchestrations = c.get_orchestration_state_by(
-            runtime_status=[client.OrchestrationStatus.FAILED],
-            fetch_inputs_and_outputs=True
-        )
+        query = client.OrchestrationQuery()
+        query.runtime_status = [client.OrchestrationStatus.FAILED]
+        query.fetch_inputs_and_outputs = True
+        failed_orchestrations = c.get_all_orchestration_states(query)
 
     assert len([o for o in completed_orchestrations if o.instance_id == completed_id]) == 1
     completed_orch = [o for o in completed_orchestrations if o.instance_id == completed_id][0]
@@ -125,17 +128,20 @@ def test_get_orchestration_state_by_time_range():
         after_creation = datetime.now(timezone.utc) + timedelta(seconds=5)
 
         # Query by time range
-        orchestrations_in_range = c.get_orchestration_state_by(
+        query = client.OrchestrationQuery(
             created_time_from=before_creation,
             created_time_to=after_creation,
             fetch_inputs_and_outputs=True
         )
+        orchestrations_in_range = c.get_all_orchestration_states(query)
 
         # Query outside time range
-        orchestrations_outside_range = c.get_orchestration_state_by(
+        query = client.OrchestrationQuery(
             created_time_from=after_creation,
-            created_time_to=after_creation + timedelta(hours=1)
+            created_time_to=after_creation + timedelta(hours=1),
+            fetch_inputs_and_outputs=True
         )
+        orchestrations_outside_range = c.get_all_orchestration_states(query)
 
     assert len([o for o in orchestrations_in_range if o.instance_id == id]) == 1
     assert len([o for o in orchestrations_outside_range if o.instance_id == id]) == 0
@@ -171,7 +177,8 @@ def emit(self, record):
             c.wait_for_orchestration_completion(id, timeout=30)
 
         # Query with max_instance_count=2
-        orchestrations = c.get_orchestration_state_by(max_instance_count=2)
+        query = client.OrchestrationQuery(max_instance_count=2)
+        orchestrations = c.get_all_orchestration_states(query)
 
     # Should return more than 2 instances since we created at least 3
     assert len(orchestrations) > 2
@@ -303,16 +310,18 @@ def counter_entity(ctx: entities.EntityContext, input):
         # Create entity
         entity_id = entities.EntityInstanceId("counter_entity", "testCounter1")
         c.signal_entity(entity_id, "add", 5)
-        time.sleep(2)  # Wait for signal to be processed
+        asyncio.run(asyncio.sleep(2))  # Wait for signal to be processed
 
         # Get all entities without state
-        all_entities = c.get_all_entities(include_state=False)
+        query = client.EntityQuery(include_state=False)
+        all_entities = c.get_all_entities(query)
         assert len([e for e in all_entities if e.id == entity_id]) == 1
         entity_without_state = [e for e in all_entities if e.id == entity_id][0]
         assert entity_without_state.get_state(int) is None
 
         # Get all entities with state
-        all_entities_with_state = c.get_all_entities(include_state=True)
+        query = client.EntityQuery(include_state=True)
+        all_entities_with_state = c.get_all_entities(query)
         assert len([e for e in all_entities_with_state if e.id == entity_id]) == 1
         entity_with_state = [e for e in all_entities_with_state if e.id == entity_id][0]
         assert entity_with_state.get_state(int) == 5
@@ -337,18 +346,20 @@ def counter_entity(ctx: entities.EntityContext, input):
 
         c.signal_entity(entity_id_1, "set", 10)
         c.signal_entity(entity_id_2, "set", 20)
-        time.sleep(2)  # Wait for signals to be processed
+        asyncio.run(asyncio.sleep(2))  # Wait for signals to be processed
 
         # Query by prefix
-        entities_prefix1 = c.get_entities_by(
+        query = client.EntityQuery(
             instance_id_starts_with="@counter_entity@prefix1",
             include_state=True
         )
+        entities_prefix1 = c.get_all_entities(query)
 
-        entities_prefix2 = c.get_entities_by(
+        query = client.EntityQuery(
             instance_id_starts_with="@counter_entity@prefix2",
             include_state=True
         )
+        entities_prefix2 = c.get_all_entities(query)
 
     assert len([e for e in entities_prefix1 if e.id == entity_id_1]) == 1
     assert len([e for e in entities_prefix1 if e.id == entity_id_2]) == 0
@@ -376,22 +387,24 @@ def simple_entity(ctx: entities.EntityContext, input):
         # Create entity
         entity_id = entities.EntityInstanceId("simple_entity", "timeTestEntity")
         c.signal_entity(entity_id, "set", "test_value")
-        time.sleep(2)  # Wait for signal to be processed
+        asyncio.run(asyncio.sleep(2))  # Wait for signal to be processed
 
         after_creation = datetime.now(timezone.utc) + timedelta(seconds=5)
 
         # Query by time range
-        entities_in_range = c.get_entities_by(
+        query = client.EntityQuery(
             last_modified_from=before_creation,
             last_modified_to=after_creation,
             include_state=True
         )
+        entities_in_range = c.get_all_entities(query)
 
         # Query outside time range
-        entities_outside_range = c.get_entities_by(
+        query = client.EntityQuery(
             last_modified_from=after_creation,
             last_modified_to=after_creation + timedelta(hours=1)
         )
+        entities_outside_range = c.get_all_entities(query)
 
     assert len([e for e in entities_in_range if e.id == entity_id]) == 1
     assert len([e for e in entities_outside_range if e.id == entity_id]) == 0
@@ -412,13 +425,14 @@ class EmptyEntity(entities.DurableEntity):
         # Create an entity and then delete its state to make it empty
         entity_id = entities.EntityInstanceId("EmptyEntity", "toClean")
         c.signal_entity(entity_id, "delete")
-        time.sleep(2)  # Wait for signal to be processed
+        asyncio.run(asyncio.sleep(2))  # Wait for signal to be processed
 
         # Clean entity storage
         result = c.clean_entity_storage(
             remove_empty_entities=True,
             release_orphaned_locks=True
         )
 
-    # Verify clean result - we expect at least the entity we just deleted to be removed
-    assert result.empty_entities_removed >= 0
+    # Verify clean result - DTS backend always returns 0, as it has its own mechanism for entity state purge
+    assert result.empty_entities_removed == 0
+    assert result.orphaned_locks_released == 0