/** * Licensed to the Apache Software Foundation (ASF) under one * or more contributor license agreements. See the NOTICE file * distributed with this work for additional information * regarding copyright ownership. The ASF licenses this file * to you under the Apache License, Version 2.0 (the * "License"); you may not use this file except in compliance * with the License. You may obtain a copy of the License at * * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, * software distributed under the License is distributed on an * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY * KIND, either express or implied. See the License for the * specific language governing permissions and limitations * under the License. */ import { useMemo, useReducer, useCallback } from 'react'; import { useDispatch } from 'react-redux'; import { t } from '@apache-superset/core/translation'; import { Table, type TableMetaData, useSchemas, useLazyTablesQuery, useLazyTableMetadataQuery, useLazyTableExtendedMetadataQuery, } from 'src/hooks/apiResources'; import { addDangerToast } from 'src/SqlLab/actions/sqlLab'; import type { TreeNodeData } from './types'; import { SupersetError } from '@superset-ui/core'; export const EMPTY_NODE_ID_PREFIX = 'empty:'; // Reducer state and actions interface TreeDataState { tableData: Record; tableSchemaData: Record; loadingNodes: Record; errorPayload: SupersetError | null; } 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 }; const initialState: TreeDataState = { tableData: {}, tableSchemaData: {}, loadingNodes: {}, errorPayload: null, }; function treeDataReducer( state: TreeDataState, action: TreeDataAction, ): TreeDataState { switch (action.type) { case 'SET_TABLE_DATA': return { ...state, errorPayload: null, tableData: { ...state.tableData, [action.key]: action.data }, }; case 'SET_TABLE_SCHEMA_DATA': return { ...state, tableSchemaData: { ...state.tableSchemaData, [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, loadingNodes: { ...state.loadingNodes, [action.nodeId]: action.loading, }, }; case 'SET_ERROR': return { ...state, errorPayload: action.errorPayload, }; default: return state; } } interface UseTreeDataParams { dbId: number | undefined; catalog: string | null | undefined; pinnedTables: Record; } interface UseTreeDataResult { treeData: TreeNodeData[]; isFetching: boolean; refetch: () => void; loadingNodes: Record; selectStarMap: Record; handleToggle: (id: string, isOpen: boolean) => Promise; handleRefreshTables: (params: { dbId: number; catalog: string | null | undefined; schema: string; }) => void; refreshTableSchema: (id: string) => void; errorPayload: SupersetError | null; } const createEmptyNode = (parentId: string): TreeNodeData => ({ id: `${EMPTY_NODE_ID_PREFIX}${parentId}`, name: t('No items'), type: 'empty', }); const useTreeData = ({ dbId, catalog, pinnedTables, }: UseTreeDataParams): UseTreeDataResult => { const reduxDispatch = useDispatch(); // Schema data from API const { currentData: schemaData, isFetching, refetch, } = useSchemas({ dbId, catalog: catalog || undefined }); // Lazy query hooks const [fetchLazyTables] = useLazyTablesQuery(); const [fetchTableMetadata] = useLazyTableMetadataQuery(); const [fetchTableExtendedMetadata] = useLazyTableExtendedMetadataQuery(); // Combined state for table data, schema data, loading nodes, and data version 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) => { // Only fetch when opening a node if (!isOpen) return; const parts = id.split(':'); const [identifier, databaseId, schema, table] = parts; const parsedDbId = Number(databaseId); if (identifier === 'schema') { const schemaKey = `${parsedDbId}:${schema}`; if (!tableData?.[schemaKey]) { dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: true }); fetchLazyTables( { dbId: parsedDbId, catalog, schema, forceRefresh: false }, true, ) .unwrap() .then(data => { if (data) { dispatch({ type: 'SET_TABLE_DATA', key: schemaKey, data }); } }) .catch(error => { dispatch({ type: 'SET_ERROR', errorPayload: error?.errors?.[0] ?? null, }); }) .finally(() => { dispatch({ type: 'SET_LOADING_NODE', nodeId: id, loading: false, }); }); } } if (identifier === 'table') { const tableKey = `${parsedDbId}:${schema}:${table}`; // Check pinnedTables first (it's stable) if (pinnedTables[tableKey]) return; if (!tableSchemaData[tableKey]) { fetchAndStoreTableSchema(id, true); } } }, [ catalog, fetchAndStoreTableSchema, fetchLazyTables, pinnedTables, tableData, tableSchemaData, ], ); // Force-refresh the table list for a schema and update the tree const handleRefreshTables = useCallback( ({ dbId: refreshDbId, catalog: refreshCatalog, schema, }: { dbId: number; catalog: string | null | undefined; schema: string; }) => { const schemaKey = `${refreshDbId}:${schema}`; const nodeId = `schema:${refreshDbId}:${schema}`; dispatch({ type: 'SET_LOADING_NODE', nodeId, loading: true }); fetchLazyTables({ dbId: refreshDbId, catalog: refreshCatalog, schema, forceRefresh: true, }) .unwrap() .then(data => { dispatch({ type: 'SET_TABLE_DATA', key: schemaKey, data }); }) .catch(error => { dispatch({ type: 'SET_ERROR', errorPayload: error?.errors?.[0] ?? null, }); }) .finally(() => { dispatch({ type: 'SET_LOADING_NODE', nodeId, loading: false }); }); }, [fetchLazyTables], ); const refreshTableSchema = useCallback( (id: string) => { fetchAndStoreTableSchema(id, false); }, [fetchAndStoreTableSchema], ); // Build tree data const treeData = useMemo((): TreeNodeData[] => { const data = schemaData?.map(schema => { const schemaKey = `${dbId}:${schema.value}`; const schemaId = `schema:${dbId}:${schema.value}`; const tablesData = tableData?.[schemaKey]; const tables = tablesData?.options; // Determine children for schema node let schemaChildren: TreeNodeData[]; if (!tablesData) { // Not loaded yet - empty array makes it expandable schemaChildren = []; } else if (tables && tables.length > 0) { // Has tables schemaChildren = tables.map(({ value: tableName, type: tableType }) => { const tableKey = `${dbId}:${schema.value}:${tableName}`; const tableId = `table:${dbId}:${schema.value}:${tableName}`; const columnsData = tableSchemaData[tableKey] ?? pinnedTables[tableKey]; const columns = columnsData?.columns; // Determine children for table node let tableChildren: TreeNodeData[]; if (!columnsData) { // Not loaded yet tableChildren = []; } else if (columns && columns.length > 0) { // Has columns tableChildren = columns.map(col => ({ id: `column:${dbId}:${schema.value}:${tableName}:${col.name}`, name: col.name, type: 'column' as const, columnData: col, })); } else { // Loaded but empty tableChildren = [createEmptyNode(tableId)]; } return { id: tableId, name: tableName, type: 'table' as const, tableType, children: tableChildren, disableCheckbox: !columnsData, }; }); } else { // Loaded but empty schemaChildren = [createEmptyNode(schemaId)]; } return { id: schemaId, name: schema.label, type: 'schema' as const, children: schemaChildren, disableCheckbox: !tablesData, }; }); return data ?? []; }, [dbId, schemaData, tableData, tableSchemaData, pinnedTables]); // Map of tableKey -> selectStar SQL from table metadata const selectStarMap = useMemo(() => { const map: Record = {}; const addEntry = (key: string, meta: TableMetaData | undefined) => { if (meta?.selectStar) { map[key] = meta.selectStar; } }; Object.entries(tableSchemaData).forEach(([key, meta]) => addEntry(key, meta), ); Object.entries(pinnedTables).forEach(([key, meta]) => addEntry(key, meta)); return map; }, [tableSchemaData, pinnedTables]); return { treeData, isFetching, refetch, loadingNodes, selectStarMap, handleToggle, handleRefreshTables, refreshTableSchema, errorPayload, }; }; export default useTreeData;