From 122d46eebc1c7f64bb465bf0d1b6a29b7ccdc5ce Mon Sep 17 00:00:00 2001 From: Marco polo Date: Wed, 1 Mar 2023 11:42:21 -0500 Subject: [PATCH] [UI] Dynamic partition creation and selection (#12615) ### Summary & Motivation Adds the ability to create dynamic partitions Screen Shot 2023-03-01 at 7 23 20 AM Screen Shot 2023-03-01 at 7 24 22 AM Screen Shot 2023-03-01 at 7 23 24 AM ### How I Tested These Changes I tested this with `dagit dev` and also in `storybook`. I need to spend some time digging through this some more to setup tests --------- Co-authored-by: bengotow --- .../core/src/assets/AssetPartitionList.tsx | 2 +- .../LaunchAssetChoosePartitionsDialog.tsx | 30 +- .../src/assets/LaunchAssetExecutionButton.tsx | 20 + .../src/assets/usePartitionHealthData.tsx | 23 +- .../src/partitions/DimensionRangeWizard.tsx | 371 +++++++++++++++++- .../types/DimensionRangeWizard.types.ts | 18 + .../ui/src/components/TagSelector.tsx | 30 +- 7 files changed, 455 insertions(+), 39 deletions(-) create mode 100644 js_modules/dagit/packages/core/src/partitions/types/DimensionRangeWizard.types.ts diff --git a/js_modules/dagit/packages/core/src/assets/AssetPartitionList.tsx b/js_modules/dagit/packages/core/src/assets/AssetPartitionList.tsx index 3c3ffdc2c200..0472bd7d9e34 100644 --- a/js_modules/dagit/packages/core/src/assets/AssetPartitionList.tsx +++ b/js_modules/dagit/packages/core/src/assets/AssetPartitionList.tsx @@ -106,7 +106,7 @@ export const AssetPartitionList: React.FC = ({ ); }; -const StateDot = ({state}: {state: PartitionState}) => ( +export const StateDot = ({state}: {state: PartitionState}) => (
Promise; } export const LaunchAssetChoosePartitionsDialog: React.FC = (props) => { - const title = `Launch runs to materialize ${ + const displayName = props.assets.length > 1 ? `${props.assets.length} assets` - : displayNameForAssetKey(props.assets[0].assetKey) - }`; + : displayNameForAssetKey(props.assets[0].assetKey); + + const title = `Launch runs to materialize ${displayName}`; return ( = ({ repoAddress, target, upstreamAssetKeys, + refetch: _refetch, }) => { const partitionedAssets = assets.filter((a) => !!a.partitionDefinition); @@ -130,7 +133,19 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC = ({ const [previewCount, setPreviewCount] = React.useState(0); const morePreviewsCount = partitionedAssets.length - previewCount; - const assetHealth = usePartitionHealthData(partitionedAssets.map((a) => a.assetKey)); + const [lastRefresh, setLastRefresh] = React.useState(Date.now()); + + const refetch = async () => { + await _refetch?.(); + setLastRefresh(Date.now()); + }; + + const assetHealth = usePartitionHealthData( + partitionedAssets.map((a) => a.assetKey), + lastRefresh.toString(), + 'immediate', + ); + const assetHealthLoading = assetHealth.length === 0; const displayedHealth = React.useMemo(() => { @@ -147,6 +162,7 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC = ({ const knownDimensions = partitionedAssets[0].partitionDefinition?.dimensionTypes || []; const [missingOnly, setMissingOnly] = React.useState(true); + const [selections, setSelections] = usePartitionDimensionSelections({ knownDimensionNames: knownDimensions.map((d) => d.name), modifyQueryString: false, @@ -364,6 +380,7 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC = ({ selections.length === 2 ? selections[1 - idx].selectedRanges : undefined, ), }} + isDynamic={displayedPartitionDefinition?.type === PartitionDefinitionType.DYNAMIC} selected={range.selectedKeys} setSelected={(selectedKeys) => setSelections( @@ -372,6 +389,9 @@ const LaunchAssetChoosePartitionsDialogBody: React.FC = ({ ), ) } + partitionDefinitionName={displayedPartitionDefinition?.name} + repoAddress={repoAddress} + refetch={refetch} /> ))} diff --git a/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx b/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx index 9587d4b00009..71e66b1dc7df 100644 --- a/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx +++ b/js_modules/dagit/packages/core/src/assets/LaunchAssetExecutionButton.tsx @@ -317,6 +317,26 @@ export const useMaterializationAction = (preferredJobName?: string) => { target={state.target} open={true} setOpen={() => setState({type: 'none'})} + refetch={async () => { + const result = await client.query< + LaunchAssetLoaderQuery, + LaunchAssetLoaderQueryVariables + >({ + query: LAUNCH_ASSET_LOADER_QUERY, + variables: {assetKeys: state.assets.map(({assetKey}) => ({path: assetKey.path}))}, + }); + const assets = result.data.assetNodes; + const next = await stateForLaunchingAssets(client, assets, false, preferredJobName); + if (next.type === 'error') { + showCustomAlert({ + title: 'Unable to Materialize', + body: next.error, + }); + setState({type: 'none'}); + return; + } + setState(next); + }} /> ); } diff --git a/js_modules/dagit/packages/core/src/assets/usePartitionHealthData.tsx b/js_modules/dagit/packages/core/src/assets/usePartitionHealthData.tsx index 69feba96dea8..2be17fae7662 100644 --- a/js_modules/dagit/packages/core/src/assets/usePartitionHealthData.tsx +++ b/js_modules/dagit/packages/core/src/assets/usePartitionHealthData.tsx @@ -436,17 +436,18 @@ function rangesCoverAll(ranges: Range[], keyCount: number) { // a sign that we should invalidate and reload previously loaded health stats. We don't // clear them immediately to avoid an empty state. // -export function usePartitionHealthData(assetKeys: AssetKey[], assetLastMaterializedAt = '') { +export function usePartitionHealthData( + assetKeys: AssetKey[], + cacheKey = '', + cacheClearStrategy: 'immediate' | 'background' = 'background', +) { const [result, setResult] = React.useState<(PartitionHealthData & {fetchedAt: string})[]>([]); const client = useApolloClient(); const assetKeyJSONs = assetKeys.map((k) => JSON.stringify(k)); const assetKeyJSON = JSON.stringify(assetKeyJSONs); const missingKeyJSON = assetKeyJSONs.find( - (k) => - !result.some( - (r) => JSON.stringify(r.assetKey) === k && r.fetchedAt === assetLastMaterializedAt, - ), + (k) => !result.some((r) => JSON.stringify(r.assetKey) === k && r.fetchedAt === cacheKey), ); React.useMemo(() => { @@ -465,16 +466,20 @@ export function usePartitionHealthData(assetKeys: AssetKey[], assetLastMateriali const loaded = buildPartitionHealthData(data, loadKey); setResult((result) => [ ...result.filter((r) => !isEqual(r.assetKey, loadKey)), - {...loaded, fetchedAt: assetLastMaterializedAt}, + {...loaded, fetchedAt: cacheKey}, ]); }; run(); - }, [client, missingKeyJSON, assetLastMaterializedAt]); + }, [client, missingKeyJSON, cacheKey]); return React.useMemo(() => { const assetKeyJSONs = JSON.parse(assetKeyJSON); - return result.filter((r) => assetKeyJSONs.includes(JSON.stringify(r.assetKey))); - }, [assetKeyJSON, result]); + return result.filter( + (r) => + assetKeyJSONs.includes(JSON.stringify(r.assetKey)) && + (r.fetchedAt === cacheKey || cacheClearStrategy === 'background'), + ); + }, [assetKeyJSON, result, cacheKey, cacheClearStrategy]); } export const PARTITION_HEALTH_QUERY = gql` diff --git a/js_modules/dagit/packages/core/src/partitions/DimensionRangeWizard.tsx b/js_modules/dagit/packages/core/src/partitions/DimensionRangeWizard.tsx index 86069cbbe911..f1ab25e6b1df 100644 --- a/js_modules/dagit/packages/core/src/partitions/DimensionRangeWizard.tsx +++ b/js_modules/dagit/packages/core/src/partitions/DimensionRangeWizard.tsx @@ -1,31 +1,83 @@ -import {Box, Button} from '@dagster-io/ui'; +import {gql, useMutation} from '@apollo/client'; +import { + Box, + Button, + Checkbox, + Colors, + Dialog, + DialogBody, + DialogFooter, + Icon, + Menu, + MenuDivider, + MenuItem, + Mono, + TagSelector, + TextInput, +} from '@dagster-io/ui'; import * as React from 'react'; +import styled from 'styled-components/macro'; +import {showCustomAlert} from '../app/CustomAlertProvider'; +import {PythonErrorInfo} from '../app/PythonErrorInfo'; +import {StateDot} from '../assets/AssetPartitionList'; import {isTimeseriesPartition} from '../assets/MultipartitioningSupport'; +import {partitionStateAtIndex, Range} from '../assets/usePartitionHealthData'; +import {repoAddressToSelector} from '../workspace/repoAddressToSelector'; +import {RepoAddress} from '../workspace/types'; import {DimensionRangeInput} from './DimensionRangeInput'; import {PartitionStatusHealthSource, PartitionStatus} from './PartitionStatus'; +import { + AddDynamicPartitionMutation, + AddDynamicPartitionMutationVariables, +} from './types/DimensionRangeWizard.types'; export const DimensionRangeWizard: React.FC<{ selected: string[]; setSelected: (selected: string[]) => void; partitionKeys: string[]; health: PartitionStatusHealthSource; -}> = ({selected, setSelected, partitionKeys, health}) => { + isDynamic?: boolean; + partitionDefinitionName?: string | null; + repoAddress?: RepoAddress; + refetch?: () => void; +}> = ({ + selected, + setSelected, + partitionKeys, + health, + isDynamic = false, + partitionDefinitionName, + repoAddress, + refetch, +}) => { const isTimeseries = isTimeseriesPartition(partitionKeys[0]); + const [showCreatePartition, setShowCreatePartition] = React.useState(false); + return ( <> - + {isDynamic ? ( + + ) : ( + + )} - {isTimeseries && ( + {isTimeseries && !isDynamic && ( @@ -35,14 +87,301 @@ export const DimensionRangeWizard: React.FC<{ - + {isDynamic ? ( + { + setShowCreatePartition(true); + }} + > + +
Create a partition
+ + ) : ( + + )} + {repoAddress && ( + { + setShowCreatePartition(false); + }} + refetch={refetch} + /> + )} + + ); +}; + +const DynamicPartitionSelector: React.FC<{ + allPartitions: string[]; + selectedPartitions: string[]; + setSelectedPartitions: (tags: string[]) => void; + health: PartitionStatusHealthSource; + setShowCreatePartition: (show: boolean) => void; +}> = ({ + allPartitions, + selectedPartitions, + setSelectedPartitions, + setShowCreatePartition, + health, +}) => { + const isAllSelected = + allPartitions.length === selectedPartitions.length && allPartitions.length > 0; + + const statusForPartitionKey = (partitionKey: string) => { + const index = allPartitions.indexOf(partitionKey); + if ('ranges' in health) { + return partitionStateAtIndex(health.ranges as Range[], index); + } else { + return health.partitionStateForKey(partitionKey, index); + } + }; + + return ( + <> + { + return ( + + ); + }} + renderDropdown={(dropdown, {width}) => { + const toggleAll = () => { + if (isAllSelected) { + setSelectedPartitions([]); + } else { + setSelectedPartitions(allPartitions); + } + }; + return ( + + + + + + Create partition + + } + onClick={() => { + setShowCreatePartition(true); + }} + /> + + + {allPartitions.length ? ( + <> + + {dropdown} + + ) : ( +
+ No matching partitions found +
+ )} + + + ); + }} + renderTagList={(tags) => { + if (tags.length > 4) { + return {tags.length} partitions selected; + } + return tags; + }} + /> ); }; + +const StyledIcon = styled(Icon)` + font-weight: 500; +`; + +const LinkText = styled(Box)` + color: ${Colors.Link}; + cursor: pointer; + &:hover { + text-decoration: underline; + } + > * { + height: 24px; + align-content: center; + line-height: 24px; + } +`; +const CreatePartitionDialog = ({ + isOpen, + partitionDefinitionName, + close, + repoAddress, + refetch, +}: { + isOpen: boolean; + partitionDefinitionName?: string | null; + close: () => void; + repoAddress: RepoAddress; + refetch?: () => void; +}) => { + const [partitionName, setPartitionName] = React.useState(''); + + const [createPartition] = useMutation< + AddDynamicPartitionMutation, + AddDynamicPartitionMutationVariables + >(CREATE_PARTITION_MUTATION); + + const handleSave = async () => { + const result = await createPartition({ + variables: { + repositorySelector: repoAddressToSelector(repoAddress), + partitionsDefName: partitionDefinitionName || '', + partitionKey: partitionName, + }, + }); + + const data = result.data?.addDynamicPartition; + switch (data?.__typename) { + case 'PythonError': { + showCustomAlert({ + title: 'Could not create environment variable', + body: , + }); + break; + } + case 'DuplicateDynamicPartitionError': { + showCustomAlert({ + title: 'Could not create partition', + body: 'A partition this name already exists.', + }); + break; + } + case 'UnauthorizedError': { + showCustomAlert({ + title: 'Could not create partition', + body: 'You do not have permission to do this.', + }); + break; + } + case 'AddDynamicPartitionSuccess': { + refetch?.(); + close(); + break; + } + default: { + showCustomAlert({ + title: 'Could not create partition', + body: 'An unknown error occurred.', + }); + break; + } + } + }; + return ( + + +
+ Create a partition + {partitionDefinitionName ? ( + <> + {' '} + for {partitionDefinitionName} + + ) : ( + '' + )} +
+ + } + > + + +
Partition name
+ setPartitionName(e.target.value)} + /> + + + + + + + + ); +}; + +export const CREATE_PARTITION_MUTATION = gql` + mutation AddDynamicPartitionMutation( + $partitionsDefName: String! + $partitionKey: String! + $repositorySelector: RepositorySelector! + ) { + addDynamicPartition( + partitionsDefName: $partitionsDefName + partitionKey: $partitionKey + repositorySelector: $repositorySelector + ) { + __typename + ... on AddDynamicPartitionSuccess { + partitionsDefName + partitionKey + } + ... on PythonError { + message + stack + } + ... on UnauthorizedError { + message + } + } + } +`; diff --git a/js_modules/dagit/packages/core/src/partitions/types/DimensionRangeWizard.types.ts b/js_modules/dagit/packages/core/src/partitions/types/DimensionRangeWizard.types.ts new file mode 100644 index 000000000000..9689f550f559 --- /dev/null +++ b/js_modules/dagit/packages/core/src/partitions/types/DimensionRangeWizard.types.ts @@ -0,0 +1,18 @@ +// Generated GraphQL types, do not edit manually. + +import * as Types from '../../graphql/types'; + +export type AddDynamicPartitionMutationVariables = Types.Exact<{ + partitionsDefName: Types.Scalars['String']; + partitionKey: Types.Scalars['String']; + repositorySelector: Types.RepositorySelector; +}>; + +export type AddDynamicPartitionMutation = { + __typename: 'DagitMutation'; + addDynamicPartition: + | {__typename: 'AddDynamicPartitionSuccess'; partitionsDefName: string; partitionKey: string} + | {__typename: 'DuplicateDynamicPartitionError'} + | {__typename: 'PythonError'; message: string; stack: Array} + | {__typename: 'UnauthorizedError'; message: string}; +}; diff --git a/js_modules/dagit/packages/ui/src/components/TagSelector.tsx b/js_modules/dagit/packages/ui/src/components/TagSelector.tsx index bd5879e5b230..a74091ca51a6 100644 --- a/js_modules/dagit/packages/ui/src/components/TagSelector.tsx +++ b/js_modules/dagit/packages/ui/src/components/TagSelector.tsx @@ -9,6 +9,7 @@ import {MenuItem, Menu} from './Menu'; import {Popover} from './Popover'; import {Tag} from './Tag'; import {TextInputStyles} from './TextInput'; +import {useViewport} from './useViewport'; type TagProps = { remove: (ev: React.SyntheticEvent) => void; @@ -17,21 +18,24 @@ type DropdownItemProps = { toggle: () => void; selected: boolean; }; +type DropdownProps = { + width: string; +}; type Props = { placeholder?: React.ReactNode; allTags: string[]; selectedTags: string[]; - setSelectedTags: (tags: React.SetStateAction) => void; + setSelectedTags: (tags: string[]) => void; renderTag?: (tag: string, tagProps: TagProps) => React.ReactNode; renderTagList?: (tags: React.ReactNode[]) => React.ReactNode; - renderDropdown?: (dropdown: React.ReactNode) => React.ReactNode; + renderDropdown?: (dropdown: React.ReactNode, dropdownProps: DropdownProps) => React.ReactNode; renderDropdownItem?: (tag: string, dropdownItemProps: DropdownItemProps) => React.ReactNode; dropdownStyles?: React.CSSProperties; }; const defaultRenderTag = (tag: string, tagProps: TagProps) => { return ( - + {tag} @@ -70,6 +74,7 @@ export const TagSelector = ({ renderTagList, }: Props) => { const [isDropdownOpen, setIsDropdownOpen] = React.useState(false); + const {viewport, containerProps} = useViewport(); const dropdown = React.useMemo(() => { const dropdownContent = ( ); if (renderDropdown) { - return renderDropdown(dropdownContent); + return renderDropdown(dropdownContent, {width: viewport.width + 'px'}); } - return {dropdownContent}; - }, [allTags, dropdownStyles, renderDropdown, renderDropdownItem, selectedTags, setSelectedTags]); + return {dropdownContent}; + }, [ + allTags, + dropdownStyles, + renderDropdown, + renderDropdownItem, + selectedTags, + setSelectedTags, + viewport.width, + ]); const dropdownContainer = React.useRef(null); @@ -108,7 +121,7 @@ export const TagSelector = ({ const tags = selectedTags.map((tag) => (renderTag || defaultRenderTag)(tag, { remove: (ev) => { - setSelectedTags((tags) => tags.filter((t) => t !== tag)); + setSelectedTags(selectedTags.filter((t) => t !== tag)); ev.stopPropagation(); }, }), @@ -117,7 +130,7 @@ export const TagSelector = ({ return renderTagList(tags); } return tags; - }, [placeholder, selectedTags, renderTag, renderTagList]); + }, [selectedTags, renderTagList, placeholder, renderTag, setSelectedTags]); return ( { setIsDropdownOpen((isOpen) => !isOpen); }} + {...containerProps} > {tagsContent}