add dynamic store connection APIs (#28655)

## Summary & Motivation
Creates pagination connection API methods for fetching partition keys
from the dynamic store (backed by the instance / event_log_storage).

## How I Tested These Changes
BK

Author: prha

python_modules/dagster/dagster/_core/definitions/partition.py

@@ -479,6 +479,14 @@ def __init__(self, instance: DagsterInstance):
     def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
         return self._instance.get_dynamic_partitions(partitions_def_name)
 
+    @cached_method
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        return self._instance.get_paginated_dynamic_partitions(
+            partitions_def_name=partitions_def_name, limit=limit, ascending=ascending, cursor=cursor
+        )
+
     @cached_method
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         return self._instance.has_dynamic_partition(partitions_def_name, partition_key)

python_modules/dagster/dagster/_core/definitions/run_request.py

@@ -28,6 +28,7 @@
     ASSET_PARTITION_RANGE_START_TAG,
     PARTITION_NAME_TAG,
 )
+from dagster._core.types.pagination import PaginatedResults
 from dagster._record import IHaveNew, LegacyNamedTupleMixin, record, record_custom
 from dagster._utils.cached_method import cached_method
 from dagster._utils.error import SerializableErrorInfo
@@ -294,6 +295,15 @@ def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
         )
         return list((partition_keys | added_partition_keys) - deleted_partition_keys)
 
+    @cached_method
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        partition_keys = self.get_dynamic_partitions(partitions_def_name)
+        return PaginatedResults.create_from_sequence(
+            seq=partition_keys, limit=limit, ascending=ascending, cursor=cursor
+        )
+
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         return partition_key not in self.deleted_partition_keys_by_partitions_def_name.get(
             partitions_def_name, set()

python_modules/dagster/dagster/_core/instance/__init__.py

@@ -81,6 +81,7 @@
     TAGS_TO_MAYBE_OMIT_ON_RETRY,
     WILL_RETRY_TAG,
 )
+from dagster._core.types.pagination import PaginatedResults
 from dagster._serdes import ConfigurableClass
 from dagster._time import get_current_datetime, get_current_timestamp
 from dagster._utils import PrintFn, is_uuid, traced
@@ -336,6 +337,11 @@ class DynamicPartitionsStore(Protocol):
     @abstractmethod
     def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]: ...
 
+    @abstractmethod
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]: ...
+
     @abstractmethod
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool: ...
 
@@ -2397,6 +2403,26 @@ def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
         check.str_param(partitions_def_name, "partitions_def_name")
         return self._event_storage.get_dynamic_partitions(partitions_def_name)
 
+    @traced
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        """Get a paginatable subset of partition keys for the specified :py:class:`DynamicPartitionsDefinition`.
+
+        Args:
+            partitions_def_name (str): The name of the `DynamicPartitionsDefinition`.
+            limit (int): Maximum number of partition keys to return.
+            ascending (bool): The order of dynamic partitions to return.
+            cursor (Optional[str]): Cursor to use for pagination. Defaults to None.
+        """
+        check.str_param(partitions_def_name, "partitions_def_name")
+        check.int_param(limit, "limit")
+        check.bool_param(ascending, "ascending")
+        check.opt_str_param(cursor, "cursor")
+        return self._event_storage.get_paginated_dynamic_partitions(
+            partitions_def_name=partitions_def_name, limit=limit, ascending=ascending, cursor=cursor
+        )
+
     @public
     @traced
     def add_dynamic_partitions(

python_modules/dagster/dagster/_core/storage/event_log/base.py

@@ -36,6 +36,7 @@
 from dagster._core.storage.partition_status_cache import get_and_update_asset_status_cache_value
 from dagster._core.storage.sql import AlembicVersion
 from dagster._core.storage.tags import MULTIDIMENSIONAL_PARTITION_PREFIX
+from dagster._core.types.pagination import PaginatedResults
 from dagster._record import record
 from dagster._utils import PrintFn
 from dagster._utils.concurrency import ConcurrencyClaimStatus, ConcurrencyKeyInfo
@@ -524,6 +525,12 @@ def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
         """Get the list of partition keys for a dynamic partitions definition."""
         raise NotImplementedError()
 
+    @abstractmethod
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        raise NotImplementedError()
+
     @abstractmethod
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         """Check if a dynamic partition exists."""

python_modules/dagster/dagster/_core/storage/event_log/sql_event_log.py

@@ -95,6 +95,7 @@
     db_select,
     db_subquery,
 )
+from dagster._core.types.pagination import PaginatedResults, StorageIdCursor
 from dagster._serdes import deserialize_value, serialize_value
 from dagster._time import datetime_from_timestamp, get_current_timestamp, utc_datetime_from_naive
 from dagster._utils import PrintFn
@@ -2027,6 +2028,47 @@ def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
 
         return [cast("str", row[1]) for row in rows]
 
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        self._check_partitions_table()
+        order_by = (
+            DynamicPartitionsTable.c.id.asc() if ascending else DynamicPartitionsTable.c.id.desc()
+        )
+        query = (
+            db_select(
+                [
+                    DynamicPartitionsTable.c.id,
+                    DynamicPartitionsTable.c.partition,
+                ]
+            )
+            .where(DynamicPartitionsTable.c.partitions_def_name == partitions_def_name)
+            .order_by(order_by)
+            .limit(limit)
+        )
+        if cursor:
+            last_storage_id = StorageIdCursor.from_cursor(cursor).storage_id
+            if ascending:
+                query = query.where(DynamicPartitionsTable.c.id > last_storage_id)
+            else:
+                query = query.where(DynamicPartitionsTable.c.id < last_storage_id)
+
+        with self.index_connection() as conn:
+            rows = conn.execute(query).fetchall()
+
+        if rows:
+            next_cursor = StorageIdCursor(storage_id=cast(int, rows[-1][0])).to_string()
+        elif cursor:
+            next_cursor = cursor
+        else:
+            next_cursor = StorageIdCursor(storage_id=-1).to_string()
+
+        return PaginatedResults(
+            results=[cast(str, row[1]) for row in rows],
+            cursor=next_cursor,
+            has_more=len(rows) == limit,
+        )
+
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         self._check_partitions_table()
         query = (

python_modules/dagster/dagster/_core/storage/legacy_storage.py

@@ -29,6 +29,7 @@
 from dagster._core.storage.runs.base import RunStorage
 from dagster._core.storage.schedules.base import ScheduleStorage
 from dagster._core.storage.sql import AlembicVersion
+from dagster._core.types.pagination import PaginatedResults
 from dagster._serdes import ConfigurableClass, ConfigurableClassData
 from dagster._utils import PrintFn
 from dagster._utils.concurrency import ConcurrencyClaimStatus, ConcurrencyKeyInfo
@@ -621,6 +622,13 @@ def get_latest_asset_partition_materialization_attempts_without_materializations
     def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
         return self._storage.event_log_storage.get_dynamic_partitions(partitions_def_name)
 
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        return self._storage.event_log_storage.get_paginated_dynamic_partitions(
+            partitions_def_name=partitions_def_name, limit=limit, ascending=ascending, cursor=cursor
+        )
+
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         return self._storage.event_log_storage.has_dynamic_partition(
             partitions_def_name, partition_key

python_modules/dagster/dagster/_core/types/pagination.py

@@ -83,3 +83,22 @@ def to_string(self) -> str:
     @classmethod
     def from_cursor(cls, cursor: str):
         return deserialize_value(base64.b64decode(cursor).decode("utf-8"), cls)
+
+
+@whitelist_for_serdes
+@record
+class StorageIdCursor:
+    storage_id: int
+
+    def __str__(self) -> str:
+        return self.to_string()
+
+    def to_string(self) -> str:
+        string_serialized = serialize_value(self)
+        return base64.b64encode(bytes(string_serialized, encoding="utf-8")).decode(
+            "utf-8"
+        )
+
+    @classmethod
+    def from_cursor(cls, cursor: str):
+        return deserialize_value(base64.b64decode(cursor).decode("utf-8"), cls)

python_modules/dagster/dagster/_utils/caching_instance_queryer.py

@@ -34,6 +34,7 @@
     RunRecord,
 )
 from dagster._core.storage.tags import PARTITION_NAME_TAG
+from dagster._core.types.pagination import PaginatedResults
 from dagster._time import get_current_datetime
 from dagster._utils.cached_method import cached_method
 
@@ -644,6 +645,23 @@ def get_dynamic_partitions(self, partitions_def_name: str) -> Sequence[str]:
             )
         return self._dynamic_partitions_cache[partitions_def_name]
 
+    def get_paginated_dynamic_partitions(
+        self, partitions_def_name: str, limit: int, ascending: bool, cursor: Optional[str] = None
+    ) -> PaginatedResults[str]:
+        if partitions_def_name not in self._dynamic_partitions_cache:
+            return self.instance.get_paginated_dynamic_partitions(
+                partitions_def_name=partitions_def_name,
+                limit=limit,
+                ascending=ascending,
+                cursor=cursor,
+            )
+
+        # the full set of partition keys are cached... create a sequence connection from the cached keys
+        partition_keys = self._dynamic_partitions_cache[partitions_def_name]
+        return PaginatedResults.create_from_sequence(
+            seq=partition_keys, limit=limit, ascending=ascending, cursor=cursor
+        )
+
     def has_dynamic_partition(self, partitions_def_name: str, partition_key: str) -> bool:
         return partition_key in self.get_dynamic_partitions(partitions_def_name)
 

python_modules/dagster/dagster_tests/definitions_tests/test_dynamic_partitions.py

@@ -83,6 +83,9 @@ def my_asset():
 
         instance.add_dynamic_partitions("foo", ["a"])
         assert partitions_def.get_partition_keys(dynamic_partitions_store=instance) == ["a"]
+        assert get_paginated_partition_keys(partitions_def, dynamic_partitions_store=instance) == [
+            "a"
+        ]
         assert materialize([my_asset], instance=instance, partition_key="a").success
         materialization = instance.get_latest_materialization_event(AssetKey("my_asset"))
         assert materialization

python_modules/dagster/dagster_tests/storage_tests/utils/event_log_storage.py

@@ -6713,3 +6713,71 @@ def test_fetch_failed_materializations(self, test_run_id, storage: EventLogStora
                 failed_record.asset_key == asset_key_1
                 for failed_record in failed_records_for_partitions.records
             )
+
+    def test_dynamic_store_pagination(self, storage: EventLogStorage):
+        assert (
+            storage.get_paginated_dynamic_partitions(
+                partitions_def_name="foo", limit=1, ascending=True
+            ).results
+            == []
+        )
+
+        def get_paginated_partitions(partitions_def_name, ascending=True):
+            all_results = []
+            cursor = None
+            has_more = True
+            while has_more:
+                results = storage.get_paginated_dynamic_partitions(
+                    partitions_def_name=partitions_def_name,
+                    limit=1,
+                    ascending=ascending,
+                    cursor=cursor,
+                )
+                cursor = results.cursor
+                has_more = results.has_more
+                all_results.extend(results.results)
+            return all_results
+
+        storage.add_dynamic_partitions(
+            partitions_def_name="foo", partition_keys=["foo", "bar", "baz"]
+        )
+
+        # paginated results
+        assert get_paginated_partitions("foo") == ["foo", "bar", "baz"]
+
+        # Test for idempotency
+        storage.add_dynamic_partitions(partitions_def_name="foo", partition_keys=["foo"])
+        assert get_paginated_partitions("foo") == ["foo", "bar", "baz"]
+
+        # add more partitions
+        storage.add_dynamic_partitions(partitions_def_name="foo", partition_keys=["foo", "qux"])
+        assert get_paginated_partitions("foo") == ["foo", "bar", "baz", "qux"]
+
+        assert get_paginated_partitions("foo", ascending=False) == ["qux", "baz", "bar", "foo"]
+
+        # partial paginated results
+        results = storage.get_paginated_dynamic_partitions(
+            partitions_def_name="foo", limit=1, ascending=True
+        )
+        assert results.results == ["foo"]
+        assert results.cursor
+        post_foo_cursor = results.cursor
+        assert results.has_more
+
+        # partial reverse paginated results
+        results = storage.get_paginated_dynamic_partitions(
+            partitions_def_name="foo", limit=1, ascending=False
+        )
+        assert results.results == ["qux"]
+        assert results.cursor
+        assert results.has_more
+
+        # try cursored fetching when the keys before and after the cursor are removed
+        storage.delete_dynamic_partition(partitions_def_name="foo", partition_key="foo")
+        storage.delete_dynamic_partition(partitions_def_name="foo", partition_key="bar")
+        assert storage.get_paginated_dynamic_partitions(
+            partitions_def_name="foo", limit=1, ascending=True, cursor=post_foo_cursor
+        ).results == ["baz"]
+        assert storage.get_paginated_dynamic_partitions(
+            partitions_def_name="foo", limit=1, ascending=True
+        ).results == ["baz"]