add table and field info for tables in the sql query to user context (#290)

Author: kektobiologist

apps/src/metabase/helpers/DOMToState.ts

@@ -1,6 +1,6 @@
 import { getParsedIframeInfo, RPCs } from 'web'
-import { getTablesWithFields } from './getDatabaseSchema';
-import { getAllRelevantModelsForSelectedDb, getDatabaseInfo, getDatabases } from './metabaseAPIHelpers';
+import { getAllRelevantTablesForSelectedDb, getRelevantTablesForSelectedDb, getTablesWithFields, validateTablesInDB } from './getDatabaseSchema';
+import { getAllRelevantModelsForSelectedDb, getDatabaseInfo, getDatabases, getTableData } from './metabaseAPIHelpers';
 import { getAndFormatOutputTable, getSqlErrorMessage } from './operations';
 import { isDashboardPageUrl } from './dashboard/util';
 import { DashboardInfo } from './dashboard/types';
@@ -106,7 +106,8 @@ export { type MetabasePageType } from '../defaultState'
 
 export type MetabaseAppState = MetabaseAppStateSQLEditor | MetabaseAppStateDashboard | MetabaseSemanticQueryAppState | MetabaseAppStateMBQLEditor;
 
-async function getRelevantEntitiesWithFields(sqlQuery: string): Promise<MetabaseTableOrModel[]> {
+// no need to fetch fields since we don't want that in limited entities
+async function getLimitedEntities(sqlQuery: string): Promise<MetabaseTableOrModel[]> {
   const appSettings = RPCs.getAppSettings();
 
   // Early return if conditions not met
@@ -143,7 +144,6 @@ async function getRelevantEntitiesWithFields(sqlQuery: string): Promise<Metabase
 
   const allModels = dbId ? await getAllRelevantModelsForSelectedDb(dbId) : [];
   const relevantModels = await getSelectedAndRelevantModels(sqlQuery || "", appSettings.selectedModels, allModels);
-  const relevantModelsWithFields = await getModelsWithFields(relevantModels);
 
   // Transform and combine tables and models with type annotations
   const relevantTablesWithFieldsAndType: MetabaseTableOrModel[] = relevantTablesWithFields.map(table => ({
@@ -154,7 +154,7 @@ async function getRelevantEntitiesWithFields(sqlQuery: string): Promise<Metabase
     description: table.description,
   }));
 
-  const relevantModelsWithFieldsAndType: MetabaseTableOrModel[] = relevantModelsWithFields.map(model => ({
+  const relevantModelsWithFieldsAndType: MetabaseTableOrModel[] = relevantModels.map(model => ({
     type: 'model',
     id: model.modelId || 0,
     name: model.name,
@@ -174,13 +174,29 @@ export async function convertDOMtoStateSQLQueryV2() : Promise<MetabaseAppStateSQ
   const currentCard = processCard(currentCardRaw);
   const metabaseOrigin = new URL(metabaseUrl).origin;
   const isEmbedded = getParsedIframeInfo().isEmbedded
-  const limitedEntities = await getRelevantEntitiesWithFields(
-    get(currentCard, 'dataset_query.native.query', '') || ''
-  );
+  const sqlQuery =  get(currentCard, 'dataset_query.native.query', '') || ''
+  const limitedEntities = await getLimitedEntities(sqlQuery);
   const dbId = await getSelectedDbId();
   const selectedDatabaseInfo = dbId ? await getDatabaseInfo(dbId) : undefined;
   const sqlErrorMessage = await getSqlErrorMessage();
-
+  // add tables in the sql as relevant tables, after fetching their fields
+  let relevantTablesWithFields: FormattedTable[] = []
+  {
+      const dbTables = await getAllRelevantTablesForSelectedDb(dbId || 0)
+      const sqlTables = await getTablesFromSqlRegex(sqlQuery)
+      const defaultSchema = selectedDatabaseInfo?.default_schema
+      // Apply default schema to tables if needed
+      if (defaultSchema) {
+        sqlTables.forEach((table) => {
+          if (table.schema === undefined || table.schema === '') {
+            table.schema = defaultSchema;
+          }
+        });
+      }
+      const validSqlTables = validateTablesInDB(sqlTables, dbTables, defaultSchema)
+      const sqlTablesWithFields = (await Promise.all(validSqlTables.map(table => getTableData(table.id)))).filter(table => table !== "missing")
+      relevantTablesWithFields = sqlTablesWithFields
+  }
   return {
     type: MetabaseAppStateType.SQLEditor,
     version: '2',
@@ -192,7 +208,8 @@ export async function convertDOMtoStateSQLQueryV2() : Promise<MetabaseAppStateSQ
     outputMarkdown,
     parameterValues,
     selectedDatabaseInfo,
-    sqlErrorMessage
+    sqlErrorMessage,
+    relevantTablesWithFields
   }
 }
 

apps/src/metabase/helpers/analystModeTypes.ts

@@ -28,6 +28,7 @@ export interface MetabaseAppStateSQLEditorV2 extends MetabaseAppStateBase {
   outputMarkdown: string
   parameterValues: ParameterValues
   sqlErrorMessage?: string
+  relevantTablesWithFields?: FormattedTable[]
 }
 
 export interface MetabaseAppStateDashboardV2 extends MetabaseAppStateBase {

apps/src/metabase/helpers/getDatabaseSchema.ts

@@ -37,7 +37,7 @@ function lowerAndDefaultSchemaAndDedupe(tables: TableAndSchema[]): TableAndSchem
   return dedupeAndCountTables(lowered);
 }
 
-const validateTablesInDB = (tables: TableAndSchema[], allDBTables: FormattedTable[], default_schema?: string) => {
+export const validateTablesInDB = (tables: TableAndSchema[], allDBTables: FormattedTable[], default_schema?: string) => {
   const allTablesAsMap = _.fromPairs(allDBTables.map(tableInfo => [getTableKey(tableInfo), tableInfo]));
   if (default_schema) {
     tables = tables.map(tableInfo => {
@@ -66,7 +66,7 @@ const addTableJoins = (tables: FormattedTable[], tableMap: Record<number | strin
   })
 }
 
-const getAllRelevantTablesForSelectedDb = async (dbId: number): Promise<FormattedTable[]> => {
+export const getAllRelevantTablesForSelectedDb = async (dbId: number): Promise<FormattedTable[]> => {
   const [userTables, {tables: allDBTables, default_schema}] = await Promise.all([
     getUserTables(dbId),
     handlePromise(getDatabaseTablesAndModelsWithoutFields(dbId), "Failed to get database tables", {