From 4d0c69c6bc01eea87f780fc7b5a637e29ce19798 Mon Sep 17 00:00:00 2001 From: Claire Lin Date: Thu, 25 Jan 2024 15:27:53 -0800 Subject: [PATCH] [ui] Skip dimension matching check when dynamic partition added (#19106) Fixes https://github.com/dagster-io/dagster/issues/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. --- .../LaunchAssetChoosePartitionsDialog.tsx | 4 +- .../src/assets/MultipartitioningSupport.tsx | 8 +- ...LaunchAssetChoosePartitionsDialog.test.tsx | 153 ++++++++++++++++++ .../src/partitions/DimensionRangeWizard.tsx | 1 + 4 files changed, 162 insertions(+), 4 deletions(-) create mode 100644 js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetChoosePartitionsDialog.test.tsx diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/LaunchAssetChoosePartitionsDialog.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/LaunchAssetChoosePartitionsDialog.tsx index 7e1e6cc6503f9..52ff2449ae0e4 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/assets/LaunchAssetChoosePartitionsDialog.tsx +++ b/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) { diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/MultipartitioningSupport.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/MultipartitioningSupport.tsx index b22e6497b4c77..96c691c750213 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/assets/MultipartitioningSupport.tsx +++ b/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', diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetChoosePartitionsDialog.test.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/__tests__/LaunchAssetChoosePartitionsDialog.test.tsx new file mode 100644 index 0000000000000..4bd3b43ba4d10 --- /dev/null +++ b/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({ + query: PARTITION_HEALTH_QUERY, + variables: {assetKey: {path: ['asset_a']}}, + data: {assetNodeOrError: assetA}, + }); + const assetBQueryMock = buildQueryMock({ + 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( + + + {}} + repoAddress={buildRepoAddress('test', 'test')} + target={{ + jobName: '__ASSET_JOB_0', + partitionSetName: '__ASSET_JOB_0_partition_set', + type: 'job', + }} + assets={[assetA, assetB]} + upstreamAssetKeys={[]} + /> + + , + ); + + 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: [], + }), + }); +} diff --git a/js_modules/dagster-ui/packages/ui-core/src/partitions/DimensionRangeWizard.tsx b/js_modules/dagster-ui/packages/ui-core/src/partitions/DimensionRangeWizard.tsx index 2f91a22a5ca62..bb9080a688e88 100644 --- a/js_modules/dagster-ui/packages/ui-core/src/partitions/DimensionRangeWizard.tsx +++ b/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')} >
Add a partition