diff --git a/superset-frontend/src/SqlLab/components/TableExploreTree/TreeNodeRenderer.tsx b/superset-frontend/src/SqlLab/components/TableExploreTree/TreeNodeRenderer.tsx index ba562348738..06605fdd601 100644 --- a/superset-frontend/src/SqlLab/components/TableExploreTree/TreeNodeRenderer.tsx +++ b/superset-frontend/src/SqlLab/components/TableExploreTree/TreeNodeRenderer.tsx @@ -79,6 +79,7 @@ export interface TreeNodeRendererProps extends NodeRendererProps { searchTerm: string; catalog: string | null | undefined; pinnedTableKeys: Set; + pinnedSchemas: Set; selectStarMap: Record; handleRefreshTables: (params: { dbId: number; @@ -91,6 +92,11 @@ export interface TreeNodeRendererProps extends NodeRendererProps { catalogName: string | null, ) => void; handleUnpinTable: (tableName: string, schemaName: string) => void; + handlePinSchema: (schemaName: string) => void; + handleUnpinSchema: (schemaName: string) => void; + refreshTableSchema: (id: string) => void; + sortedTables: Record; + toggleSortColumns: (tableId: string) => void; } const TreeNodeRenderer: React.FC = ({ @@ -101,19 +107,23 @@ const TreeNodeRenderer: React.FC = ({ searchTerm, catalog, pinnedTableKeys, + pinnedSchemas, selectStarMap, handleRefreshTables, handlePinTable, handleUnpinTable, + handlePinSchema, + handleUnpinSchema, + refreshTableSchema, + sortedTables, + toggleSortColumns, }) => { const theme = useTheme(); const { data } = node; const parts = data.id.split(':'); const [identifier, _dbId, schema, tableName] = parts; - // Use manually tracked open state for icon display - // This prevents search auto-expansion from affecting the icon - const isManuallyOpen = manuallyOpenedNodes[data.id] ?? false; + const isManuallyOpen = node.isOpen && !node.data.disableCheckbox; const isLoading = loadingNodes[data.id] ?? false; const renderIcon = () => { @@ -135,12 +145,7 @@ const TreeNodeRenderer: React.FC = ({ ? Icons.FunctionOutlined : Icons.TableOutlined; if (isLoading) { - return ( - <> - - - - ); + return ; } return ; } @@ -233,7 +238,27 @@ const TreeNodeRenderer: React.FC = ({ {highlightText(data.name, searchTerm)} {identifier === 'schema' && ( -
+
e.stopPropagation()} + > + {pinnedSchemas.has(schema) && ( +
+ + } + onClick={() => handleUnpinSchema(schema)} + /> +
+ )}
{ @@ -246,6 +271,30 @@ const TreeNodeRenderer: React.FC = ({ }} tooltipContent={t('Force refresh table list')} /> + + ) : ( + + ) + } + onClick={() => + pinnedSchemas.has(schema) + ? handleUnpinSchema(schema) + : handlePinSchema(schema) + } + />
)} @@ -288,6 +337,31 @@ const TreeNodeRenderer: React.FC = ({ } /> )} + + } + onClick={() => toggleSortColumns(data.id)} + /> + } + onClick={() => refreshTableSchema(data.id)} + /> `${dbId ?? ''}:${catalog ?? ''}`; + +const getPinnedSchemasFromStorage = ( + dbId: number | undefined, + catalog: string | null | undefined, +): Set => { + if (!dbId) return new Set(); + const stored = getItem(LocalStorageKeys.SqllabPinnedSchemas, {}); + const key = getPinnedSchemasStorageKey(dbId, catalog); + const schemas = stored[key]; + return Array.isArray(schemas) ? new Set(schemas) : new Set(); +}; + +const savePinnedSchemasToStorage = ( + dbId: number | undefined, + catalog: string | null | undefined, + schemas: Set, +) => { + if (!dbId) return; + const stored = getItem(LocalStorageKeys.SqllabPinnedSchemas, {}); + const key = getPinnedSchemasStorageKey(dbId, catalog); + setItem(LocalStorageKeys.SqllabPinnedSchemas, { + ...stored, + [key]: [...schemas], + }); +}; + const TableExploreTree: React.FC = ({ queryEditorId }) => { const dispatch = useDispatch(); const theme = useTheme(); @@ -161,6 +196,7 @@ const TableExploreTree: React.FC = ({ queryEditorId }) => { selectStarMap, handleToggle, handleRefreshTables, + refreshTableSchema, errorPayload, } = useTreeData({ dbId, @@ -199,6 +235,83 @@ const TableExploreTree: React.FC = ({ queryEditorId }) => { }, [dispatch, tables, editorId, dbId], ); + const [pinnedSchemas, setPinnedSchemas] = useState>(() => + getPinnedSchemasFromStorage(dbId, catalog), + ); + + const previousDbIdRef = useRef(dbId); + const previousCatalogRef = useRef(catalog); + + // Single effect handles both loading and persisting pinned schemas. + // Using refs to detect source changes avoids the race condition where the + // persist branch would run with stale pinnedSchemas right after a dbId/catalog + // change, corrupting the new source's stored pins. + useEffect(() => { + const dbChanged = previousDbIdRef.current !== dbId; + const catalogChanged = previousCatalogRef.current !== catalog; + + if (dbChanged || catalogChanged) { + previousDbIdRef.current = dbId; + previousCatalogRef.current = catalog; + setPinnedSchemas(getPinnedSchemasFromStorage(dbId, catalog)); + return; + } + + savePinnedSchemasToStorage(dbId, catalog, pinnedSchemas); + }, [dbId, catalog, pinnedSchemas]); + + const handlePinSchema = useCallback((schemaName: string) => { + setPinnedSchemas(prev => new Set([...prev, schemaName])); + }, []); + + const handleUnpinSchema = useCallback((schemaName: string) => { + setPinnedSchemas(prev => { + const next = new Set(prev); + next.delete(schemaName); + return next; + }); + }, []); + + const sortedTreeData = useMemo(() => { + if (pinnedSchemas.size === 0) return treeData; + const pinned = treeData.filter(node => pinnedSchemas.has(node.name)); + const rest = treeData.filter(node => !pinnedSchemas.has(node.name)); + return [...pinned, ...rest]; + }, [treeData, pinnedSchemas]); + + const [sortedTables, setSortedTables] = useState>({}); + + useEffect(() => { + setSortedTables({}); + }, [dbId, catalog]); + + const toggleSortColumns = useCallback((tableId: string) => { + setSortedTables(prev => ({ ...prev, [tableId]: !prev[tableId] })); + }, []); + + const displayTreeData = useMemo(() => { + const activeSorted = Object.keys(sortedTables).filter( + id => sortedTables[id], + ); + if (activeSorted.length === 0) return sortedTreeData; + + const sortedSet = new Set(activeSorted); + return sortedTreeData.map(schemaNode => ({ + ...schemaNode, + children: schemaNode.children?.map(tableNode => { + if (tableNode.type !== 'table' || !sortedSet.has(tableNode.id)) { + return tableNode; + } + const { children } = tableNode; + if (!children || children.length <= 1) return tableNode; + return { + ...tableNode, + children: [...children].sort((a, b) => a.name.localeCompare(b.name)), + }; + }), + })); + }, [sortedTreeData, sortedTables]); + const [searchTerm, setSearchTerm] = useState(''); const handleSearchChange = useCallback( ({ target }: ChangeEvent) => setSearchTerm(target.value), @@ -270,8 +383,8 @@ const TableExploreTree: React.FC = ({ queryEditorId }) => { return false; }; - return treeData.some(node => checkNode(node)); - }, [searchTerm, treeData]); + return displayTreeData.some(node => checkNode(node)); + }, [searchTerm, displayTreeData]); // Node renderer for react-arborist const renderNode = useCallback( @@ -283,19 +396,31 @@ const TableExploreTree: React.FC = ({ queryEditorId }) => { searchTerm={searchTerm} catalog={catalog} pinnedTableKeys={pinnedTableKeys} + pinnedSchemas={pinnedSchemas} selectStarMap={selectStarMap} handleRefreshTables={handleRefreshTables} handlePinTable={handlePinTable} handleUnpinTable={handleUnpinTable} + handlePinSchema={handlePinSchema} + handleUnpinSchema={handleUnpinSchema} + refreshTableSchema={refreshTableSchema} + sortedTables={sortedTables} + toggleSortColumns={toggleSortColumns} /> ), [ catalog, pinnedTableKeys, + pinnedSchemas, selectStarMap, handleRefreshTables, handlePinTable, handleUnpinTable, + handlePinSchema, + handleUnpinSchema, + refreshTableSchema, + sortedTables, + toggleSortColumns, loadingNodes, manuallyOpenedNodes, searchTerm, @@ -369,7 +494,7 @@ const TableExploreTree: React.FC = ({ queryEditorId }) => { return ( ref={treeRef} - data={treeData} + data={displayTreeData} width="100%" height={height || 500} rowHeight={ROW_HEIGHT} diff --git a/superset-frontend/src/SqlLab/components/TableExploreTree/useTreeData.ts b/superset-frontend/src/SqlLab/components/TableExploreTree/useTreeData.ts index 9cc54ce599c..a61ea25c6ca 100644 --- a/superset-frontend/src/SqlLab/components/TableExploreTree/useTreeData.ts +++ b/superset-frontend/src/SqlLab/components/TableExploreTree/useTreeData.ts @@ -17,6 +17,7 @@ * under the License. */ import { useMemo, useReducer, useCallback } from 'react'; +import { useDispatch } from 'react-redux'; import { t } from '@apache-superset/core/translation'; import { Table, @@ -26,6 +27,7 @@ import { useLazyTableMetadataQuery, useLazyTableExtendedMetadataQuery, } from 'src/hooks/apiResources'; +import { addDangerToast } from 'src/SqlLab/actions/sqlLab'; import type { TreeNodeData } from './types'; import { SupersetError } from '@superset-ui/core'; @@ -42,6 +44,7 @@ interface TreeDataState { type TreeDataAction = | { type: 'SET_TABLE_DATA'; key: string; data: { options: Table[] } } | { type: 'SET_TABLE_SCHEMA_DATA'; key: string; data: TableMetaData } + | { type: 'CLEAR_TABLE_SCHEMA_DATA'; key: string } | { type: 'SET_LOADING_NODE'; nodeId: string; loading: boolean } | { type: 'SET_ERROR'; errorPayload: SupersetError | null }; @@ -71,6 +74,10 @@ function treeDataReducer( [action.key]: action.data, }, }; + case 'CLEAR_TABLE_SCHEMA_DATA': { + const { [action.key]: _, ...rest } = state.tableSchemaData; + return { ...state, tableSchemaData: rest }; + } case 'SET_LOADING_NODE': return { ...state, @@ -108,6 +115,7 @@ interface UseTreeDataResult { catalog: string | null | undefined; schema: string; }) => void; + refreshTableSchema: (id: string) => void; errorPayload: SupersetError | null; } @@ -122,6 +130,7 @@ const useTreeData = ({ catalog, pinnedTables, }: UseTreeDataParams): UseTreeDataResult => { + const reduxDispatch = useDispatch(); // Schema data from API const { currentData: schemaData, @@ -137,6 +146,64 @@ const useTreeData = ({ const [state, dispatch] = useReducer(treeDataReducer, initialState); const { tableData, tableSchemaData, loadingNodes, errorPayload } = state; + // Shared helper: fetch table metadata + extended metadata and store in state. + // preferCacheValue=true on initial open (use cached data if available), + // preferCacheValue=false on explicit refresh (bypass cache). + const fetchAndStoreTableSchema = useCallback( + (id: string, preferCacheValue: boolean) => { + if (loadingNodes[id]) return; + + const parts = id.split(':'); + const [, databaseId, schema, table] = parts; + const parsedDbId = Number(databaseId); + const tableKey = `${parsedDbId}:${schema}:${table}`; + + dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: true }); + + // .unwrap() causes RTK Query to reject on error so .catch() fires. + // Without it RTK Query resolves with { error } instead of rejecting. + Promise.all([ + fetchTableMetadata( + { dbId: parsedDbId, catalog, schema, table }, + preferCacheValue, + ).unwrap(), + fetchTableExtendedMetadata( + { dbId: parsedDbId, catalog, schema, table }, + preferCacheValue, + ).unwrap(), + ]) + .then(([tableMetadata, tableExtendedMetadata]) => { + if (tableMetadata) { + dispatch({ + type: 'SET_TABLE_SCHEMA_DATA', + key: tableKey, + data: { ...tableMetadata, ...tableExtendedMetadata }, + }); + } + }) + .catch(() => { + reduxDispatch( + addDangerToast( + t( + 'An error occurred while fetching table metadata for %s', + table, + ), + ), + ); + }) + .finally(() => { + dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: false }); + }); + }, + [ + catalog, + fetchTableExtendedMetadata, + fetchTableMetadata, + loadingNodes, + reduxDispatch, + ], + ); + // Handle async loading when node is toggled open const handleToggle = useCallback( async (id: string, isOpen: boolean) => { @@ -150,20 +217,14 @@ const useTreeData = ({ if (identifier === 'schema') { const schemaKey = `${parsedDbId}:${schema}`; if (!tableData?.[schemaKey]) { - // Set loading state dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: true }); - // Fetch tables asynchronously fetchLazyTables( - { - dbId: parsedDbId, - catalog, - schema, - forceRefresh: false, - }, + { dbId: parsedDbId, catalog, schema, forceRefresh: false }, true, ) - .then(({ data }) => { + .unwrap() + .then(data => { if (data) { dispatch({ type: 'SET_TABLE_DATA', key: schemaKey, data }); } @@ -191,59 +252,14 @@ const useTreeData = ({ if (pinnedTables[tableKey]) return; if (!tableSchemaData[tableKey]) { - // Set loading state - dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: true }); - - // Fetch metadata asynchronously - Promise.all([ - fetchTableMetadata( - { - dbId: parsedDbId, - catalog, - schema, - table, - }, - true, - ), - fetchTableExtendedMetadata( - { - dbId: parsedDbId, - catalog, - schema, - table, - }, - true, - ), - ]) - .then( - ([{ data: tableMetadata }, { data: tableExtendedMetadata }]) => { - if (tableMetadata) { - dispatch({ - type: 'SET_TABLE_SCHEMA_DATA', - key: tableKey, - data: { - ...tableMetadata, - ...tableExtendedMetadata, - }, - }); - } - }, - ) - .finally(() => { - dispatch({ - type: 'SET_LOADING_NODE', - nodeId: id, - loading: false, - }); - }); + fetchAndStoreTableSchema(id, true); } } }, [ catalog, + fetchAndStoreTableSchema, fetchLazyTables, - fetchTableExtendedMetadata, - fetchTableMetadata, pinnedTables, tableData, tableSchemaData, @@ -289,6 +305,13 @@ const useTreeData = ({ [fetchLazyTables], ); + const refreshTableSchema = useCallback( + (id: string) => { + fetchAndStoreTableSchema(id, false); + }, + [fetchAndStoreTableSchema], + ); + // Build tree data const treeData = useMemo((): TreeNodeData[] => { const data = schemaData?.map(schema => { @@ -378,6 +401,7 @@ const useTreeData = ({ selectStarMap, handleToggle, handleRefreshTables, + refreshTableSchema, errorPayload, }; }; diff --git a/superset-frontend/src/utils/localStorageHelpers.ts b/superset-frontend/src/utils/localStorageHelpers.ts index 61e1bbc3505..3c59b04b9c4 100644 --- a/superset-frontend/src/utils/localStorageHelpers.ts +++ b/superset-frontend/src/utils/localStorageHelpers.ts @@ -51,6 +51,7 @@ export enum LocalStorageKeys { */ SqllabIsAutocompleteEnabled = 'sqllab__is_autocomplete_enabled', SqllabIsRenderHtmlEnabled = 'sqllab__is_render_html_enabled', + SqllabPinnedSchemas = 'sqllab__pinned_schemas', ExploreDataTableOriginalFormattedTimeColumns = 'explore__data_table_original_formatted_time_columns', DashboardCustomFilterBarWidths = 'dashboard__custom_filter_bar_widths', DashboardExploreContext = 'dashboard__explore_context', @@ -71,6 +72,7 @@ export type LocalStorageValues = { homepage_activity_filter: TableTab | null; sqllab__is_autocomplete_enabled: boolean; sqllab__is_render_html_enabled: boolean; + sqllab__pinned_schemas: Record; explore__data_table_original_formatted_time_columns: Record; dashboard__custom_filter_bar_widths: Record; dashboard__explore_context: Record;