Skip to content

Commit a7c535e

Browse files
ppsreejithppsreejith
authored
Feature/improved parameters filters tags (#277)
* Give current card context in tool response & input * Create V2 of Metabase SQL Editor State * Add toggle to flip state processing to V2 * Fix EditQuery action * Add SetQueryParameterValues tool * Execute query upon setting parameter values
1 parent b3d17b1 commit a7c535e

File tree

11 files changed

+201
-68
lines changed

11 files changed

+201
-68
lines changed

‎apps/src/metabase/actionDescriptions.ts‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,32 @@ export const ACTION_DESCRIPTIONS_PLANNER: ActionDescription[] = [
8181
description: `Executes the SQL query in the metabase SQL editor. This also sets the "queryExecuted" state to true after execution.
8282
`,
8383
},
84+
{
85+
name: 'setQueryParameterValues',
86+
args: {
87+
parameterValues: {
88+
type: 'array',
89+
items: {
90+
type: 'object',
91+
properties: {
92+
id: {
93+
type: 'string',
94+
description: "The id of the parameter to set the value for."
95+
},
96+
value: {
97+
type: 'array',
98+
items: {
99+
type: 'string'
100+
},
101+
description: "The value to set for the parameter."
102+
}
103+
}
104+
},
105+
description: "The parameters to set in the SQL editor."
106+
}
107+
},
108+
description: "Sets the values of the parameters in the SQL editor. This is useful for setting the values of variables in the SQL query. The parameterValues should be an array of objects, each containing the id of the parameter and the value to set.",
109+
},
84110
{
85111
name: 'setVisualizationType',
86112
args: {

‎apps/src/metabase/appController.ts‎

Lines changed: 29 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import {
77
MetabaseAppStateSQLEditor,
88
MetabaseSemanticQueryAppState,
99
MetabaseAppStateMBQLEditor,
10-
MetabaseAppStateType,
1110
MetabasePageType,
1211
} from "./helpers/DOMToState";
1312
import {
@@ -30,6 +29,7 @@ import {
3029
primaryVisualizationTypes,
3130
Card,
3231
toLowerVisualizationType,
32+
ParameterValues,
3333
} from "./helpers/types";
3434
import {
3535
getTemplateTags as getTemplateTagsForVars,
@@ -41,7 +41,7 @@ import {
4141
SQLEdits
4242
} from "./helpers/sqlQuery";
4343
import axios from 'axios'
44-
import { getSelectedDbId, getCurrentUserInfo as getUserInfo, getSnippets, getCurrentCard, getDashboardState } from "./helpers/metabaseStateAPI";
44+
import { getSelectedDbId, getCurrentUserInfo as getUserInfo, getSnippets, getCurrentCard, getDashboardState, getCurrentQuery, getParameterValues } from "./helpers/metabaseStateAPI";
4545
import { runSQLQueryFromDashboard } from "./helpers/dashboard/runSqlQueryFromDashboard";
4646
import { getAllRelevantModelsForSelectedDb, getTableData } from "./helpers/metabaseAPIHelpers";
4747
import { processSQLWithCtesOrModels, dispatch, updateIsDevToolsOpen, updateDevToolsTabName, addMemory } from "web";
@@ -360,6 +360,22 @@ export class MetabaseController extends AppController<MetabaseAppState> {
360360
return await this._executeSQLQueryInternal();
361361
}
362362

363+
@Action({
364+
labelRunning: "Setting parameter values for a query",
365+
labelDone: "Parameter values set",
366+
labelTask: "Parameter values set",
367+
description: "Sets parameter values for a query in the Metabase SQL editor and execute.",
368+
renderBody: ({ parameterValues }: { parameterValues: Array<{id: string, value: string[]}> }) => {
369+
return {text: null, code: JSON.stringify({ parameterValues })}
370+
}
371+
})
372+
async setQueryParameterValues({ parameterValues }: { parameterValues: Array<{id: string, value: string[]}> }) {
373+
await Promise.all(parameterValues.map(async ({id, value}) => {
374+
return RPCs.dispatchMetabaseAction('metabase/qb/SET_PARAMETER_VALUE', { id, value });
375+
}));
376+
return this._executeSQLQueryInternal()
377+
}
378+
363379
@Action({
364380
labelRunning: "Executes the SQL query with parameters",
365381
labelDone: "Executed query",
@@ -422,8 +438,7 @@ export class MetabaseController extends AppController<MetabaseAppState> {
422438
}
423439
})
424440
async EditAndExecuteQuery({ sql_edits, _ctes = [], explanation = "", template_tags={}, parameters=[] }: { sql_edits: SQLEdits, _ctes?: CTE[], explanation?: string, template_tags?: object, parameters?: any[] }) {
425-
const appState = (await this.app.getState()) as MetabaseAppStateSQLEditor;
426-
let sql = appState.sqlQuery || "";
441+
let sql = await getCurrentQuery() || ""
427442
sql = applySQLEdits(sql, sql_edits);
428443
return await this.ExecuteQuery({ sql, _ctes, explanation, template_tags, parameters });
429444
}
@@ -927,14 +942,20 @@ export class MetabaseController extends AppController<MetabaseAppState> {
927942
};
928943
await this.uClick({ query: "run_query" });
929944
await waitForQueryExecution();
945+
const [currentCard, currentParameterValues] = await Promise.all([
946+
getCurrentCard(),
947+
getParameterValues()
948+
]) as [Card, ParameterValues];
949+
const cardState = `<CURRENT_CARD>${JSON.stringify(currentCard)}</CURRENT_CARD>`
950+
const parameterValuesState = `<CURRENT_PARAMETER_VALUES>${JSON.stringify(currentParameterValues)}</CURRENT_PARAMETER_VALUES>`;
930951
const sqlErrorMessage = await getSqlErrorMessage();
931952
if (sqlErrorMessage) {
932-
actionContent.content = `<ERROR>${sqlErrorMessage}</ERROR>`;
953+
actionContent.content = `${cardState}${parameterValuesState}<ERROR>${sqlErrorMessage}</ERROR>`;
933954
} else {
934955
// table output
935-
let tableOutput = ""
936-
tableOutput = await getAndFormatOutputTable(_type);
937-
actionContent.content = tableOutput;
956+
let output = ""
957+
output = await getAndFormatOutputTable(_type);
958+
actionContent.content = `${cardState}${parameterValuesState}<OUTPUT>${output}<OUTPUT>`;
938959
}
939960
return actionContent;
940961
}

‎apps/src/metabase/helpers/DOMToState.ts‎

Lines changed: 33 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import { canUseModelsModeForCatalog } from '../../../../web/src/helpers/catalogA
1818
import { getMBQLAppState } from './mbql/appState';
1919
import { isMBQLPageUrl, MBQLInfo } from './mbql/utils';
2020
import { getModelsWithFields, getSelectedAndRelevantModels, modifySqlForMetabaseModels} from './metabaseModels';
21+
import { MetabaseAppStateSQLEditorV2, MetabaseAppStateType, processCard } from './analystModeTypes';
2122

2223
const {modifySqlForMxModels} = catalogAsModels
2324

@@ -40,12 +41,6 @@ interface ExtractedTable {
4041
id: number;
4142
}
4243

43-
export enum MetabaseAppStateType {
44-
SQLEditor = 'metabaseSQLEditor',
45-
Dashboard = 'metabaseDashboard',
46-
SemanticQuery = 'metabaseSemanticQuery',
47-
MBQLEditor = 'metabaseMBQLEditor'
48-
}
4944
export interface MetabaseAppStateSQLEditor {
5045
type: MetabaseAppStateType.SQLEditor;
5146
availableDatabases?: string[];
@@ -71,6 +66,7 @@ export interface MetabaseAppStateSQLEditor {
7166
metabaseUrl?: string;
7267
isEmbedded: boolean;
7368
relevantEntitiesWithFields?: FormattedTable[];
69+
currentCard?: Card;
7470
}
7571

7672
// make this DashboardInfo
@@ -108,6 +104,30 @@ export { type MetabasePageType } from '../defaultState'
108104

109105
export type MetabaseAppState = MetabaseAppStateSQLEditor | MetabaseAppStateDashboard | MetabaseSemanticQueryAppState | MetabaseAppStateMBQLEditor;
110106

107+
export async function convertDOMtoStateSQLQueryV2() : Promise<MetabaseAppStateSQLEditorV2> {
108+
const [metabaseUrl, currentCardRaw, outputMarkdown, parameterValues] = await Promise.all([
109+
RPCs.queryURL(),
110+
getCurrentCard() as Promise<Card>,
111+
getAndFormatOutputTable(),
112+
getParameterValues() as Promise<ParameterValues>
113+
]);
114+
const currentCard = processCard(currentCardRaw);
115+
const metabaseOrigin = new URL(metabaseUrl).origin;
116+
const isEmbedded = getParsedIframeInfo().isEmbedded
117+
const relevantEntitiesWithFields: FormattedTable[] = []
118+
return {
119+
type: MetabaseAppStateType.SQLEditor,
120+
version: '2',
121+
metabaseOrigin,
122+
metabaseUrl,
123+
isEmbedded,
124+
relevantEntitiesWithFields,
125+
currentCard,
126+
outputMarkdown,
127+
parameterValues,
128+
}
129+
}
130+
111131
export async function convertDOMtoStateSQLQuery() {
112132
// CAUTION: This one does not update when changed via ui for some reason
113133
// const dbId = _.get(hashMetadata, 'dataset_query.database');
@@ -153,12 +173,14 @@ export async function convertDOMtoStateSQLQuery() {
153173
const vizType = await getVisualizationType()
154174
const visualizationSettings = await getVisualizationSettings() as visualizationSettings
155175
const sqlVariables = await getSqlVariables();
176+
const currentCard = await getCurrentCard() as Card;
156177
const metabaseAppStateSQLEditor: MetabaseAppStateSQLEditor = {
157178
type: MetabaseAppStateType.SQLEditor,
158179
availableDatabases,
159180
selectedDatabaseInfo,
160181
relevantTables: relevantTablesWithFields,
161182
sqlQuery: sqlQuery || '',
183+
currentCard,
162184
queryExecuted,
163185
sqlEditorState: nativeEditorOpen ? 'open' : 'closed',
164186
visualizationType: showingRawTable ? 'table' : vizType,
@@ -244,8 +266,13 @@ export async function convertDOMtoState() {
244266
// if(appSettings.semanticPlanner) {
245267
// return await semanticQueryState();
246268
// }
269+
const appSettings = RPCs.getAppSettings()
270+
if (appSettings.useV2States) {
271+
return await convertDOMtoStateSQLQueryV2();
272+
}
247273
return await convertDOMtoStateSQLQuery();
248274
}
275+
249276
async function getSqlVariables() {
250277
const currentCard = await getCurrentCard() as Card;
251278
if (!currentCard) {

‎apps/src/metabase/helpers/analystModeTypes.ts‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
import { omit, pick } from "lodash";
2+
import { DashboardInfo } from "./dashboard/types";
3+
import { Card, FormattedTable, ParameterValues } from "./types";
4+
5+
export enum MetabaseAppStateType {
6+
SQLEditor = 'metabaseSQLEditor',
7+
Dashboard = 'metabaseDashboard',
8+
SemanticQuery = 'metabaseSemanticQuery',
9+
MBQLEditor = 'metabaseMBQLEditor'
10+
}
11+
12+
interface MetabaseAppStateBase {
13+
type: string
14+
version: string
15+
metabaseOrigin: string;
16+
metabaseUrl: string;
17+
isEmbedded: boolean;
18+
relevantEntitiesWithFields: FormattedTable[];
19+
}
20+
21+
export interface MetabaseAppStateSQLEditorV2 extends MetabaseAppStateBase {
22+
type: MetabaseAppStateType.SQLEditor
23+
version: '2'
24+
currentCard: Card
25+
outputMarkdown: string
26+
parameterValues: ParameterValues
27+
}
28+
29+
export interface MetabaseAppStateDashboardV2 extends MetabaseAppStateBase {
30+
// state.dashboard's dashboards[dashboardId], remove unused fields
31+
// state.dashboard's parameterValues
32+
// state.dashboard's dashcards[selected card ids]'s processed card, parameter_mappings, size_x, size_y, row, col, created_at, outputMarkdown
33+
currentDashboard: DashboardInfo
34+
parameterValues: ParameterValues
35+
}
36+
37+
const fieldsToRemove = [
38+
'cache_invalidated_at',
39+
'archived',
40+
'collection_position',
41+
'source_card_id',
42+
'result_metadata',
43+
'creator',
44+
'initially_published_at',
45+
'enable_embedding',
46+
'collection_id',
47+
'made_public_by_id',
48+
'embedding_params',
49+
'cache_ttl',
50+
'archived_directly',
51+
'collection_preview'
52+
];
53+
54+
export function processCard(card: Card): Card {
55+
// Remove unused fields from the card
56+
const cleanCard: any = omit(card, fieldsToRemove);
57+
58+
// Process nested fields - keep only specific properties
59+
if (cleanCard['last-edit-info']) {
60+
cleanCard['last-edit-info'] = pick(cleanCard['last-edit-info'], ['timestamp']);
61+
}
62+
63+
if (cleanCard.collection) {
64+
cleanCard.collection = pick(cleanCard.collection, ['slug']);
65+
}
66+
67+
return cleanCard;
68+
}

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { DashboardInfo, DashboardMetabaseState } from './types';
22
import _, { forEach, reduce, template, values } from 'lodash';
3-
import { MetabaseAppStateDashboard, MetabaseAppStateType} from '../DOMToState';
3+
import { MetabaseAppStateDashboard} from '../DOMToState';
44
import { getTablesWithFields } from '../getDatabaseSchema';
55
import { getAllRelevantModelsForSelectedDb, getDatabaseInfo, getFieldResolvedName } from '../metabaseAPIHelpers';
66
import { getDashboardState, getSelectedDbId } from '../metabaseStateAPI';
@@ -11,6 +11,7 @@ import { find, get } from 'lodash';
1111
import { getTablesFromSqlRegex, TableAndSchema } from '../parseSql';
1212
import { getTableContextYAML } from '../catalog';
1313
import { getModelsFromSql, getModelsWithFields, modifySqlForMetabaseModels, replaceLLMFriendlyIdentifiersInSqlWithModels } from '../metabaseModels';
14+
import { MetabaseAppStateType } from '../analystModeTypes';
1415

1516
// Removed: const { getMetabaseState } = RPCs - using centralized state functions instead
1617

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
import { MetabaseAppStateMBQLEditor, MetabaseAppStateType} from '../DOMToState';
1+
import { MetabaseAppStateMBQLEditor} from '../DOMToState';
22
import { getParsedIframeInfo, RPCs } from 'web';
33
import { MBQLInfo, getSourceTableIds } from './utils';
44
import { getMBQLState, getSelectedDbId } from '../metabaseStateAPI';
@@ -8,6 +8,7 @@ import { getTableContextYAML } from '../catalog';
88
import { getTablesWithFields } from '../getDatabaseSchema';
99
import { getModelsWithFields, getSelectedAndRelevantModels } from '../metabaseModels';
1010
import { getAndFormatOutputTable } from '../operations';
11+
import { MetabaseAppStateType } from '../analystModeTypes';
1112

1213

1314
export async function getMBQLAppState(): Promise<MetabaseAppStateMBQLEditor | null> {

‎apps/src/metabase/helpers/metabaseAPIHelpers.ts‎

Lines changed: 2 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ import {
3333
fetchDatabaseFields
3434
} from './metabaseAPI';
3535
import { Card, SearchApiResponse } from './types';
36+
import { processCard } from './analystModeTypes';
3637

3738
// =============================================================================
3839
// HELPER FUNCTIONS FOR DATA EXTRACTION
@@ -237,39 +238,7 @@ export async function getAllCardsAndModels(forceRefresh = false) {
237238
sortedCards.length = 1000;
238239
}
239240

240-
// Remove unnecessary fields and keep only required fields
241-
const fieldsToRemove = [
242-
'cache_invalidated_at',
243-
'archived',
244-
'collection_position',
245-
'source_card_id',
246-
'result_metadata',
247-
'creator',
248-
'initially_published_at',
249-
'enable_embedding',
250-
'collection_id',
251-
'made_public_by_id',
252-
'embedding_params',
253-
'cache_ttl',
254-
'archived_directly',
255-
'collection_preview'
256-
];
257-
258-
const processedCards = map(sortedCards, (card: any) => {
259-
// Remove unwanted fields
260-
const cleanCard: any = omit(card, fieldsToRemove);
261-
262-
// Process nested fields - keep only specific properties
263-
if (cleanCard['last-edit-info']) {
264-
cleanCard['last-edit-info'] = pick(cleanCard['last-edit-info'], ['timestamp']);
265-
}
266-
267-
if (cleanCard.collection) {
268-
cleanCard.collection = pick(cleanCard.collection, ['slug']);
269-
}
270-
271-
return cleanCard;
272-
});
241+
const processedCards = map(sortedCards, processCard);
273242

274243
const processedModelFields = filteredModels.flatMap((model) => {
275244
const result_metadata = get(model, 'result_metadata', [])

‎apps/src/metabase/helpers/sqlQuery.ts‎

Lines changed: 22 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -184,22 +184,28 @@ export const applySQLEdits = (sql: string, edits: SQLEdits): string => {
184184
}
185185

186186

187-
let updatedSQL = sql;
188-
for (const edit of edits) {
189-
const { old_sql, new_sql, replace_all } = edit;
190-
if (old_sql === "<ENTIRE_QUERY>") {
191-
// Replace the entire SQL query
192-
updatedSQL = new_sql;
193-
} else if (old_sql === "") {
194-
// If old_sql is empty, set new_sql as the updated SQL
195-
updatedSQL = new_sql;
196-
} else if (replace_all) {
197-
// Replace all occurrences of old_sql with new_sql
198-
updatedSQL = updatedSQL.split(old_sql).join(new_sql);
199-
} else {
200-
// Replace only the first occurrence of old_sql with new_sql
201-
updatedSQL = updatedSQL.replace(old_sql, new_sql);
187+
try {
188+
let updatedSQL = sql;
189+
for (const edit of edits) {
190+
const { old_sql, new_sql, replace_all } = edit;
191+
if (old_sql === "<ENTIRE_QUERY>") {
192+
// Replace the entire SQL query
193+
updatedSQL = new_sql;
194+
} else if (old_sql === "") {
195+
// If old_sql is empty, set new_sql as the updated SQL
196+
updatedSQL = new_sql;
197+
} else if (replace_all) {
198+
// Replace all occurrences of old_sql with new_sql
199+
updatedSQL = updatedSQL.split(old_sql).join(new_sql);
200+
} else {
201+
// Replace only the first occurrence of old_sql with new_sql
202+
updatedSQL = updatedSQL.replace(old_sql, new_sql);
203+
}
202204
}
205+
return updatedSQL;
206+
} catch (error) {
207+
console.error("Error applying sql_edits:", error);
208+
return sql; // Return the original SQL if applying edits fails
203209
}
204-
return updatedSQL;
210+
205211
};

0 commit comments

Comments
 (0)