Support edit sql tool (#260)

* support edit sql tool

* reuse other tools

Author: nuwandavek

apps/src/metabase/appController.ts

@@ -36,7 +36,9 @@ import {
   getParameters,
   getVariablesAndUuidsInQuery,
   MetabaseStateSnippetsDict,
-  getAllTemplateTagsInQuery
+  getAllTemplateTagsInQuery,
+  applySQLEdits,
+  SQLEdits
 } from "./helpers/sqlQuery";
 import axios from 'axios'
 import { getSelectedDbId, getCurrentUserInfo as getUserInfo, getSnippets, getCurrentCard, getDashboardState } from "./helpers/metabaseStateAPI";
@@ -400,6 +402,24 @@ export class MetabaseController extends AppController<MetabaseAppState> {
     }
   }
 
+  @Action({
+    labelRunning: "Executes the SQL query with parameters",
+    labelDone: "Executed query",
+    labelTask: "Executed SQL query",
+    description: "Executes the SQL query in the Metabase SQL editor with support for template tags and parameters.",
+    renderBody: ({ explanation, sql_edits}: { explanation: string, sql_edits: SQLEdits }, appState: MetabaseAppStateSQLEditor) => {
+      const sqlQuery = appState?.sqlQuery
+      const newQuery = applySQLEdits(sqlQuery, sql_edits);
+      return {text: explanation, code: newQuery, oldCode: sqlQuery, language: "sql"}
+    }
+  })
+  async EditAndExecuteQuery({ sql_edits, _ctes = [], explanation = "", template_tags={}, parameters=[] }: { sql_edits: SQLEdits, _ctes?: CTE[], explanation?: string, template_tags?: object, parameters?: any[] }) {
+    const appState = (await this.app.getState()) as MetabaseAppStateSQLEditor;
+    let sql = appState.sqlQuery || "";
+    sql = applySQLEdits(sql, sql_edits);
+    return await this.ExecuteQuery({ sql, _ctes, explanation, template_tags, parameters });
+  }
+
   @Action({
     labelRunning: "Executes the SQL query",
     labelDone: "Executed query",

apps/src/metabase/helpers/sqlQuery.ts

@@ -150,4 +150,52 @@ export const getAllTemplateTagsInQuery = (query: string, allSnippets?: MetabaseS
   const snippetTags = allSnippets ? getSnippetsInQuery(query, allSnippets) : {};
   const modelTags = getModelsInQuery(query);
   return { ...snippetTags, ...modelTags };
-}
+}
+
+
+export interface SQLEdit {
+    old_sql: string;
+    new_sql: string;
+    replace_all: boolean; // replace all occurrences if true, otherwise replace only the first occurrence
+}
+export type SQLEdits = SQLEdit[];
+
+
+export const applySQLEdits = (sql: string, edits: SQLEdits): string => {
+    //Rules
+    // - The edits will be applied in the order they are provided.
+    // - Use sql_edits to apply edits to the SQL query (have to provide old_sql, new_sql, and replace_all flag)
+    // - If you want to replace the entire current SQL query, set old_sql to <ENTIRE_QUERY> and new_sql to the new query.
+    // - If you want to replace a specific part of the SQL query, set old_sql to the part you want to replace and new_sql to the new part.
+    // - If the current sql query is empty, you can set old_sql to an empty string and new_sql to the new query.
+    // - If you want to not change the SQL query, you can set old_sql to an empty string and new_sql to an empty string.
+    // - If the diff is too complicated (>5 edits), consider just replacing the entire query with a new one.
+  try {
+    if (typeof edits === 'string') {
+        edits = JSON.parse(edits) as SQLEdits;
+    }
+  } catch (error) {
+      console.error("Error parsing sql_edits:", error);
+      return sql; // Return the original SQL if parsing fails
+  }
+
+
+  let updatedSQL = sql;
+  for (const edit of edits) {
+    const { old_sql, new_sql, replace_all } = edit;
+    if (old_sql === "<ENTIRE_QUERY>") {
+      // Replace the entire SQL query
+      updatedSQL = new_sql;
+    } else if (old_sql === "") {
+      // If old_sql is empty, set new_sql as the updated SQL
+      updatedSQL = new_sql;
+    } else if (replace_all) {
+      // Replace all occurrences of old_sql with new_sql
+      updatedSQL = updatedSQL.split(old_sql).join(new_sql);
+    } else {
+      // Replace only the first occurrence of old_sql with new_sql
+      updatedSQL = updatedSQL.replace(old_sql, new_sql);
+    }
+  }
+  return updatedSQL;
+};

web/src/components/common/ActionStack.tsx

@@ -25,7 +25,7 @@ import { get, isEmpty } from 'lodash';
 
 // Todo: Vivek: Hardcoding here, need to fix this later
 // This is a list of actions that are undo/redoable
-const UNDO_REDO_ACTIONS = ['ExecuteSQLClient', 'ExecuteQuery']
+const UNDO_REDO_ACTIONS = ['ExecuteSQLClient', 'ExecuteQuery', 'EditAndExecuteQuery']
 
 
 function removeThinkingTags(input: string): string {
@@ -96,7 +96,7 @@ export const ActionStack: React.FC<{status: string, actions: Array<ActionStatusV
             event.stopPropagation();
             executeAction({
                 index: -1,
-                function: fn,
+                function: 'ExecuteQuery',
                 args: {sql: sql},
             });
         };