[1/n] Split RootWorkspaceQuery by code location (#22296)

## Summary & Motivation

This PR speeds up workspace fetching by splitting the RootWorkspaceQuery
into multiple parallel fetches by code location. We now rely on the
CodeLocationStatusQuery to determine these locations. To minimize
delays, we cache the code location status query in IndexedDB, enabling
immediate fetches after the initial load. It also enables us to fetch
only changed code locations rather than needing to fetch all locations
when a single one changes.

The WorkspaceProvider is now responsible for refetching code locations
when they change by comparing status updates to detect changes. This
responsibility was previously with useCodeLocationsStatus, but to avoid
a cyclical dependency, it has been moved to the WorkspaceProvider.
useCodeLocationsStatus will continue handling toast updates and
returning status information for downstream consumers, but now it uses
data from the workspace context instead of making direct queries. (In a
future PR we could collapse some of this functionality but I wanted to
minimize the number of code changes in this PR).

The skip parameter has been removed from useCodeLocationsStatus since it
was never set to true in either cloud or OSS environments, both of which
poll for code location updates.

Tests now require mocking more queries to create the desired workspace
context. To simplify this, a buildWorkspaceMocks function has been
added. It accepts location entries as arguments and creates the
necessary code location statuses query mock and workspace query mocks by
location.

Had to update testing-library/react to get tests to pass because the
version we were on had some goofy react act stuff that was fixed.

Next PRs in stack will be:

- [Cloud PR updating workspace mocks /
imports](dagster-io/internal#10126)
- Fetching schedule/sensor state independently in LeftNavItem and
useDaemonStatus (currently they depend on the workspace context but this
doesnt refresh anymore)
- Making OverviewSensors/OverviewSchedules/OverviewJobs etc. use the
root workspace instead of fetching everything independently from scratch
every time.

## How I Tested These Changes

- Ran local host cloud and updated code locations. Tested cached loads
with locations different and with locations the same making sure we
didn't make extra queries if not necessary
- Ran dagster-dev and made sure manual code location reloading still
worked
- Existing indexeddb test still pass covering the changes in that file
- Relying on existing tests using WorkspaceProvider to make sure product
functionality is still intact.

A cloud PR will follow updating workspace mocks there

js_modules/dagster-ui/packages/ui-components/package.json

@@ -66,7 +66,7 @@
     "@storybook/react-webpack5": "^7.2.0",
     "@testing-library/dom": "^10.0.0",
     "@testing-library/jest-dom": "^6.4.2",
-    "@testing-library/react": "^15.0.3",
+    "@testing-library/react": "^16.0.0",
     "@testing-library/user-event": "^14.5.2",
     "@types/babel__core": "^7",
     "@types/babel__preset-env": "^7",

js_modules/dagster-ui/packages/ui-core/package.json

@@ -107,7 +107,7 @@
     "@storybook/react-webpack5": "^7.6.7",
     "@testing-library/dom": "^10.0.0",
     "@testing-library/jest-dom": "^6.4.2",
-    "@testing-library/react": "^15.0.3",
+    "@testing-library/react": "^16.0.0",
     "@testing-library/user-event": "^14.5.2",
     "@types/codemirror": "^5.60.5",
     "@types/color": "^3.0.2",

js_modules/dagster-ui/packages/ui-core/src/assets/__fixtures__/AssetViewDefinition.fixtures.ts

@@ -6,14 +6,9 @@ import {
   buildRepository,
   buildRepositoryLocation,
   buildResourceRequirement,
-  buildWorkspace,
   buildWorkspaceLocationEntry,
 } from '../../graphql/types';
-import {ROOT_WORKSPACE_QUERY} from '../../workspace/WorkspaceContext';
-import {
-  RootWorkspaceQuery,
-  RootWorkspaceQueryVariables,
-} from '../../workspace/types/WorkspaceContext.types';
+import {buildWorkspaceMocks} from '../../workspace/__fixtures__/Workspace.fixtures';
 import {ASSET_VIEW_DEFINITION_QUERY} from '../AssetView';
 import {buildQueryMock} from '../AutoMaterializePolicyPage/__fixtures__/AutoMaterializePolicyPage.fixtures';
 import {
@@ -153,30 +148,19 @@ export const AssetViewDefinitionSDA = buildQueryMock<
   },
 });
 
-export const RootWorkspaceWithOneLocation = buildQueryMock<
-  RootWorkspaceQuery,
-  RootWorkspaceQueryVariables
->({
-  query: ROOT_WORKSPACE_QUERY,
-  variables: {},
-  data: {
-    workspaceOrError: buildWorkspace({
-      locationEntries: [
-        buildWorkspaceLocationEntry({
-          locationOrLoadError: buildRepositoryLocation({
-            repositories: [
-              buildRepository({
-                id: '4d0b1967471d9a4682ccc97d12c1c508d0d9c2e1',
-                name: 'repo',
-                location: buildRepositoryLocation({
-                  id: 'test.py',
-                  name: 'test.py',
-                }),
-              }),
-            ],
+export const RootWorkspaceWithOneLocation = buildWorkspaceMocks([
+  buildWorkspaceLocationEntry({
+    locationOrLoadError: buildRepositoryLocation({
+      repositories: [
+        buildRepository({
+          id: '4d0b1967471d9a4682ccc97d12c1c508d0d9c2e1',
+          name: 'repo',
+          location: buildRepositoryLocation({
+            id: 'test.py',
+            name: 'test.py',
           }),
         }),
       ],
     }),
-  },
-});
+  }),
+]);

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/AssetTables.test.tsx

@@ -4,9 +4,8 @@ import userEvent from '@testing-library/user-event';
 import {MemoryRouter} from 'react-router';
 import {RecoilRoot} from 'recoil';
 
-import {buildWorkspace} from '../../graphql/types';
-import {buildWorkspaceContextMockedResponse} from '../../runs/__fixtures__/RunsFilterInput.fixtures';
 import {WorkspaceProvider} from '../../workspace/WorkspaceContext';
+import {buildWorkspaceMocks} from '../../workspace/__fixtures__/Workspace.fixtures';
 import {AssetsCatalogTable} from '../AssetsCatalogTable';
 import {
   AssetCatalogGroupTableMock,
@@ -17,7 +16,7 @@ import {
   SingleAssetQueryTrafficDashboard,
 } from '../__fixtures__/AssetTables.fixtures';
 
-const workspaceMock = buildWorkspaceContextMockedResponse(buildWorkspace({}));
+const workspaceMocks = buildWorkspaceMocks([]);
 
 const MOCKS = [
   AssetCatalogTableMock,
@@ -26,7 +25,7 @@ const MOCKS = [
   SingleAssetQueryMaterializedWithLatestRun,
   SingleAssetQueryMaterializedStaleAndLate,
   SingleAssetQueryLastRunFailed,
-  workspaceMock,
+  ...workspaceMocks,
 ];
 
 // This file must be mocked because Jest can't handle `import.meta.url`.

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/AssetView.test.tsx

@@ -59,7 +59,7 @@ describe('AssetView', () => {
       <RecoilRoot>
         <MockedProvider
           mocks={[
-            RootWorkspaceWithOneLocation,
+            ...RootWorkspaceWithOneLocation,
             AssetViewDefinitionSDA,
             AssetViewDefinitionNonSDA,
             AssetViewDefinitionSourceAsset,

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetChoosePartitionsDialog.test.tsx

@@ -12,15 +12,15 @@ import {
   buildDimensionPartitionKeys,
   buildMultiPartitionStatuses,
   buildPartitionDefinition,
-  buildWorkspace,
 } from '../../graphql/types';
 import {CREATE_PARTITION_MUTATION} from '../../partitions/CreatePartitionDialog';
 import {
   AddDynamicPartitionMutation,
   AddDynamicPartitionMutationVariables,
 } from '../../partitions/types/CreatePartitionDialog.types';
-import {buildWorkspaceContextMockedResponse} from '../../runs/__fixtures__/RunsFilterInput.fixtures';
 import {buildMutationMock, buildQueryMock, getMockResultFn} from '../../testing/mocking';
+import {WorkspaceProvider} from '../../workspace/WorkspaceContext';
+import {buildWorkspaceMocks} from '../../workspace/__fixtures__/Workspace.fixtures';
 import {buildRepoAddress} from '../../workspace/buildRepoAddress';
 import {LaunchAssetChoosePartitionsDialog} from '../LaunchAssetChoosePartitionsDialog';
 import {
@@ -29,7 +29,7 @@ import {
 } from '../types/usePartitionHealthData.types';
 import {PARTITION_HEALTH_QUERY} from '../usePartitionHealthData';
 
-const workspaceMock = buildWorkspaceContextMockedResponse(buildWorkspace({}));
+const workspaceMocks = buildWorkspaceMocks([]);
 
 describe('launchAssetChoosePartitionsDialog', () => {
   it('Adding a dynamic partition when multiple assets selected', async () => {
@@ -92,21 +92,23 @@ describe('launchAssetChoosePartitionsDialog', () => {
             assetASecondQueryMock,
             assetBSecondQueryMock,
             addPartitionMock,
-            workspaceMock,
+            ...workspaceMocks,
           ]}
         >
-          <LaunchAssetChoosePartitionsDialog
-            open={true}
-            setOpen={(_open: boolean) => {}}
-            repoAddress={buildRepoAddress('test', 'test')}
-            target={{
-              jobName: '__ASSET_JOB_0',
-              partitionSetName: '__ASSET_JOB_0_partition_set',
-              type: 'job',
-            }}
-            assets={[assetA, assetB]}
-            upstreamAssetKeys={[]}
-          />
+          <WorkspaceProvider>
+            <LaunchAssetChoosePartitionsDialog
+              open={true}
+              setOpen={(_open: boolean) => {}}
+              repoAddress={buildRepoAddress('test', 'test')}
+              target={{
+                jobName: '__ASSET_JOB_0',
+                partitionSetName: '__ASSET_JOB_0_partition_set',
+                type: 'job',
+              }}
+              assets={[assetA, assetB]}
+              upstreamAssetKeys={[]}
+            />
+          </WorkspaceProvider>
         </MockedProvider>
       </MemoryRouter>,
     );

js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetExecutionButton.test.tsx

@@ -6,12 +6,11 @@ import userEvent from '@testing-library/user-event';
 import {CustomAlertProvider} from '../../app/CustomAlertProvider';
 import {CustomConfirmationProvider} from '../../app/CustomConfirmationProvider';
 import {displayNameForAssetKey} from '../../asset-graph/Utils';
-import {buildWorkspace} from '../../graphql/types';
 import {LaunchPartitionBackfillMutation} from '../../instance/backfill/types/BackfillUtils.types';
-import {buildWorkspaceContextMockedResponse} from '../../runs/__fixtures__/RunsFilterInput.fixtures';
 import {LaunchPipelineExecutionMutation} from '../../runs/types/RunUtils.types';
 import {TestProvider} from '../../testing/TestProvider';
 import * as WorkspaceContext from '../../workspace/WorkspaceContext';
+import {buildWorkspaceMocks} from '../../workspace/__fixtures__/Workspace.fixtures';
 import {ADDITIONAL_REQUIRED_KEYS_WARNING} from '../AssetDefinedInMultipleReposNotice';
 import {
   AssetsInScope,
@@ -46,7 +45,7 @@ import {
 } from '../__fixtures__/LaunchAssetExecutionButton.fixtures';
 import {asAssetKeyInput} from '../asInput';
 
-const workspaceMock = buildWorkspaceContextMockedResponse(buildWorkspace({}));
+const workspaceMocks = buildWorkspaceMocks([]);
 
 // This file must be mocked because Jest can't handle `import.meta.url`.
 jest.mock('../../graph/asyncGraphLayout', () => ({}));
@@ -653,7 +652,7 @@ function renderButton({
     }),
     buildLaunchAssetLoaderMock([MULTI_ASSET_OUT_1.assetKey, MULTI_ASSET_OUT_2.assetKey]),
     buildLaunchAssetLoaderMock(assetKeys),
-    workspaceMock,
+    ...workspaceMocks,
     ...(launchMock ? [launchMock] : []),
   ];
 

js_modules/dagster-ui/packages/ui-core/src/instance/DeploymentStatusProvider.tsx

@@ -24,7 +24,7 @@ interface Props {
 export const DeploymentStatusProvider = (props: Props) => {
   const {children, include} = props;
 
-  const codeLocations = useCodeLocationsStatus(!include.has('code-locations'));
+  const codeLocations = useCodeLocationsStatus();
   const daemons = useDaemonStatus(!include.has('daemons'));
 
   const value = React.useMemo(() => ({codeLocations, daemons}), [daemons, codeLocations]);

js_modules/dagster-ui/packages/ui-core/src/instance/backfill/__tests__/BackfillTable.test.tsx

@@ -1,6 +1,5 @@
 import {MockedProvider} from '@apollo/client/testing';
-import {render, screen} from '@testing-library/react';
-import userEvent from '@testing-library/user-event';
+import {act, render, screen, waitFor} from '@testing-library/react';
 import {MemoryRouter} from 'react-router-dom';
 
 import * as Alerting from '../../../app/CustomAlertProvider';
@@ -28,7 +27,11 @@ describe('BackfillTable', () => {
 
     expect(screen.getByRole('table')).toBeVisible();
     const statusLabel = await screen.findByText('Failed');
-    await userEvent.click(statusLabel);
-    expect(Alerting.showCustomAlert).toHaveBeenCalled();
+    act(() => {
+      statusLabel.click();
+    });
+    await waitFor(() => {
+      expect(Alerting.showCustomAlert).toHaveBeenCalled();
+    });
   });
 });

js_modules/dagster-ui/packages/ui-core/src/instance/flattenCodeLocationRows.tsx

@@ -1,5 +1,5 @@
 import {CodeLocationRowType} from '../workspace/VirtualizedCodeLocationRow';
-import {WorkspaceLocationNodeFragment} from '../workspace/types/WorkspaceContext.types';
+import {WorkspaceLocationNodeFragment} from '../workspace/types/WorkspaceQueries.types';
 
 const flatten = (locationEntries: WorkspaceLocationNodeFragment[]) => {
   // Consider each loaded repository to be a "code location".