From 4ee7bb26e6e0aea0038195c74dac4c3534498ae5 Mon Sep 17 00:00:00 2001
From: prha <philrha@gmail.com>
Date: Tue, 7 Jan 2025 20:19:37 -0800
Subject: [PATCH] add pools to graphs

---
 .../packages/ui-core/src/graphql/schema.graphql           | 3 +++
 .../dagster-ui/packages/ui-core/src/graphql/types.ts      | 6 ++++++
 .../dagster-graphql/dagster_graphql/schema/solids.py      | 8 ++++++++
 3 files changed, 17 insertions(+)

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
index 1be5771d58b50..c9685c7341aa6 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
+++ b/js_modules/dagster-ui/packages/ui-core/src/graphql/schema.graphql
@@ -3098,6 +3098,7 @@ type CompositeSolidDefinition implements ISolidDefinition & SolidContainer {
   inputDefinitions: [InputDefinition!]!
   outputDefinitions: [OutputDefinition!]!
   assetNodes: [AssetNode!]!
+  pools: [String!]!
   id: ID!
   solids: [Solid!]!
   solidHandle(handleID: String!): SolidHandle
@@ -3133,6 +3134,7 @@ interface ISolidDefinition {
   inputDefinitions: [InputDefinition!]!
   outputDefinitions: [OutputDefinition!]!
   assetNodes: [AssetNode!]!
+  pools: [String!]!
 }
 
 type Output {
@@ -3183,6 +3185,7 @@ type SolidDefinition implements ISolidDefinition {
   inputDefinitions: [InputDefinition!]!
   outputDefinitions: [OutputDefinition!]!
   assetNodes: [AssetNode!]!
+  pools: [String!]!
   configField: ConfigTypeField
   requiredResources: [ResourceRequirement!]!
   pool: String
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
index b70227407e8f2..52c3cf5d643df 100644
--- a/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
+++ b/js_modules/dagster-ui/packages/ui-core/src/graphql/types.ts
@@ -824,6 +824,7 @@ export type CompositeSolidDefinition = ISolidDefinition &
     name: Scalars['String']['output'];
     outputDefinitions: Array<OutputDefinition>;
     outputMappings: Array<OutputMapping>;
+    pools: Array<Scalars['String']['output']>;
     solidHandle: Maybe<SolidHandle>;
     solidHandles: Array<SolidHandle>;
     solids: Array<Solid>;
@@ -1839,6 +1840,7 @@ export type ISolidDefinition = {
   metadata: Array<MetadataItemDefinition>;
   name: Scalars['String']['output'];
   outputDefinitions: Array<OutputDefinition>;
+  pools: Array<Scalars['String']['output']>;
 };
 
 export type Input = {
@@ -5194,6 +5196,7 @@ export type SolidDefinition = ISolidDefinition & {
   name: Scalars['String']['output'];
   outputDefinitions: Array<OutputDefinition>;
   pool: Maybe<Scalars['String']['output']>;
+  pools: Array<Scalars['String']['output']>;
   requiredResources: Array<ResourceRequirement>;
 };
 
@@ -7308,6 +7311,7 @@ export const buildCompositeSolidDefinition = (
         : [],
     outputMappings:
       overrides && overrides.hasOwnProperty('outputMappings') ? overrides.outputMappings! : [],
+    pools: overrides && overrides.hasOwnProperty('pools') ? overrides.pools! : [],
     solidHandle:
       overrides && overrides.hasOwnProperty('solidHandle')
         ? overrides.solidHandle!
@@ -8828,6 +8832,7 @@ export const buildISolidDefinition = (
       overrides && overrides.hasOwnProperty('outputDefinitions')
         ? overrides.outputDefinitions!
         : [],
+    pools: overrides && overrides.hasOwnProperty('pools') ? overrides.pools! : [],
   };
 };
 
@@ -14463,6 +14468,7 @@ export const buildSolidDefinition = (
         ? overrides.outputDefinitions!
         : [],
     pool: overrides && overrides.hasOwnProperty('pool') ? overrides.pool! : 'voluptates',
+    pools: overrides && overrides.hasOwnProperty('pools') ? overrides.pools! : [],
     requiredResources:
       overrides && overrides.hasOwnProperty('requiredResources')
         ? overrides.requiredResources!
diff --git a/python_modules/dagster-graphql/dagster_graphql/schema/solids.py b/python_modules/dagster-graphql/dagster_graphql/schema/solids.py
index d7c79d27818a4..faa67b6d6a636 100644
--- 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)