Skip to content

Commit

Permalink
refactor yuge
Browse files Browse the repository at this point in the history
  • Loading branch information
@salazarm
salazarm committed Jan 28, 2025
1 parent 2480c65 commit b8b7551
Show file tree
Hide file tree
Showing 14 changed files with 1,173 additions and 1,749 deletions.
2 changes: 1 addition & 1 deletion js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/AssetSelectionInput.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ interface AssetSelectionInputProps {
linter?: Linter<any>;
useAssetSelectionAutoCompleteProvider?: (
assets: AssetGraphQueryItem[],
) => SelectionAutoCompleteProvider<any>;
) => SelectionAutoCompleteProvider;
}

const defaultLinter = createSelectionLinter({
Expand Down
167 changes: 12 additions & 155 deletions .../packages/ui-core/src/asset-selection/input/useAssetSelectionAutoCompleteProvider.oss.tsx
View file
Original file line number Diff line number Diff line change
@@ -1,107 +1,35 @@
import {BodySmall, Box, Colors, Icon, IconName, MonoSmall} from '@dagster-io/ui-components';
import {IconName} from '@dagster-io/ui-components';
import {useMemo} from 'react';

import {getAttributesMap} from './util';
import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
import {useUpdatingRef} from '../../hooks/useUpdatingRef';
import {createSelectionAutoComplete} from '../../selection/SelectionAutoComplete';
import {
BaseSuggestion,
SelectionAutoCompleteProvider,
createProvider,
} from '../../selection/SelectionAutoCompleteProvider';
import {createSelectionAutoCompleteProviderFromAttributeMap} from '../../selection/SelectionAutoCompleteProviderFromAttributeMap';

const FUNCTIONS = ['sinks', 'roots'];

type Attribute = 'kind' | 'code_location' | 'group' | 'owner' | 'tag' | 'status';

type Suggestion =
| {
text: string;
displayText: string;
type: 'function' | 'attribute-value' | 'attribute-with-value';
attributeName?: string;
}
| {
text: string;
displayText: string;
type: 'attribute';
attributeName: Attribute | 'key';
}
| {
text: string;
type: 'substring';
value: string;
};
type Attribute = 'kind' | 'code_location' | 'group' | 'owner' | 'tag' | 'status' | 'key';

export function useAssetSelectionAutoCompleteProvider(
assets: AssetGraphQueryItem[],
): SelectionAutoCompleteProvider<Suggestion> {
const attributesMapRef = useUpdatingRef(getAttributesMap(assets));
): SelectionAutoCompleteProvider {
const attributesMap = useMemo(() => getAttributesMap(assets), [assets]);

const baseProvider = useMemo(
() =>
createSelectionAutoCompleteProviderFromAttributeMap<
typeof attributesMapRef.current,
'key',
Suggestion
>({
nameBase: 'key',
attributesMapRef,
functions: FUNCTIONS,
doesValueIncludeQuery: (_attribute, value, query) => value.includes(query),
createAttributeSuggestion: (attribute, textCallback) => {
const text = `${attribute}:`;
return {
text: textCallback ? textCallback(text) : text,
displayText: text,
type: 'attribute',
attributeName: attribute,
nameBase: attribute === 'key',
};
},
createAttributeValueSuggestion: (attribute, value, textCallback) => {
const text = `"${value}"`;
return {
text: textCallback ? textCallback(text) : text,
displayText: value,
type: 'attribute-value',
attributeName: attribute,
};
},
createFunctionSuggestion: (func, textCallback, options) => {
const text = options?.includeParenthesis ? `${func}()` : func;
return {
text: textCallback ? textCallback(text) : text,
displayText: `${func}()`,
type: 'function',
};
},
createSubstringSuggestion: (query, textCallback) => {
const text = `key_substring:"${query}"`;
return {
text: textCallback ? textCallback(text) : text,
value: query,
type: 'substring',
};
},
createAttributeValueIncludeAttributeSuggestion: (attribute, value, textCallback) => {
const text = `${attribute}:"${value}"`;
return {
text: textCallback ? textCallback(text) : text,
displayText: `${attribute}:${value}`,
type: 'attribute-with-value',
attributeName: attribute,
};
},
createSelectionAutoCompleteProviderFromAttributeMap({
attributesMap,
...createProvider<typeof attributesMap, 'key'>({nameBase: 'key', attributeToIcon: iconMap}),
}),
[attributesMapRef],
[attributesMap],
);
const selectionHint = useMemo(() => createSelectionAutoComplete(baseProvider), [baseProvider]);

return {
...baseProvider,
useAutoComplete: (line, cursorIndex) => {
useAutoComplete: ({line, cursorIndex}) => {
const autoCompleteResults = useMemo(
() => selectionHint(line, cursorIndex),
[line, cursorIndex],
Expand All @@ -111,86 +39,15 @@ export function useAssetSelectionAutoCompleteProvider(
loading: false,
};
},
renderResult: (suggestion) => <SuggestionItem suggestion={suggestion} />,
};
}

const attributeToIcon: Record<Attribute, IconName> = {
export const iconMap: Record<Attribute, IconName> = {
key: 'magnify_glass',
kind: 'compute_kind',
code_location: 'code_location',
group: 'asset_group',
owner: 'owner',
tag: 'tag',
status: 'status',
};

export const SuggestionItem = ({suggestion}: {suggestion: Suggestion | BaseSuggestion}) => {
let label;
let icon: IconName | null = null;
let value: string | null = 'displayText' in suggestion ? suggestion.displayText : null;
if (suggestion.type === 'attribute' && suggestion.attributeName === 'key') {
if (suggestion.text.endsWith('_substring:')) {
icon = 'magnify_glass_checked';
label = 'Contains match';
} else {
icon = 'magnify_glass';
label = 'Exact match';
}
} else if (suggestion.type === 'down-traversal' || suggestion.type === 'up-traversal') {
icon = 'curly_braces';
label =
suggestion.type === 'down-traversal'
? 'Include downstream dependencies'
: 'Include upstream dependencies';
} else if (suggestion.type === 'logical_operator') {
icon = 'curly_braces';
label = suggestion.displayText.toUpperCase();
} else if (suggestion.type === 'parenthesis') {
icon = 'curly_braces';
label = 'Parenthesis';
value = suggestion.text;
} else if (suggestion.type === 'function') {
if (suggestion.displayText === 'roots()') {
label = 'Roots';
icon = 'arrow_upward';
} else if (suggestion.displayText === 'sinks()') {
label = 'Sinks';
icon = 'arrow_indent';
}
} else if (suggestion.type === 'attribute') {
if (suggestion.attributeName === 'key') {
label = suggestion.displayText.replace(':', '').replace('_', ' ');
label = label.charAt(0).toUpperCase() + label.slice(1);
} else if ('attributeName' in suggestion && suggestion.attributeName) {
icon = attributeToIcon[suggestion.attributeName]!;
label = suggestion.displayText.replace(':', '').replace('_', ' ');
label = label.charAt(0).toUpperCase() + label.slice(1);
}
} else if (suggestion.type === 'attribute-with-value') {
const firstColon = suggestion.displayText.indexOf(':');
const attributeKey = suggestion.displayText.slice(0, firstColon);
const attributeValue = suggestion.displayText.slice(firstColon + 1);
label = (
<Box flex={{direction: 'row', alignItems: 'center', gap: 2}}>
<MonoSmall color={Colors.textLight()}>{attributeKey}:</MonoSmall>
<MonoSmall>{attributeValue}</MonoSmall>
</Box>
);
value = null;
} else if (suggestion.type === 'attribute-value') {
label = suggestion.displayText;
value = null;
} else if (suggestion.type === 'substring') {
label = `Asset key contains "${suggestion.value}"`;
value = `key_substring:${suggestion.value}`;
}
return (
<Box flex={{direction: 'row', alignItems: 'center', justifyContent: 'space-between', gap: 24}}>
<Box flex={{direction: 'row', alignItems: 'center', gap: 6}}>
{icon ? <Icon name={icon} size={12} style={{margin: 0}} /> : null}
<BodySmall>{label}</BodySmall>
</Box>
<MonoSmall>{value}</MonoSmall>
</Box>
);
};
16 changes: 8 additions & 8 deletions js_modules/dagster-ui/packages/ui-core/src/asset-selection/input/util.ts
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
import {assertUnreachable} from '../../app/Util';
import {AssetGraphQueryItem} from '../../asset-graph/useAssetGraphData';
import {isKindTag} from '../../graph/KindTags';
import {weakMapMemoize} from '../../util/weakMapMemoize';
import {buildRepoPathForHuman} from '../../workspace/buildRepoAddress';

export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
const assetNamesSet: Set<string> = new Set();
const tagNamesSet: Set<string> = new Set();
const tagNamesSet: Set<{key: string; value: string}> = new Set();
const ownersSet: Set<string> = new Set();
const groupsSet: Set<string> = new Set();
const kindsSet: Set<string> = new Set();
Expand All @@ -17,13 +18,7 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
if (isKindTag(tag)) {
return;
}
if (tag.key && tag.value) {
// We add quotes around the equal sign here because the auto-complete suggestion already wraps the entire value in quotes.
// So wer end up with tag:"key"="value" as the final suggestion
tagNamesSet.add(`${tag.key}"="${tag.value}`);
} else {
tagNamesSet.add(tag.key);
}
tagNamesSet.add(memoizedTag(tag.key, tag.value));
});
asset.node.owners.forEach((owner) => {
switch (owner.__typename) {
Expand Down Expand Up @@ -66,3 +61,8 @@ export const getAttributesMap = (assets: AssetGraphQueryItem[]) => {
code_location: codeLocations,
};
};

const memoizedTag = weakMapMemoize((key: string, value: string) => ({
key,
value,
}));
Loading

0 comments on commit b8b7551

Please sign in to comment.