Fix reloadWorkspace nested resolver context

js_modules/ui-core/src/nav/useRepositoryLocationReload.tsx

@@ -310,36 +310,40 @@ export const reloadFnForWorkspace = async (client: ApolloClient<any>): Promise<A
   ) {
     return {type: 'error', error: data?.reloadWorkspace, errorLocationId: null};
   }
+  let result;
+  try {
+    result = await client.query<
+      RepositoryLocationStatusQuery,
+      RepositoryLocationStatusQueryVariables
+    >({
+      query: REPOSITORY_LOCATION_STATUS_QUERY,
+      fetchPolicy: 'no-cache',
+    });
+  } catch (e) {
+    return {
+      type: 'error',
+      error: {message: e instanceof ApolloError ? e.message : 'An unexpected error occurred'},
+      errorLocationId: null,
+    };
+  }
+  const workspaceOrError = result.data?.workspaceOrError;
+  if (!workspaceOrError) {
+    return {type: 'error', error: {message: 'Unable to load definitions'}, errorLocationId: null};
+  }
+  if (workspaceOrError.__typename === 'PythonError') {
+    return {type: 'error', error: workspaceOrError, errorLocationId: null};
+  }
   return {
     type: 'finish-mutation-and-start-polling',
-    locationIds: data.reloadWorkspace.locationEntries.map((l) => l.id),
+    locationIds: workspaceOrError.locationEntries.map((l) => l.id),
   };
 };
 
 const RELOAD_WORKSPACE_MUTATION = gql`
   mutation ReloadWorkspaceMutation {
     reloadWorkspace {
-      ... on Workspace {
-        id
-        locationEntries {
-          name
-          id
-          loadStatus
-          locationOrLoadError {
-            ... on RepositoryLocation {
-              id
-              repositories {
-                id
-                name
-                pipelines {
-                  id
-                  name
-                }
-              }
-            }
-            ...PythonErrorFragment
-          }
-        }
+      ... on ReloadWorkspaceSuccess {
+        success
       }
       ...UnauthorizedErrorFragment
       ...PythonErrorFragment

python_modules/dagster-graphql/dagster_graphql/implementation/external.py

@@ -165,11 +165,11 @@ def fetch_workspace(
     )
 
     nodes = [
-        GrapheneWorkspaceLocationEntry(entry)
+        GrapheneWorkspaceLocationEntry(entry, workspace_request_context)
         for entry in workspace_request_context.get_code_location_entries().values()
     ]
 
-    return GrapheneWorkspace(locationEntries=nodes)
+    return GrapheneWorkspace(locationEntries=nodes, workspace_context=workspace_request_context)
 
 
 def fetch_location_statuses(

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

@@ -16,7 +16,7 @@
     LocationStateSubscriber,
 )
 from dagster._core.remote_representation.handle import RepositoryHandle
-from dagster._core.workspace.context import WorkspaceProcessContext
+from dagster._core.workspace.context import BaseWorkspaceRequestContext, WorkspaceProcessContext
 from dagster._core.workspace.workspace import (
     CodeLocationEntry,
     CodeLocationLoadStatus,
@@ -102,8 +102,14 @@ class GrapheneRepositoryLocation(graphene.ObjectType):
     class Meta:
         name = "RepositoryLocation"
 
-    def __init__(self, name: str, location: CodeLocation | None = None):
+    def __init__(
+        self,
+        name: str,
+        location: CodeLocation | None = None,
+        workspace_context: BaseWorkspaceRequestContext | None = None,
+    ):
         self._location = location
+        self._workspace_context = workspace_context
         super().__init__(
             name=name,
         )
@@ -113,12 +119,13 @@ def resolve_id(self, _) -> str:
 
     def get_location(self, graphene_info: ResolveInfo) -> CodeLocation:
         if self._location is None:
-            self._location = graphene_info.context.get_code_location(self.name)
+            context = self._workspace_context or graphene_info.context
+            self._location = context.get_code_location(self.name)
         return self._location
 
     def resolve_repositories(self, graphene_info: ResolveInfo):
         return [
-            GrapheneRepository(repository.handle)
+            GrapheneRepository(repository.handle, workspace_context=self._workspace_context)
             for repository in self.get_location(graphene_info).get_repositories().values()
         ]
 
@@ -226,8 +233,13 @@ class GrapheneWorkspaceLocationEntry(graphene.ObjectType):
     class Meta:
         name = "WorkspaceLocationEntry"
 
-    def __init__(self, location_entry: CodeLocationEntry):
+    def __init__(
+        self,
+        location_entry: CodeLocationEntry,
+        workspace_context: BaseWorkspaceRequestContext | None = None,
+    ):
         self._location_entry = check.inst_param(location_entry, "location_entry", CodeLocationEntry)
+        self._workspace_context = workspace_context
         super().__init__(
             name=self._location_entry.origin.location_name,
             definitionsSource=self._location_entry.definitions_source,
@@ -241,6 +253,7 @@ def resolve_locationOrLoadError(self, _: ResolveInfo):
             return GrapheneRepositoryLocation(
                 self._location_entry.code_location.name,
                 self._location_entry.code_location,
+                workspace_context=self._workspace_context,
             )
 
         error = self._location_entry.load_error
@@ -276,7 +289,8 @@ def resolve_versionKey(self, _) -> str:
         return self._location_entry.version_key
 
     def resolve_permissions(self, graphene_info):
-        permissions = graphene_info.context.permissions_for_location(location_name=self.name)
+        context = self._workspace_context or graphene_info.context
+        permissions = context.permissions_for_location(location_name=self.name)
         return [GraphenePermission(permission, value) for permission, value in permissions.items()]
 
     def resolve_featureFlags(self, graphene_info):
@@ -319,23 +333,29 @@ class Meta:
     def __init__(
         self,
         handle: RepositoryHandle,
+        workspace_context: BaseWorkspaceRequestContext | None = None,
     ):
         # Warning! GrapheneAssetNode contains a GrapheneRepository. Any computation in this
         # __init__ will be done **once per asset**. Ensure that any expensive work is done
         # elsewhere or cached.
         self._handle = handle
+        self._workspace_context = workspace_context
 
         self._batch_loader = None
 
         super().__init__(name=handle.repository_name)
 
+    def get_context(self, graphene_info: ResolveInfo) -> BaseWorkspaceRequestContext:
+        return self._workspace_context or graphene_info.context
+
     def get_repository(self, graphene_info: ResolveInfo) -> RemoteRepository:
-        return graphene_info.context.get_repository(self._handle.to_selector())
+        return self.get_context(graphene_info).get_repository(self._handle.to_selector())
 
     def get_batch_loader(self, graphene_info: ResolveInfo):
         if self._batch_loader is None:
+            context = self.get_context(graphene_info)
             self._batch_loader = RepositoryScopedBatchLoader(
-                graphene_info.context.instance, self.get_repository(graphene_info)
+                context.instance, self.get_repository(graphene_info)
             )
         return self._batch_loader
 
@@ -347,7 +367,10 @@ def resolve_origin(self, _graphene_info: ResolveInfo):
         return GrapheneRepositoryOrigin(origin)
 
     def resolve_location(self, graphene_info: ResolveInfo):
-        return GrapheneRepositoryLocation(self._handle.location_name)
+        return GrapheneRepositoryLocation(
+            self._handle.location_name,
+            workspace_context=self._workspace_context,
+        )
 
     def resolve_schedules(self, graphene_info: ResolveInfo):
         batch_loader = self.get_batch_loader(graphene_info)
@@ -407,7 +430,9 @@ def resolve_usedSolids(self, graphene_info: ResolveInfo):
         return get_solids(self.get_repository(graphene_info))
 
     def resolve_partitionSets(self, graphene_info: ResolveInfo):
-        partition_sets = graphene_info.context.get_partition_sets(self._handle.to_selector())
+        partition_sets = self.get_context(graphene_info).get_partition_sets(
+            self._handle.to_selector()
+        )
         return (GraphenePartitionSet(partition_set) for partition_set in partition_sets)
 
     def resolve_displayMetadata(self, graphene_info: ResolveInfo):
@@ -462,9 +487,9 @@ def resolve_assetNodesConnection(
 
     def resolve_assetManifest(self, graphene_info: ResolveInfo) -> list:
         repository = self.get_repository(graphene_info)
-        base_deployment_asset_graph = graphene_info.context.get_base_deployment_asset_graph(
-            self._handle.to_selector()
-        )
+        base_deployment_asset_graph = self.get_context(
+            graphene_info
+        ).get_base_deployment_asset_graph(self._handle.to_selector())
         asset_graph_differ = (
             AssetGraphDiffer(
                 branch_asset_graph=repository.asset_graph,
@@ -538,6 +563,13 @@ class GrapheneWorkspace(graphene.ObjectType):
     class Meta:
         name = "Workspace"
 
+    def __init__(
+        self,
+        locationEntries,
+        workspace_context: BaseWorkspaceRequestContext | None = None,
+    ):
+        super().__init__(locationEntries=locationEntries)
+
     def resolve_id(self, _graphene_info: ResolveInfo):
         return "Workspace"
 

python_modules/dagster-graphql/dagster_graphql/schema/roots/mutation.py

@@ -35,7 +35,7 @@
     launch_pipeline_reexecution,
     launch_reexecution_from_parent_run,
 )
-from dagster_graphql.implementation.external import fetch_workspace, get_full_remote_job_or_raise
+from dagster_graphql.implementation.external import get_full_remote_job_or_raise
 from dagster_graphql.implementation.fetch_app_managed_components import (
     delete_app_managed_component,
     set_app_managed_component,
@@ -76,7 +76,7 @@
     GrapheneUnauthorizedError,
     GrapheneUnsupportedOperationError,
 )
-from dagster_graphql.schema.external import GrapheneWorkspace, GrapheneWorkspaceLocationEntry
+from dagster_graphql.schema.external import GrapheneWorkspaceLocationEntry
 from dagster_graphql.schema.inputs import (
     GrapheneExecutionParams,
     GrapheneLaunchBackfillParams,
@@ -763,7 +763,8 @@ def mutate(
         # an updated WorkspaceRequestContext for us to use.
         new_context = graphene_info.context.reload_code_location(repositoryLocationName)
         return GrapheneWorkspaceLocationEntry(
-            check.not_none(new_context.get_location_entry(repositoryLocationName))
+            check.not_none(new_context.get_location_entry(repositoryLocationName)),
+            workspace_context=new_context,
         )
 
 
@@ -800,12 +801,21 @@ def mutate(
         )
 
 
+class GrapheneReloadWorkspaceSuccess(graphene.ObjectType):
+    """Output indicating that the workspace was reloaded."""
+
+    success = graphene.NonNull(graphene.Boolean)
+
+    class Meta:
+        name = "ReloadWorkspaceSuccess"
+
+
 class GrapheneReloadWorkspaceMutationResult(graphene.Union):
     """The output from reloading the workspace."""
 
     class Meta:
         types = (
-            GrapheneWorkspace,
+            GrapheneReloadWorkspaceSuccess,
             GrapheneUnauthorizedError,
             GraphenePythonError,
         )
@@ -823,8 +833,8 @@ class Meta:
     @capture_error
     @check_permission(Permissions.RELOAD_WORKSPACE)
     def mutate(self, graphene_info: ResolveInfo):
-        new_context = graphene_info.context.reload_workspace()
-        return fetch_workspace(new_context)
+        graphene_info.context.reload_workspace()
+        return GrapheneReloadWorkspaceSuccess(success=True)
 
 
 class GrapheneAssetWipeSuccess(graphene.ObjectType):