[ui] add filters for storage kind (#22041)

## Summary Add frontend filters for storage kind to the asset list and asset graph pages. <img width="562" alt="Screenshot 2024-05-22 at 2 14 30 PM" src="https://github.com/dagster-io/dagster/assets/10215173/4a201721-3517-4021-9243-67482722c6d9"> ## Test Plan Tested locally
dagster-io · May 28, 2024 · d47c5b3 · d47c5b3 · github-actions · May 28, 2024
1 parent 5ad3b82
commit d47c5b3
Show file tree

Hide file tree

Showing 4 changed files with 121 additions and 4 deletions.
diff --git a/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorerFilters.tsx b/js_modules/dagster-ui/packages/ui-core/src/asset-graph/AssetGraphExplorerFilters.tsx
@@ -4,6 +4,7 @@ import React, {useContext, useEffect, useState} from 'react';
 import {GraphNode} from './Utils';
 import {CloudOSSContext} from '../app/CloudOSSContext';
 import {AssetFilterState} from '../assets/useAssetDefinitionFilterState';
+import {isCanonicalStorageKindTag} from '../graph/KindTags';
 import {ChangeReason} from '../graphql/types';
 import {useFilters} from '../ui/Filters';
 import {useAssetGroupFilter, useAssetGroupsForAssets} from '../ui/Filters/useAssetGroupFilter';
@@ -16,6 +17,7 @@ import {
   useComputeKindTagFilter,
 } from '../ui/Filters/useComputeKindTagFilter';
 import {FilterObject, FilterTag, FilterTagHighlightedText} from '../ui/Filters/useFilter';
+import {useStorageKindFilter} from '../ui/Filters/useStorageKindFilter';
 import {WorkspaceContext} from '../workspace/WorkspaceContext';
 
 type Props = {
@@ -36,6 +38,7 @@ const defaultState: AssetFilterState = {
     repos: [],
     selectAllFilters: [],
     tags: [],
+    storageKindTags: [],
   },
   setAssetTags: () => {},
   setChangedInBranch: () => {},
@@ -46,6 +49,7 @@ const defaultState: AssetFilterState = {
   setRepos: () => {},
   setSelectAllFilters: () => {},
   filterFn: () => true,
+  setStorageKindTags: () => {},
 };
 
 export function useAssetGraphExplorerFilters({
@@ -61,14 +65,24 @@ export function useAssetGraphExplorerFilters({
   const {allRepos} = useContext(WorkspaceContext);
 
   const {
-    filters: {changedInBranch, computeKindTags, repos, owners, groups, tags, selectAllFilters},
+    filters: {
+      changedInBranch,
+      computeKindTags,
+      repos,
+      owners,
+      groups,
+      tags,
+      storageKindTags,
+      selectAllFilters,
+    },
     setAssetTags,
     setChangedInBranch,
     setComputeKindTags,
     setGroups,
     setOwners,
     setRepos,
     setSelectAllFilters,
+    setStorageKindTags,
   } = assetFilterState || defaultState;
 
   const reposFilter = useCodeLocationFilter(repos ? {repos, setRepos} : undefined);
@@ -93,12 +107,23 @@ export function useAssetGraphExplorerFilters({
     setComputeKindTags,
   });
 
+  const allStorageKindTags = allAssetTags.filter(isCanonicalStorageKindTag);
+  const allNonStorageKindTags = allAssetTags.filter((tag) => !isCanonicalStorageKindTag(tag));
+
   const tagsFilter = useAssetTagFilter({
-    allAssetTags,
+    allAssetTags: allNonStorageKindTags,
     tags: selectAllFilters.includes('tags') ? allAssetTags : tags,
     setTags: setAssetTags,
   });
 
+  const storageKindTagsFilter = useStorageKindFilter({
+    allAssetStorageKindTags: allStorageKindTags,
+    storageKindTags: selectAllFilters.includes('storageKindTags')
+      ? allStorageKindTags
+      : storageKindTags,
+    setStorageKindTags,
+  });
+
   const allAssetOwners = useAssetOwnersForAssets(nodes);
   const ownerFilter = useAssetOwnerFilter({
     allAssetOwners,
@@ -177,6 +202,7 @@ export function useAssetGraphExplorerFilters({
     filters.push(changedFilter);
   }
   filters.push(kindTagsFilter);
+  filters.push(storageKindTagsFilter);
   filters.push(tagsFilter);
   filters.push(ownerFilter);
   const {button, activeFiltersJsx} = useFilters({filters});

diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/useAssetCatalogFiltering.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/useAssetCatalogFiltering.tsx
@@ -6,6 +6,7 @@ import {AssetTableFragment} from './types/AssetTableFragment.types';
 import {useAssetDefinitionFilterState} from './useAssetDefinitionFilterState';
 import {useAssetSearch} from './useAssetSearch';
 import {CloudOSSContext} from '../app/CloudOSSContext';
+import {isCanonicalStorageKindTag} from '../graph/KindTags';
 import {useQueryPersistedState} from '../hooks/useQueryPersistedState';
 import {useFilters} from '../ui/Filters';
 import {useAssetGroupFilter} from '../ui/Filters/useAssetGroupFilter';
@@ -18,6 +19,7 @@ import {
   useComputeKindTagFilter,
 } from '../ui/Filters/useComputeKindTagFilter';
 import {FilterObject} from '../ui/Filters/useFilter';
+import {useStorageKindFilter} from '../ui/Filters/useStorageKindFilter';
 import {WorkspaceContext} from '../workspace/WorkspaceContext';
 
 const EMPTY_ARRAY: any[] = [];
@@ -37,6 +39,7 @@ export function useAssetCatalogFiltering(
     setGroups,
     setOwners,
     setRepos,
+    setStorageKindTags,
   } = useAssetDefinitionFilterState();
 
   const searchPath = (search || '')
@@ -77,13 +80,29 @@ export function useAssetCatalogFiltering(
     owners: filters.owners,
     setOwners,
   });
+
+  const tags = useAssetTagsForAssets(pathMatches);
+  const storageKindTags = tags.filter(isCanonicalStorageKindTag);
+  const nonStorageKindTags = tags.filter((tag) => !isCanonicalStorageKindTag(tag));
+
   const tagsFilter = useAssetTagFilter({
-    allAssetTags: useAssetTagsForAssets(pathMatches),
+    allAssetTags: nonStorageKindTags,
     tags: filters.tags,
     setTags: setAssetTags,
   });
+  const storageKindFilter = useStorageKindFilter({
+    allAssetStorageKindTags: storageKindTags,
+    storageKindTags: filters.storageKindTags,
+    setStorageKindTags,
+  });
 
-  const uiFilters: FilterObject[] = [groupsFilter, computeKindFilter, ownersFilter, tagsFilter];
+  const uiFilters: FilterObject[] = [
+    groupsFilter,
+    computeKindFilter,
+    storageKindFilter,
+    ownersFilter,
+    tagsFilter,
+  ];
   const {isBranchDeployment} = React.useContext(CloudOSSContext);
   if (isBranchDeployment) {
     uiFilters.push(changedInBranchFilter);

diff --git a/js_modules/dagster-ui/packages/ui-core/src/assets/useAssetDefinitionFilterState.tsx b/js_modules/dagster-ui/packages/ui-core/src/assets/useAssetDefinitionFilterState.tsx
@@ -2,6 +2,7 @@ import isEqual from 'lodash/isEqual';
 import {SetStateAction, useCallback, useMemo} from 'react';
 
 import {buildAssetGroupSelector} from './AssetGroupSuggest';
+import {isCanonicalStorageKindTag} from '../graph/KindTags';
 import {
   AssetGroupSelector,
   AssetNode,
@@ -25,6 +26,7 @@ type FilterableAssetDefinition = Partial<
 export type AssetFilterBaseType = {
   groups: AssetGroupSelector[];
   computeKindTags: string[];
+  storageKindTags: DefinitionTag[];
   changedInBranch: ChangeReason[];
   owners: AssetOwner[];
   tags: DefinitionTag[];
@@ -40,6 +42,7 @@ export const useAssetDefinitionFilterState = () => {
     encode: ({
       groups,
       computeKindTags,
+      storageKindTags,
       changedInBranch,
       owners,
       tags,
@@ -48,6 +51,7 @@ export const useAssetDefinitionFilterState = () => {
     }) => ({
       groups: groups?.length ? JSON.stringify(groups) : undefined,
       computeKindTags: computeKindTags?.length ? JSON.stringify(computeKindTags) : undefined,
+      storageKindTags: storageKindTags?.length ? JSON.stringify(storageKindTags) : undefined,
       changedInBranch: changedInBranch?.length ? JSON.stringify(changedInBranch) : undefined,
       owners: owners?.length ? JSON.stringify(owners) : undefined,
       tags: tags?.length ? JSON.stringify(tags) : undefined,
@@ -57,6 +61,7 @@ export const useAssetDefinitionFilterState = () => {
     decode: (qs) => ({
       groups: qs.groups ? JSON.parse(qs.groups) : [],
       computeKindTags: qs.computeKindTags ? JSON.parse(qs.computeKindTags) : [],
+      storageKindTags: qs.storageKindTags ? JSON.parse(qs.storageKindTags) : [],
       changedInBranch: qs.changedInBranch ? JSON.parse(qs.changedInBranch) : [],
       owners: qs.owners ? JSON.parse(qs.owners) : [],
       tags: qs.tags ? JSON.parse(qs.tags) : [],
@@ -76,6 +81,7 @@ export const useAssetDefinitionFilterState = () => {
 
   const {
     setComputeKindTags,
+    setStorageKindTags,
     setGroups,
     setChangedInBranch,
     setOwners,
@@ -93,6 +99,7 @@ export const useAssetDefinitionFilterState = () => {
     }
     return {
       setComputeKindTags: makeSetter('computeKindTags'),
+      setStorageKindTags: makeSetter('storageKindTags'),
       setGroups: makeSetter('groups'),
       setChangedInBranch: makeSetter('changedInBranch'),
       setOwners: makeSetter('owners'),
@@ -107,6 +114,7 @@ export const useAssetDefinitionFilterState = () => {
     setFilters,
     filterFn,
     setComputeKindTags,
+    setStorageKindTags,
     setGroups,
     setChangedInBranch,
     setOwners,
@@ -173,6 +181,21 @@ export function filterAssetDefinition(
     }
   }
 
+  const isAllStorageKindTagsSelected = filters.selectAllFilters?.includes('storageKindTags');
+  const storageKindTag = definition?.tags?.find(isCanonicalStorageKindTag);
+  if (isAllStorageKindTagsSelected) {
+    if (!storageKindTag) {
+      return false;
+    }
+  } else if (filters.storageKindTags?.length) {
+    if (
+      !storageKindTag ||
+      !doesFilterArrayMatchValueArray(filters.storageKindTags, [storageKindTag])
+    ) {
+      return false;
+    }
+  }
+
   const isAllChangedInBranchSelected = filters.selectAllFilters?.includes('changedInBranch');
   if (isAllChangedInBranchSelected) {
     if (!definition?.changedReasons?.length) {

diff --git a/js_modules/dagster-ui/packages/ui-core/src/ui/Filters/useStorageKindFilter.tsx b/js_modules/dagster-ui/packages/ui-core/src/ui/Filters/useStorageKindFilter.tsx
@@ -0,0 +1,49 @@
+import {Box, Icon} from '@dagster-io/ui-components';
+import {useMemo} from 'react';
+
+import {buildDefinitionTag} from './useAssetTagFilter';
+import {useStaticSetFilter} from './useStaticSetFilter';
+import {DefinitionTag} from '../../graphql/types';
+import {TruncatedTextWithFullTextOnHover} from '../../nav/getLeftNavItemsForOption';
+
+const emptyArray: any[] = [];
+
+export const useStorageKindFilter = ({
+  allAssetStorageKindTags,
+  storageKindTags,
+  setStorageKindTags,
+}: {
+  allAssetStorageKindTags: DefinitionTag[];
+  storageKindTags?: null | DefinitionTag[];
+  setStorageKindTags?: null | ((s: DefinitionTag[]) => void);
+}) => {
+  const memoizedState = useMemo(() => storageKindTags?.map(buildDefinitionTag), [storageKindTags]);
+  return useStaticSetFilter<DefinitionTag>({
+    name: 'Storage kind',
+    icon: 'tag',
+    allValues: useMemo(
+      () =>
+        allAssetStorageKindTags.map((value) => ({
+          value,
+          match: [value.key + ':' + value.value],
+        })),
+      [allAssetStorageKindTags],
+    ),
+    menuWidth: '300px',
+    renderLabel: ({value: tag}) => {
+      return (
+        <Box flex={{direction: 'row', gap: 4, alignItems: 'center'}}>
+          <Icon name="tag" />
+          <TruncatedTextWithFullTextOnHover tooltipText={tag.value} text={tag.value} />
+        </Box>
+      );
+    },
+    getStringValue: ({value}) => value,
+    state: memoizedState ?? emptyArray,
+    onStateChanged: (values) => {
+      setStorageKindTags?.(Array.from(values));
+    },
+    matchType: 'all-of',
+    canSelectAll: false,
+  });
+};