Dbid unify (#316)

* add database to taskui context

* do not compute dbid in appstate calculation

* pass db id to all downstream checks

* bugfix useappstore

* Speed up metabase state updates

* add db id in one place

* Remove unused functions

* Fix imports

---------

Co-authored-by: Sreejith <1743700+ppsreejith@users.noreply.github.com>

Author: nuwandavek

apps/src/metabase/appState.ts

@@ -128,11 +128,13 @@ export class MetabaseState extends DefaultAppState<MetabaseAppState> {
   actionController = new MetabaseController(this);
 
   private async triggerMetabaseStateUpdate(url: string, elements: DOMQueryMapResponse) {
-    const queryType = await RPCs.getMetabaseState('qb.card.dataset_query.type') as string;
-    const state = this.useStore().getState();
+    const [queryType, dbId, allDBs] = await Promise.all([
+      RPCs.getMetabaseState('qb.card.dataset_query.type') as Promise<string>,
+      getSelectedDbId(),
+      RPCs.getMetabaseState('entities.databases')
+    ]);
+    
     const getState = this.useStore().getState
-    const dbId = await getSelectedDbId();
-    const allDBs = await RPCs.getMetabaseState('entities.databases') as object;
     const minifiedDBs = Object.values(allDBs || {}).map((db: any) => ({ id: db.id, name: db.name }))
 
     let toolEnabledNew = shouldEnable(elements, url);
@@ -460,7 +462,8 @@ Here's what I need modified:
   }
 
   public async getState(): Promise<MetabaseAppState> {
-    return await convertDOMtoState();
+    const currentDBId = this.useStore().getState().toolContext?.dbId || undefined;
+    return await convertDOMtoState(currentDBId);
   }
 
   public async getPlannerConfig() {
@@ -490,7 +493,7 @@ Here's what I need modified:
       }))
       const isCancelled = () => taskStatus.status === 'cancelled';
       const [relevantTables, dbInfo] = await Promise.all([
-        handlePromise(abortable(getRelevantTablesForSelectedDb(), isCancelled), "Failed to get relevant tables", []),
+        handlePromise(abortable(getRelevantTablesForSelectedDb(dbId), isCancelled), "Failed to get relevant tables", []),
         handlePromise(abortable(getDatabaseTablesAndModelsWithoutFields(dbId), isCancelled), "Failed to get database info", DB_INFO_DEFAULT)
       ])
       state.getState().update((oldState) => ({
@@ -505,7 +508,7 @@ Here's what I need modified:
       // Perf caching
       if (!isCancelled() && dbId !== oldDbId) {
         console.log('Running perf caching')
-        processAllMetadata()
+        processAllMetadata(false, dbId)
         getDatabaseInfo(dbId)
       }
     })

apps/src/metabase/helpers/DOMToState.ts

@@ -104,7 +104,7 @@ export type MetabaseAppState = MetabaseAppStateSQLEditor | MetabaseAppStateDashb
 
 // no need to fetch fields since we don't want that in limited entities
 
-export async function convertDOMtoStateSQLQueryV2(pageType: MetabasePageType) : Promise<MetabaseAppStateSQLEditorV2> {
+export async function convertDOMtoStateSQLQueryV2(pageType: MetabasePageType, currentDBId: number) : Promise<MetabaseAppStateSQLEditorV2> {
   const [metabaseUrl, currentCardRaw, outputMarkdown, parameterValues] = await Promise.all([
     RPCs.queryURL(),
     getCurrentCard() as Promise<Card>,
@@ -115,8 +115,8 @@ export async function convertDOMtoStateSQLQueryV2(pageType: MetabasePageType) :
   const metabaseOrigin = new URL(metabaseUrl).origin;
   const isEmbedded = getParsedIframeInfo().isEmbedded
   const sqlQuery =  get(currentCard, 'dataset_query.native.query', '') || ''
-  const limitedEntities = await getLimitedEntities(sqlQuery);
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId
+  const limitedEntities = await getLimitedEntities(sqlQuery, currentDBId);
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined;
   const sqlErrorMessage = await getSqlErrorMessage();
   const appStateType = pageType === 'sql' ? MetabaseAppStateType.SQLEditor : MetabaseAppStateType.RandomPage;
@@ -153,27 +153,27 @@ export async function convertDOMtoStateSQLQueryV2(pageType: MetabasePageType) :
   }
 }
 
-export async function convertDOMtoStateSQLQuery() {
+export async function convertDOMtoStateSQLQuery(currentDBId: number) {
   // CAUTION: This one does not update when changed via ui for some reason
   // const dbId = _.get(hashMetadata, 'dataset_query.database');
   const fullUrl = await RPCs.queryURL();
   const url = new URL(fullUrl).origin;
   const availableDatabases = (await getDatabases())?.data?.map(({ name }) => name);
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId;
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined;
   const defaultSchema = selectedDatabaseInfo?.default_schema;
-  let sqlQuery = await getCurrentQuery()
-  const appSettings = RPCs.getAppSettings()
-  const cache = RPCs.getCache()
-  const sqlTables = getTablesFromSqlRegex(sqlQuery)
+  let sqlQuery = await getCurrentQuery();
+  const appSettings = RPCs.getAppSettings();
+  const cache = RPCs.getCache();
+  const sqlTables = getTablesFromSqlRegex(sqlQuery);
   if (defaultSchema) {
     sqlTables.forEach((table) => {
       if (table.schema === undefined || table.schema === '') {
         table.schema = defaultSchema
       }
     })
   }
-  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, false, sqlTables, [])
+  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, false, sqlTables, [], dbId);
   // add defaultSchema back to relevantTablesWithFields. kind of hacky but whatever
   relevantTablesWithFields = relevantTablesWithFields.map(table => {
     if (table.schema === undefined || table.schema === '') {
@@ -241,13 +241,13 @@ export async function convertDOMtoStateSQLQuery() {
   return metabaseAppStateSQLEditor;
 }
 
-export async function convertDOMtoStateMBQLQuery() {
-    return await getMBQLAppState() as MetabaseAppStateMBQLEditor;
+export async function convertDOMtoStateMBQLQuery(currentDBId: number): Promise<MetabaseAppStateMBQLEditor> {
+    return await getMBQLAppState(currentDBId) as MetabaseAppStateMBQLEditor;
 }
 
 // check if on dashboard page
-export async function convertDOMtoStateDashboard(): Promise<MetabaseAppStateDashboard> {
-    const dashboardInfo = await getDashboardAppState();
+export async function convertDOMtoStateDashboard(currentDBId: number): Promise<MetabaseAppStateDashboard> {
+    const dashboardInfo = await getDashboardAppState(currentDBId);
     return dashboardInfo as MetabaseAppStateDashboard;
 };
 
@@ -272,21 +272,21 @@ export async function semanticQueryState() {
   return metabaseSemanticQueryAppState;
 }
 
-export async function convertDOMtoState() {
+export async function convertDOMtoState(currentDBId: number) {
   const url = await queryURL();
   const qlType = await getQLType();
   const metabasePageType = determineMetabasePageType({}, url, qlType);
   if (metabasePageType === 'dashboard') {
-    return await convertDOMtoStateDashboard();
+    return await convertDOMtoStateDashboard(currentDBId);
   }
   if (metabasePageType === 'mbql') {
-    return await convertDOMtoStateMBQLQuery();
+    return await convertDOMtoStateMBQLQuery(currentDBId);
   }
   const appSettings = RPCs.getAppSettings()
   if (appSettings.useV2States && appSettings.analystMode) {
-    return await convertDOMtoStateSQLQueryV2(metabasePageType);
+    return await convertDOMtoStateSQLQueryV2(metabasePageType, currentDBId);
   }
-  return await convertDOMtoStateSQLQuery();
+  return await convertDOMtoStateSQLQuery(currentDBId);
 }
 
 async function getSqlVariables() {

apps/src/metabase/helpers/dashboard/appState.ts

@@ -1,7 +1,6 @@
 import { DashboardInfo, DashboardMetabaseState } from './types';
 import _, { forEach, reduce, template, values } from 'lodash';
 import { MetabaseAppStateDashboard} from '../DOMToState';
-import { getTablesWithFields } from '../getDatabaseSchema';
 import { getAllRelevantModelsForSelectedDb, getDatabaseInfo, getFieldResolvedName } from '../metabaseAPIHelpers';
 import { getDashboardState, getSelectedDbId } from '../metabaseStateAPI';
 import { getParsedIframeInfo, RPCs } from 'web';
@@ -332,14 +331,11 @@ async function substituteParameters(
   return sql;
 };
 
-export async function getDashboardAppState(): Promise<MetabaseAppStateDashboard | null> {
+export async function getDashboardAppState(currentDBId: number): Promise<MetabaseAppStateDashboard | null> {
   const fullUrl = await RPCs.queryURL();
   const url = new URL(fullUrl).origin;
-  const appSettings = RPCs.getAppSettings();
-  const selectedCatalog = get(find(appSettings.availableCatalogs, { name: appSettings.selectedCatalog }), 'content')
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined
-  const defaultSchema = selectedDatabaseInfo?.default_schema;
   const dashboardMetabaseState: DashboardMetabaseState = await getDashboardState() as DashboardMetabaseState;
   if (!dashboardMetabaseState || !dashboardMetabaseState.dashboards || !dashboardMetabaseState.dashboardId) {
     console.warn('Could not get dashboard info');
@@ -364,12 +360,14 @@ export async function getDashboardAppState(): Promise<MetabaseAppStateDashboard
     const limitedEntitiesSQL = await getLimitedEntitiesFromQueries(
         filteredCards.flatMap(card => 
             card?.dataset_query?.native?.query ? [card.dataset_query.native.query] : []
-        )
+        ),
+        dbId
     );
     const limitedEntitiesMBQL = await getLimitedEntitiesFromMBQLQueries(
         filteredCards.flatMap(card => 
             card?.dataset_query?.query ? [card.dataset_query.query] : []
-        )
+        ),
+        dbId
     );
     const limitedEntities = [...limitedEntitiesSQL, ...limitedEntitiesMBQL];
     // remove duplicates based on id and type

apps/src/metabase/helpers/getDatabaseSchema.ts

@@ -1,7 +1,6 @@
 import { FormattedTable } from './types';
 import { getTablesFromSqlRegex, TableAndSchema } from './parseSql';
 import _ from 'lodash';
-import { getSelectedDbId } from './metabaseStateAPI';
 import { getUserTables, searchUserQueries, getDatabaseTablesAndModelsWithoutFields } from './metabaseAPIHelpers';
 import { applyTableDiffs, handlePromise } from '../../common/utils';
 import { TableDiff } from 'web/types';
@@ -102,8 +101,8 @@ export const searchTables = async (userId: number, dbId: number, query: string):
   return dedupedTables
 }
 
-export const getTablesWithFields = async (tableDiff?: TableDiff, drMode = false, isCatalogSelected: boolean = false, sqlTables: TableAndSchema[] = [], mbqlTableIds: number[] = []) => {
-  const dbId = await getSelectedDbId();
+export const getTablesWithFields = async (tableDiff?: TableDiff, drMode = false, isCatalogSelected: boolean = false, sqlTables: TableAndSchema[] = [], mbqlTableIds: number[] = [], currentDBId?: number) => {
+  const dbId = currentDBId;
   if (!dbId) {
     console.warn("[minusx] No database selected when getting tables with fields");
     return [];
@@ -128,8 +127,8 @@ export const getTablesWithFields = async (tableDiff?: TableDiff, drMode = false,
 
 
 
-export const getRelevantTablesForSelectedDb = async (): Promise<FormattedTable[]> => {
-  const dbId = await getSelectedDbId();
+export const getRelevantTablesForSelectedDb = async (currentDBId?: number): Promise<FormattedTable[]> => {
+  const dbId = currentDBId;
   if (!dbId) {
     console.warn("[minusx] No database selected when getting relevant tables");
     return [];

apps/src/metabase/helpers/mbql/appState.ts

@@ -11,13 +11,13 @@ import { getAndFormatOutputTable } from '../operations';
 import { MetabaseAppStateType } from '../analystModeTypes';
 
 
-export async function getMBQLAppState(): Promise<MetabaseAppStateMBQLEditor | null> {
+export async function getMBQLAppState(currentDBId: number): Promise<MetabaseAppStateMBQLEditor | null> {
   const fullUrl = await RPCs.queryURL();
   const url = new URL(fullUrl).origin;
 
   const appSettings = RPCs.getAppSettings();
   const selectedCatalog = get(find(appSettings.availableCatalogs, { name: appSettings.selectedCatalog }), 'content')
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined
   const defaultSchema = selectedDatabaseInfo?.default_schema;
   const mbqlState = await getMBQLState();
@@ -32,7 +32,7 @@ export async function getMBQLAppState(): Promise<MetabaseAppStateMBQLEditor | nu
   const relevantModels = await getSelectedAndRelevantModels('', appSettings.selectedModels, allModels, sourceTableModelIds)
   const relevantModelsWithFields = await getModelsWithFields(relevantModels)
 
-  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, [], sourceTableModelIds)
+  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, [], sourceTableModelIds, dbId);
   relevantTablesWithFields = relevantTablesWithFields.map(table => {
     if (table.schema === undefined || table.schema === '') {
       table.schema = defaultSchema || 'unknown'

apps/src/metabase/helpers/metabaseAPIHelpers.ts

@@ -200,16 +200,16 @@ type ResultMetadataElm = {
   field_ref: [string, string, { "base-type": string }];
 };
 
-export async function getAllCardsAndModels(forceRefresh = false) {
+export async function getAllCardsAndModels(forceRefresh = false, currentDBId: number) {
   const cards = await handlePromise(
     forceRefresh ? fetchCards.refresh({}) : fetchCards({}),
     "[minusx] Error getting all cards",
     []
   );
 
   // Get selected database ID
-  const selectedDbId = await getSelectedDbId();
-  
+  const selectedDbId = currentDBId;
+
   console.log('[minusx] getAllCards - Total cards:', cards.length);
 
   // Filter cards by database_id only
@@ -299,14 +299,9 @@ export async function getAllCardsAndModels(forceRefresh = false) {
   return { cards: processedCards, tables: relevantTables, modelFields: processedModelFields };
 }
 
-export async function getAllCardsLegacy() {
-  const result = await getAllCardsAndModels();
-  return result.cards;
-}
-
-export async function getAllFields() {
+export async function getAllFields(currentDBId: number) {
   // Get selected database ID
-  const selectedDbId = await getSelectedDbId();
+  const selectedDbId = currentDBId;
 
   if (!selectedDbId) {
     console.log('[minusx] getAllFields - No database selected');
@@ -329,14 +324,14 @@ export async function getAllFields() {
   return limitedFields;
 }
 
-export async function getAllFieldsFiltered(tableNames: string[]) {
+export async function getAllFieldsFiltered(tableNames: string[], currentDBId: number) {
   if (!tableNames || tableNames.length === 0) {
     console.log('[minusx] getAllFieldsFiltered - No table names provided, returning empty array');
     return [];
   }
 
   // Get selected database ID
-  const selectedDbId = await getSelectedDbId();
+  const selectedDbId = currentDBId
 
   if (!selectedDbId) {
     console.log('[minusx] getAllFieldsFiltered - No database selected');

apps/src/metabase/helpers/utils.ts

@@ -52,23 +52,23 @@ export function determineMetabasePageType(elements: DOMQueryMapResponse, url: st
     return 'unknown';
 }
 
-export async function getLimitedEntitiesFromMBQLQueries(mbqlQueries: any): Promise<MetabaseTableOrModel[]> {
+export async function getLimitedEntitiesFromMBQLQueries(mbqlQueries: any, currentDBId: number): Promise<MetabaseTableOrModel[]> {
     const entities : MetabaseTableOrModel[] = [];
     for (const mbqlQuery of mbqlQueries) {
-        const limitedEntities = await getLimitedMBQLEntities(mbqlQuery);
+        const limitedEntities = await getLimitedMBQLEntities(mbqlQuery, currentDBId);
         entities.push(...limitedEntities);
     }
     // Remove duplicates based on id and type
     const uniqueEntities = Array.from(new Map(entities.map(entity => [entity.id, entity])).values());
     return uniqueEntities;
 }
 
-async function getLimitedMBQLEntities(mbqlQuery: any): Promise<MetabaseTableOrModel[]> {
+async function getLimitedMBQLEntities(mbqlQuery: any, currentDBId: number): Promise<MetabaseTableOrModel[]> {
     const appSettings = RPCs.getAppSettings();
   if (!appSettings.analystMode || !appSettings.manuallyLimitContext) {
     return [];
   }
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId;
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined
   const defaultSchema = selectedDatabaseInfo?.default_schema;
   const selectedCatalogObj = find(appSettings.availableCatalogs, { name: appSettings.selectedCatalog });
@@ -80,7 +80,7 @@ async function getLimitedMBQLEntities(mbqlQuery: any): Promise<MetabaseTableOrMo
     const relevantModels = await getSelectedAndRelevantModels('', appSettings.selectedModels, allModels, sourceTableModelIds)
     const relevantModelsWithFields = await getModelsWithFields(relevantModels)
 
-    let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, [], sourceTableModelIds)
+    let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, [], sourceTableModelIds, dbId)
     relevantTablesWithFields = relevantTablesWithFields.map(table => {
       if (table.schema === undefined || table.schema === '') {
         table.schema = defaultSchema || 'unknown'
@@ -97,26 +97,26 @@ async function getLimitedMBQLEntities(mbqlQuery: any): Promise<MetabaseTableOrMo
 }
 
 
-export async function getLimitedEntitiesFromQueries(sqlQueries: string[]): Promise<MetabaseTableOrModel[]> {
+export async function getLimitedEntitiesFromQueries(sqlQueries: string[], currentDBId: number): Promise<MetabaseTableOrModel[]> {
     const entities : MetabaseTableOrModel[] = [];
     for (const sqlQuery of sqlQueries) {
-        const limitedEntities = await getLimitedEntities(sqlQuery);
+        const limitedEntities = await getLimitedEntities(sqlQuery, currentDBId);
         entities.push(...limitedEntities);
     }
     // Remove duplicates based on id and type
     const uniqueEntities = Array.from(new Map(entities.map(entity => [entity.id, entity])).values());
     return uniqueEntities;
 }
 
-export async function getLimitedEntities(sqlQuery: string): Promise<MetabaseTableOrModel[]> {
+export async function getLimitedEntities(sqlQuery: string, currentDBId: number): Promise<MetabaseTableOrModel[]> {
   const appSettings = RPCs.getAppSettings();
 
   // Early return if conditions not met
   if (!appSettings.analystMode || !appSettings.manuallyLimitContext) {
     return [];
   }
 
-  const dbId = await getSelectedDbId();
+  const dbId = currentDBId;
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined;
   const defaultSchema = selectedDatabaseInfo?.default_schema;
 
@@ -133,7 +133,7 @@ export async function getLimitedEntities(sqlQuery: string): Promise<MetabaseTabl
     });
   }
 
-  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, sqlTables, []);
+  let relevantTablesWithFields = await getTablesWithFields(appSettings.tableDiff, appSettings.drMode, !!selectedCatalog, sqlTables, [], dbId);
 
   // Add defaultSchema back to relevantTablesWithFields
   relevantTablesWithFields = relevantTablesWithFields.map(table => {

apps/src/package.ts

@@ -1,7 +1,7 @@
 export { getAppStateConfigs } from "./appStateConfigs";
 export { applyTableDiffs } from "./common/utils";
 export { getTableContextYAML, filterTablesByCatalog } from "./metabase/helpers/catalog";
-export { getTableData, getDatabaseTablesAndModelsWithoutFields, getAllCardsAndModels as getAllCards, getAllCardsLegacy, getAllFields, getAllFieldsFiltered } from "./metabase/helpers/metabaseAPIHelpers";
+export { getAllCardsAndModels, getTableData, getDatabaseTablesAndModelsWithoutFields, getAllCardsAndModels as getAllCards, getAllFields, getAllFieldsFiltered } from "./metabase/helpers/metabaseAPIHelpers";
 export { fetchModelInfo } from "./metabase/helpers/metabaseAPI";
 export { getAllTemplateTagsInQuery, applySQLEdits, type SQLEdits } from "./metabase/helpers/sqlQuery";
 export { getModelsWithFields, getSelectedAndRelevantModels, modifySqlForMetabaseModels, replaceLLMFriendlyIdentifiersInSqlWithModels } from "./metabase/helpers/metabaseModels";