add asset graph

Author: prha

python_modules/dagster-graphql/dagster_graphql/schema/asset_graph.py

@@ -17,7 +17,12 @@
     AutomationConditionSnapshot,
 )
 from dagster._core.definitions.freshness import FreshnessState
-from dagster._core.definitions.partition import CachingDynamicPartitionsLoader, PartitionsDefinition
+from dagster._core.definitions.partition import (
+    CachingDynamicPartitionsLoader,
+    PartitionLoadingContext,
+    PartitionsDefinition,
+    TemporalContext,
+)
 from dagster._core.definitions.partition_mapping import PartitionMapping
 from dagster._core.definitions.remote_asset_graph import RemoteAssetNode, RemoteWorkspaceAssetNode
 from dagster._core.definitions.selector import JobSelector
@@ -44,6 +49,7 @@
 from dagster._core.storage.tags import KIND_PREFIX
 from dagster._core.utils import is_valid_email
 from dagster._core.workspace.permissions import Permissions
+from dagster._time import get_current_datetime
 from packaging import version
 
 from dagster_graphql.implementation.events import iterate_metadata_entries
@@ -103,6 +109,7 @@
     GrapheneObservationEvent,
 )
 from dagster_graphql.schema.metadata import GrapheneMetadataEntry
+from dagster_graphql.schema.partition_keys import GraphenePartitionKeyConnection
 from dagster_graphql.schema.partition_mappings import GraphenePartitionMapping
 from dagster_graphql.schema.partition_sets import (
     GrapheneDimensionPartitionKeys,
@@ -300,6 +307,12 @@ class GrapheneAssetNode(graphene.ObjectType):
     opVersion = graphene.String()
     partitionDefinition = graphene.Field(GraphenePartitionDefinition)
     partitionKeys = non_null_list(graphene.String)
+    partitionKeyConnection = graphene.Field(
+        GraphenePartitionKeyConnection,
+        limit=graphene.Argument(graphene.NonNull(graphene.Int)),
+        ascending=graphene.Argument(graphene.NonNull(graphene.Boolean)),
+        cursor=graphene.Argument(graphene.String),
+    )
     partitionKeysByDimension = graphene.Field(
         non_null_list(GrapheneDimensionPartitionKeys),
         startIdx=graphene.Int(),
@@ -1542,6 +1555,39 @@ def resolve_partitionKeysByDimension(
     def resolve_partitionKeys(self, _graphene_info: ResolveInfo) -> Sequence[str]:
         return self.get_partition_keys()
 
+    def resolve_partitionKeyConnection(
+        self,
+        graphene_info: ResolveInfo,
+        limit: int,
+        ascending: bool,
+        cursor: Optional[str] = None,
+    ) -> Optional[GraphenePartitionKeyConnection]:
+        if not self._dynamic_partitions_loader:
+            check.failed("dynamic_partitions_loader must be provided to get partition keys")
+
+        if not self._remote_node.is_partitioned:
+            return None
+
+        partitions_def = self._get_partitions_def()
+        context = PartitionLoadingContext(
+            TemporalContext(
+                effective_dt=get_current_datetime(),
+                last_event_id=graphene_info.context.instance.event_log_storage.get_maximum_record_id(),
+            ),
+            dynamic_partitions_store=self._dynamic_partitions_loader,
+        )
+        conn = partitions_def.get_partition_key_connection(
+            context=context,
+            limit=limit,
+            ascending=ascending,
+            cursor=cursor,
+        )
+        return GraphenePartitionKeyConnection(
+            results=conn.results,
+            cursor=conn.cursor,
+            hasMore=conn.has_more,
+        )
+
     def resolve_partitionDefinition(
         self, _graphene_info: ResolveInfo
     ) -> Optional[GraphenePartitionDefinition]:

python_modules/dagster-graphql/dagster_graphql/schema/partition_keys.py

@@ -4,6 +4,15 @@
 from dagster_graphql.schema.util import non_null_list
 
 
+class GraphenePartitionKeyConnection(graphene.ObjectType):
+    results = non_null_list(graphene.String)
+    cursor = graphene.NonNull(graphene.String)
+    hasMore = graphene.NonNull(graphene.Boolean)
+
+    class Meta:
+        name = "PartitionKeyConnection"
+
+
 class GraphenePartitionKeys(graphene.ObjectType):
     partitionKeys = non_null_list(graphene.String)
 

python_modules/dagster-graphql/dagster_graphql_tests/graphql/test_assets.py

@@ -337,6 +337,19 @@
     }
 """
 
+GET_ASSET_PARTITIONS_CONNECTION = """
+    query AssetNodeQuery($assetKeys: [AssetKeyInput!], $limit: Int!, $ascending: Boolean!, $cursor: String) {
+        assetNodes(assetKeys: $assetKeys) {
+            id
+            partitionKeyConnection(limit: $limit, ascending: $ascending, cursor: $cursor) {
+                results
+                cursor
+                hasMore
+            }
+        }
+    }
+"""
+
 GET_PARTITIONS_BY_DIMENSION = """
     query AssetNodeQuery($assetKeys: [AssetKeyInput!], $startIdx: Int, $endIdx: Int) {
         assetNodes(assetKeys: $assetKeys) {
@@ -1418,6 +1431,70 @@ def test_asset_partitions_in_pipeline(self, graphql_context: WorkspaceRequestCon
         assert asset_node["partitionKeys"][0] == "2021-05-05-01:00"
         assert asset_node["partitionKeys"][1] == "2021-05-05-02:00"
 
+    def test_asset_partition_key_connection(self, graphql_context: WorkspaceRequestContext):
+        result = execute_dagster_graphql(
+            graphql_context,
+            GET_ASSET_PARTITIONS_CONNECTION,
+            variables={
+                "assetKeys": [
+                    {"path": ["not_included_asset"]},
+                    {"path": ["upstream_static_partitioned_asset"]},
+                ],
+                "limit": 2,
+                "ascending": True,
+            },
+        )
+
+        assert result.data and result.data["assetNodes"]
+        assert len(result.data["assetNodes"]) == 2
+        unpartitioned_asset_node = result.data["assetNodes"][0]
+        assert unpartitioned_asset_node["partitionKeyConnection"] is None
+        partitioned_asset_node = result.data["assetNodes"][1]
+        partition_key_conn = partitioned_asset_node["partitionKeyConnection"]
+        assert partition_key_conn["hasMore"]
+        assert len(partition_key_conn["results"]) == 2
+        assert partition_key_conn["results"] == ["a", "b"]
+        cursor = partition_key_conn["cursor"]
+
+        result = execute_dagster_graphql(
+            graphql_context,
+            GET_ASSET_PARTITIONS_CONNECTION,
+            variables={
+                "assetKeys": [
+                    {"path": ["upstream_static_partitioned_asset"]},
+                ],
+                "limit": 10,
+                "ascending": True,
+                "cursor": cursor,
+            },
+        )
+        assert result.data and result.data["assetNodes"]
+        assert len(result.data["assetNodes"]) == 1
+        partitioned_asset_node = result.data["assetNodes"][0]
+        partition_key_conn = partitioned_asset_node["partitionKeyConnection"]
+        assert not partition_key_conn["hasMore"]
+        assert len(partition_key_conn["results"]) == 4
+        assert partition_key_conn["results"] == ["c", "d", "e", "f"]
+
+        result = execute_dagster_graphql(
+            graphql_context,
+            GET_ASSET_PARTITIONS_CONNECTION,
+            variables={
+                "assetKeys": [
+                    {"path": ["upstream_static_partitioned_asset"]},
+                ],
+                "limit": 10,
+                "ascending": False,
+            },
+        )
+        assert result.data and result.data["assetNodes"]
+        assert len(result.data["assetNodes"]) == 1
+        partitioned_asset_node = result.data["assetNodes"][0]
+        partition_key_conn = partitioned_asset_node["partitionKeyConnection"]
+        assert not partition_key_conn["hasMore"]
+        assert len(partition_key_conn["results"]) == 6
+        assert partition_key_conn["results"] == ["f", "e", "d", "c", "b", "a"]
+
     def test_latest_materialization_per_partition(self, graphql_context: WorkspaceRequestContext):
         _create_partitioned_run(graphql_context, "partition_materialization_job", partition_key="c")
         _create_partitioned_run(graphql_context, "partition_materialization_job", partition_key="d")