add pools to graphs

dagster-io · Jan 21, 2025 · 529046e · 529046e
1 parent eca2527
commit 529046e
Show file tree

Hide file tree

Showing 4 changed files with 22 additions and 5 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql b/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
diff --git a/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts b/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/asset_graph.py b/python_modules/dagster-graphql/dagster_graphql/schema/asset_graph.py
@@ -241,7 +241,6 @@ class GrapheneAssetNode(graphene.ObjectType):
     backfillPolicy = graphene.Field(GrapheneBackfillPolicy)
     changedReasons = graphene.Field(non_null_list(GrapheneAssetChangedReason))
     computeKind = graphene.String()
-    pools = non_null_list(graphene.String)
     configField = graphene.Field(GrapheneConfigTypeField)
     dataVersion = graphene.Field(graphene.String(), partition=graphene.String())
     dataVersionByPartition = graphene.Field(
@@ -288,6 +287,7 @@ class GrapheneAssetNode(graphene.ObjectType):
         startIdx=graphene.Int(),
         endIdx=graphene.Int(),
     )
+    pools = non_null_list(graphene.String)
     repository = graphene.NonNull(lambda: external.GrapheneRepository)
     required_resources = non_null_list(GrapheneResourceRequirement)
     staleStatus = graphene.Field(GrapheneAssetStaleStatus, partition=graphene.String())
@@ -681,9 +681,6 @@ def _get_config_type(key: str):
             field_snap=node_def_snap.config_field_snap,
         )
 
-    def resolve_pools(self, _graphene_info: ResolveInfo) -> Sequence[str]:
-        return sorted([pool for pool in self._asset_node_snap.pools or set()])
-
     def resolve_computeKind(self, _graphene_info: ResolveInfo) -> Optional[str]:
         return self._asset_node_snap.compute_kind
 
@@ -1219,6 +1216,9 @@ def resolve_partitionDefinition(
             return GraphenePartitionDefinition(partitions_snap)
         return None
 
+    def resolve_pools(self, _graphene_info: ResolveInfo) -> Sequence[str]:
+        return sorted([pool for pool in self._asset_node_snap.pools or set()])
+
     def resolve_repository(self, graphene_info: ResolveInfo) -> "GrapheneRepository":
         return external.GrapheneRepository(self._repository_handle)
 

diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/solids.py b/python_modules/dagster-graphql/dagster_graphql/schema/solids.py
@@ -374,6 +374,7 @@ class GrapheneISolidDefinition(graphene.Interface):
     input_definitions = non_null_list(GrapheneInputDefinition)
     output_definitions = non_null_list(GrapheneOutputDefinition)
     assetNodes = non_null_list("dagster_graphql.schema.asset_graph.GrapheneAssetNode")
+    pools = non_null_list(graphene.String)
 
     class Meta:
         name = "ISolidDefinition"
@@ -454,6 +455,13 @@ def resolve_assetNodes(self, graphene_info: ResolveInfo) -> Sequence["GrapheneAs
                 for remote_node in remote_nodes
             ]
 
+    def resolve_pools(self, _graphene_info) -> Sequence[str]:
+        if isinstance(self._solid_def_snap, OpDefSnap):
+            return [self._solid_def_snap.pool] if self._solid_def_snap.pool else []
+        if isinstance(self._solid_def_snap, GraphDefSnap):
+            return list(self._solid_def_snap.pools)
+        return []
+
 
 class GrapheneSolidDefinition(graphene.ObjectType, ISolidDefinitionMixin):
     config_field = graphene.Field(GrapheneConfigTypeField)