From 913b4b6e4d4552268a527b71463a0f63a9782ace Mon Sep 17 00:00:00 2001
From: Chad Retz <chad@temporal.io>
Date: Wed, 17 Jul 2024 12:54:17 -0500
Subject: [PATCH] Support query for listing schedules (#581)

Fixes #573
---
 temporalio/client.py |  7 +++++++
 tests/test_client.py | 25 +++++++++++++++++++++++--
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/temporalio/client.py b/temporalio/client.py
index 8b203207..b7a4e93f 100644
--- a/temporalio/client.py
+++ b/temporalio/client.py
@@ -966,6 +966,7 @@ def get_schedule_handle(self, id: str) -> ScheduleHandle:
 
     async def list_schedules(
         self,
+        query: Optional[str] = None,
         *,
         page_size: int = 1000,
         next_page_token: Optional[bytes] = None,
@@ -982,6 +983,9 @@ async def list_schedules(
 
         Args:
             page_size: Maximum number of results for each page.
+            query: A Temporal visibility list filter. See Temporal documentation
+                concerning visibility list filters including behavior when left
+                unset.
             next_page_token: A previously obtained next page token if doing
                 pagination. Usually not needed as the iterator automatically
                 starts from the beginning.
@@ -998,6 +1002,7 @@ async def list_schedules(
                 next_page_token=next_page_token,
                 rpc_metadata=rpc_metadata,
                 rpc_timeout=rpc_timeout,
+                query=query,
             )
         )
 
@@ -4195,6 +4200,7 @@ async def fetch_next_page(self, *, page_size: Optional[int] = None) -> None:
                 namespace=self._client.namespace,
                 maximum_page_size=page_size or self._input.page_size,
                 next_page_token=self._next_page_token or b"",
+                query=self._input.query or "",
             ),
             retry=True,
             metadata=self._input.rpc_metadata,
@@ -4704,6 +4710,7 @@ class ListSchedulesInput:
     next_page_token: Optional[bytes]
     rpc_metadata: Mapping[str, str]
     rpc_timeout: Optional[timedelta]
+    query: Optional[str] = None
 
 
 @dataclass
diff --git a/tests/test_client.py b/tests/test_client.py
index d596facb..3c41409d 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -909,10 +909,20 @@ async def update_desc_get_action_count() -> int:
     )
 
     # Create 4 more schedules of the same type and confirm they are in list
-    # eventually
+    # eventually. Two of them we will create with search attributes.
+    keyword_attr_key = SearchAttributeKey.for_keyword("python-test-schedule-keyword")
+    await ensure_search_attributes_present(client, keyword_attr_key)
     expected_ids = [handle.id]
     for i in range(4):
-        new_handle = await client.create_schedule(f"{handle.id}-{i + 1}", desc.schedule)
+        new_handle = await client.create_schedule(
+            f"{handle.id}-{i + 1}",
+            desc.schedule,
+            search_attributes=TypedSearchAttributes(
+                [SearchAttributePair(keyword_attr_key, "some-schedule-attr")]
+            )
+            if i >= 2
+            else None,
+        )
         expected_ids.append(new_handle.id)
 
     async def list_ids() -> List[str]:
@@ -925,6 +935,17 @@ async def list_ids() -> List[str]:
 
     await assert_eq_eventually(expected_ids, list_ids)
 
+    # Now do a list w/ query for certain search attributes and confirm
+    list_descs = [
+        d
+        async for d in await client.list_schedules(
+            "`python-test-schedule-keyword` = 'some-schedule-attr'"
+        )
+    ]
+    assert len(list_descs) == 2
+    assert list_descs[0].id in [f"{handle.id}-3", f"{handle.id}-4"]
+    assert list_descs[1].id in [f"{handle.id}-3", f"{handle.id}-4"]
+
     # Delete all of the schedules
     for id in await list_ids():
         await client.get_schedule_handle(id).delete()