Add timeout support for filtered orchestration purge

CHANGELOG.md

@@ -7,6 +7,12 @@ adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## Unreleased
 
+ADDED
+
+- Added an optional timeout to filtered orchestration purges. Callers can now
+limit a purge operation's duration and inspect `PurgeInstancesResult.is_complete`
+to determine whether it finished.
+
 ## v1.9.0
 
 ADDED

durabletask/client.py

@@ -8,7 +8,7 @@
 import uuid
 from collections.abc import AsyncIterable, Iterable, Sequence
 from dataclasses import dataclass, field
-from datetime import datetime
+from datetime import datetime, timedelta
 from enum import Enum
 from typing import Any, Generic, Protocol, TypeVar, cast, overload
 
@@ -846,13 +846,15 @@ def purge_orchestrations_by(self,
                                 created_time_from: datetime | None = None,
                                 created_time_to: datetime | None = None,
                                 runtime_status: list[OrchestrationStatus] | None = None,
-                                recursive: bool = False) -> PurgeInstancesResult:
+                                recursive: bool = False,
+                                timeout: timedelta | None = None) -> 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}")
-        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive)
+                          f"recursive={recursive}, "
+                          f"timeout={timeout}")
+        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive, timeout)
         resp: pb.PurgeInstancesResponse = self._stub.PurgeInstances(req)
         return PurgeInstancesResult(resp.deletedInstanceCount, resp.isComplete.value)
 
@@ -1380,13 +1382,15 @@ async def purge_orchestrations_by(self,
                                       created_time_from: datetime | None = None,
                                       created_time_to: datetime | None = None,
                                       runtime_status: list[OrchestrationStatus] | None = None,
-                                      recursive: bool = False) -> PurgeInstancesResult:
+                                      recursive: bool = False,
+                                      timeout: timedelta | None = None) -> 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}")
-        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive)
+                          f"recursive={recursive}, "
+                          f"timeout={timeout}")
+        req = build_purge_by_filter_req(created_time_from, created_time_to, runtime_status, recursive, timeout)
         resp: pb.PurgeInstancesResponse = await self._get_stub().PurgeInstances(req)
         return PurgeInstancesResult(resp.deletedInstanceCount, resp.isComplete.value)
 

durabletask/internal/client_helpers.py

@@ -6,10 +6,10 @@
 import logging
 import uuid
 from collections.abc import Sequence
-from datetime import datetime, timezone
+from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, TypeVar
 
-from google.protobuf import wrappers_pb2
+from google.protobuf import duration_pb2, wrappers_pb2
 
 import durabletask.internal.helpers as helpers
 import durabletask.internal.orchestrator_service_pb2 as pb
@@ -114,14 +114,24 @@ def build_purge_by_filter_req(
         created_time_from: datetime | None,
         created_time_to: datetime | None,
         runtime_status: list[OrchestrationStatus] | None,
-        recursive: bool) -> pb.PurgeInstancesRequest:
+        recursive: bool,
+        timeout: timedelta | None = None) -> pb.PurgeInstancesRequest:
     """Build a PurgeInstancesRequest for purging orchestrations by filter."""
+    if timeout is not None and timeout <= timedelta():
+        raise ValueError("timeout must be greater than zero.")
+
+    purge_filter = 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,
+        runtimeStatus=[status.value for status in runtime_status] if runtime_status else None
+    )
+    if timeout is not None:
+        timeout_duration = duration_pb2.Duration()
+        timeout_duration.FromTimedelta(timeout)
+        purge_filter.timeout.CopyFrom(timeout_duration)
+
     return pb.PurgeInstancesRequest(
-        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,
-            runtimeStatus=[status.value for status in runtime_status] if runtime_status else None
-        ),
+        purgeInstanceFilter=purge_filter,
         recursive=recursive
     )
 

tests/durabletask/test_purge_timeout.py

@@ -0,0 +1,88 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from datetime import timedelta
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+from google.protobuf import wrappers_pb2
+
+import durabletask.internal.orchestrator_service_pb2 as pb
+from durabletask.client import AsyncTaskHubGrpcClient, TaskHubGrpcClient
+
+
+def test_sync_filtered_purge_serializes_timeout():
+    stub = MagicMock()
+    stub.PurgeInstances.return_value = pb.PurgeInstancesResponse(
+        deletedInstanceCount=1,
+        isComplete=wrappers_pb2.BoolValue(value=False),
+    )
+
+    with (
+        patch("durabletask.client.shared.get_grpc_channel", return_value=MagicMock()),
+        patch("durabletask.client.stubs.TaskHubSidecarServiceStub", return_value=stub),
+    ):
+        client = TaskHubGrpcClient()
+        result = client.purge_orchestrations_by(timeout=timedelta(seconds=1, microseconds=500000))
+
+    request = stub.PurgeInstances.call_args.args[0]
+    assert request.purgeInstanceFilter.timeout.seconds == 1
+    assert request.purgeInstanceFilter.timeout.nanos == 500000000
+    assert result.is_complete is False
+
+
+def test_sync_filtered_purge_omits_timeout_when_not_supplied():
+    stub = MagicMock()
+    stub.PurgeInstances.return_value = pb.PurgeInstancesResponse(
+        isComplete=wrappers_pb2.BoolValue(value=True),
+    )
+
+    with (
+        patch("durabletask.client.shared.get_grpc_channel", return_value=MagicMock()),
+        patch("durabletask.client.stubs.TaskHubSidecarServiceStub", return_value=stub),
+    ):
+        client = TaskHubGrpcClient()
+        client.purge_orchestrations_by()
+
+    request = stub.PurgeInstances.call_args.args[0]
+    assert not request.purgeInstanceFilter.HasField("timeout")
+
+
+@pytest.mark.parametrize("timeout", [timedelta(), timedelta(seconds=-1)])
+def test_sync_filtered_purge_rejects_non_positive_timeout(timeout):
+    client = TaskHubGrpcClient(channel=MagicMock())
+
+    with pytest.raises(ValueError, match="timeout must be greater than zero"):
+        client.purge_orchestrations_by(timeout=timeout)
+
+
+@pytest.mark.asyncio
+async def test_async_filtered_purge_serializes_timeout():
+    stub = MagicMock()
+    stub.PurgeInstances = AsyncMock(return_value=pb.PurgeInstancesResponse(
+        deletedInstanceCount=1,
+        isComplete=wrappers_pb2.BoolValue(value=False),
+    ))
+    channel = MagicMock()
+    channel.close = AsyncMock()
+
+    with (
+        patch("durabletask.client.shared.get_async_grpc_channel", return_value=channel),
+        patch("durabletask.client.stubs.TaskHubSidecarServiceStub", return_value=stub),
+    ):
+        client = AsyncTaskHubGrpcClient()
+        result = await client.purge_orchestrations_by(timeout=timedelta(milliseconds=250))
+        await client.close()
+
+    request = stub.PurgeInstances.call_args.args[0]
+    assert request.purgeInstanceFilter.timeout.seconds == 0
+    assert request.purgeInstanceFilter.timeout.nanos == 250000000
+    assert result.is_complete is False
+
+
+@pytest.mark.asyncio
+async def test_async_filtered_purge_rejects_non_positive_timeout():
+    client = AsyncTaskHubGrpcClient(channel=MagicMock())
+
+    with pytest.raises(ValueError, match="timeout must be greater than zero"):
+        await client.purge_orchestrations_by(timeout=timedelta())