[ui] Skip dimension matching check when dynamic partition added (#19106)

Fixes #17380. 

Users were seeing the `Attempting to show unified asset health for
assets with dimension of different lengths` error because the
`usePartitionHealthData` query makes a separate request for partitions
per asset, updating the state after each asset's request. When 2 assets
are selected, the first request returns a response including the new
partition, causing the error because the second asset's request has not
returned yet.

Because the `LaunchAssetChoosePartitionsDialog` component already has an
existing check to assert that each root partitions def in the asset
selection is the same, we can skip the internal check. This PR adds a
flag to bypass the "dimensions of different lengths" check.

Adds a test to preserve this behavior.

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

@@ -171,7 +171,7 @@ const LaunchAssetChoosePartitionsDialogBody = ({
       return mergedAssetHealth([]);
     }
     if (target.type === 'job' || assetHealthLoading) {
-      return mergedAssetHealth(assetHealth);
+      return mergedAssetHealth(assetHealth, true);
     }
     return assetHealth.find(itemWithAssetKey(target.anchorAssetKey)) || mergedAssetHealth([]);
   }, [assetHealth, assetHealthLoading, target]);
@@ -425,7 +425,7 @@ const LaunchAssetChoosePartitionsDialogBody = ({
     if (target.type === 'pureWithAnchorAsset') {
       notices.push(
         `Dagster will materialize all partitions downstream of the ` +
-          `selected partitions for the selected assets, using separate runs 
+          `selected partitions for the selected assets, using separate runs
                 ${backfillPolicyVaries ? `and obeying backfill policies.` : `as needed.`}`,
       );
     } else if (backfillPolicyVaries) {

js_modules/dagster-ui/packages/ui-core/src/assets/MultipartitioningSupport.tsx

@@ -35,7 +35,10 @@ the asset health bar you see is a flattened representation of the health of all
 "show per-asset health" button beneath.
 
 */
-export function mergedAssetHealth(assetHealth: PartitionHealthData[]): PartitionHealthDataMerged {
+export function mergedAssetHealth(
+  assetHealth: PartitionHealthData[],
+  skipDimensionLengthsMatchingCheck: boolean = false,
+): PartitionHealthDataMerged {
   if (!assetHealth.length) {
     return {
       dimensions: [],
@@ -56,7 +59,8 @@ export function mergedAssetHealth(assetHealth: PartitionHealthData[]): Partition
       h.dimensions.every(
         (dim, idx) => dim.partitionKeys.length === dimensions[idx]!.partitionKeys.length,
       ),
-    )
+    ) &&
+    !skipDimensionLengthsMatchingCheck
   ) {
     throw new Error(
       'Attempting to show unified asset health for assets with dimension of different lengths',

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

@@ -0,0 +1,153 @@
+import {MockedProvider} from '@apollo/client/testing';
+import {render, screen, waitFor} from '@testing-library/react';
+import userEvent from '@testing-library/user-event';
+import React from 'react';
+import {MemoryRouter} from 'react-router';
+
+import {
+  PartitionDefinitionType,
+  buildAddDynamicPartitionSuccess,
+  buildAssetKey,
+  buildAssetNode,
+  buildDimensionPartitionKeys,
+  buildMultiPartitionStatuses,
+  buildPartitionDefinition,
+} from '../../graphql/types';
+import {CREATE_PARTITION_MUTATION} from '../../partitions/CreatePartitionDialog';
+import {
+  AddDynamicPartitionMutation,
+  AddDynamicPartitionMutationVariables,
+} from '../../partitions/types/CreatePartitionDialog.types';
+import {buildMutationMock, buildQueryMock, getMockResultFn} from '../../testing/mocking';
+import {buildRepoAddress} from '../../workspace/buildRepoAddress';
+import {LaunchAssetChoosePartitionsDialog} from '../LaunchAssetChoosePartitionsDialog';
+import {
+  PartitionHealthQuery,
+  PartitionHealthQueryVariables,
+} from '../types/usePartitionHealthData.types';
+import {PARTITION_HEALTH_QUERY} from '../usePartitionHealthData';
+
+describe('launchAssetChoosePartitionsDialog', () => {
+  it('Adding a dynamic partition when multiple assets selected', async () => {
+    const assetA = buildAsset('asset_a', ['test']);
+    const assetB = buildAsset('asset_b', ['test']);
+
+    const assetAQueryMock = buildQueryMock<PartitionHealthQuery, PartitionHealthQueryVariables>({
+      query: PARTITION_HEALTH_QUERY,
+      variables: {assetKey: {path: ['asset_a']}},
+      data: {assetNodeOrError: assetA},
+    });
+    const assetBQueryMock = buildQueryMock<PartitionHealthQuery, PartitionHealthQueryVariables>({
+      query: PARTITION_HEALTH_QUERY,
+      variables: {assetKey: {path: ['asset_b']}},
+      data: {assetNodeOrError: assetB},
+    });
+    const assetASecondQueryMock = buildQueryMock<
+      PartitionHealthQuery,
+      PartitionHealthQueryVariables
+    >({
+      query: PARTITION_HEALTH_QUERY,
+      variables: {assetKey: {path: ['asset_a']}},
+      data: {assetNodeOrError: buildAsset('asset_a', ['test', 'test2'])},
+    });
+    const assetBSecondQueryMock = buildQueryMock<
+      PartitionHealthQuery,
+      PartitionHealthQueryVariables
+    >({
+      query: PARTITION_HEALTH_QUERY,
+      variables: {assetKey: {path: ['asset_b']}},
+      data: {assetNodeOrError: buildAsset('asset_b', ['test', 'test2'])},
+      delay: 5000,
+    });
+
+    const addPartitionMock = buildMutationMock<
+      AddDynamicPartitionMutation,
+      AddDynamicPartitionMutationVariables
+    >({
+      query: CREATE_PARTITION_MUTATION,
+      variables: {
+        repositorySelector: {repositoryName: 'test', repositoryLocationName: 'test'},
+        partitionsDefName: 'foo',
+        partitionKey: 'test2',
+      },
+      data: {
+        addDynamicPartition: buildAddDynamicPartitionSuccess(),
+      },
+    });
+
+    const assetAQueryMockResult = getMockResultFn(assetAQueryMock);
+    const assetBQueryMockResult = getMockResultFn(assetBQueryMock);
+    const assetASecondQueryMockResult = getMockResultFn(assetASecondQueryMock);
+    const assetBSecondQueryMockResult = getMockResultFn(assetBSecondQueryMock);
+    render(
+      <MemoryRouter>
+        <MockedProvider
+          mocks={[
+            assetAQueryMock,
+            assetBQueryMock,
+            assetASecondQueryMock,
+            assetBSecondQueryMock,
+            addPartitionMock,
+          ]}
+        >
+          <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={[]}
+          />
+        </MockedProvider>
+      </MemoryRouter>,
+    );
+
+    await waitFor(() => {
+      expect(assetAQueryMockResult).toHaveBeenCalled();
+      expect(assetBQueryMockResult).toHaveBeenCalled();
+    });
+
+    const link = await waitFor(() => screen.getByTestId('add-partition-link'));
+    userEvent.click(link);
+    const partitionInput = await waitFor(() => screen.getByTestId('partition-input'));
+    await userEvent.type(partitionInput, 'test2');
+    expect(assetASecondQueryMockResult).not.toHaveBeenCalled();
+    expect(assetBSecondQueryMockResult).not.toHaveBeenCalled();
+    const savePartitionButton = screen.getByTestId('save-partition-button');
+    userEvent.click(savePartitionButton);
+
+    await waitFor(() => {
+      expect(assetASecondQueryMockResult).toHaveBeenCalled();
+    });
+  });
+});
+
+function buildAsset(name: string, dynamicPartitionKeys: string[]) {
+  return buildAssetNode({
+    assetKey: buildAssetKey({path: [name]}),
+    id: `repro_dynamic_in_multipartitions_bug.py.__repository__.["${name}"]`,
+    partitionKeysByDimension: [
+      buildDimensionPartitionKeys({
+        name: 'a',
+        type: PartitionDefinitionType.DYNAMIC,
+        partitionKeys: dynamicPartitionKeys,
+      }),
+      buildDimensionPartitionKeys({
+        name: 'b',
+        type: PartitionDefinitionType.TIME_WINDOW,
+        partitionKeys: ['2024-01-01'],
+      }),
+    ],
+    partitionDefinition: buildPartitionDefinition({
+      name: 'foo',
+    }),
+    assetPartitionStatuses: buildMultiPartitionStatuses({
+      primaryDimensionName: 'b',
+      ranges: [],
+    }),
+  });
+}

js_modules/dagster-ui/packages/ui-core/src/partitions/DimensionRangeWizard.tsx

@@ -91,6 +91,7 @@ export const DimensionRangeWizard = ({
             onClick={() => {
               setShowCreatePartition(true);
             }}
+            data-testid={testId('add-partition-link')}
           >
             <StyledIcon name="add" size={24} />
             <div>Add a partition</div>