CodeLens Buttons to Run SQL

package.json

@@ -125,6 +125,12 @@
         "title": "Build Strict",
         "category": "Evidence",
         "enablement": "evidence.hasProject"
+      },
+      {
+        "command": "evidence.runQuery",
+        "title": "Run Query",
+        "category": "Evidence",
+        "enablement": "evidence.hasProject"
       }
     ],
     "keybindings": [
@@ -362,6 +368,7 @@
   ],
   "dependencies": {
     "tiged": "^2.12.5",
-    "vsce": "^2.11.0"
+    "vsce": "^2.11.0",
+    "@sqltools/types": "0.1.7"
   }
 }

src/commands/commands.ts

@@ -20,6 +20,7 @@ import { preview } from './preview';
 import { openSettingsFile, viewExtensionSettings, viewAppSettings } from './settings';
 import { clearCache} from './cache';
 import { showOutput } from '../output';
+import { runQuery } from './query';
 
 /**
  * VSCode and Evidence extension commands.
@@ -55,6 +56,7 @@ export const enum Commands {
   ShowOutput = 'evidence.showOutput',
   OpenIndex = 'evidence.openIndex',
   OpenEvidenceWalkthrough = 'evidence.openWalkthrough',
+  RunQuery = 'evidence.runQuery',
   OpenSimpleBrowser = 'simpleBrowser.api.open'
 }
 
@@ -85,6 +87,7 @@ export function registerCommands(context: ExtensionContext) {
   registerCommand(Commands.ShowOutput, showOutput);
   registerCommand(Commands.OpenIndex, openIndex);
   registerCommand(Commands.OpenEvidenceWalkthrough, openWalkthrough);
+  registerCommand(Commands.RunQuery, runQuery);
 }
 
 /**

src/commands/query.ts

@@ -0,0 +1,139 @@
+import { extensions, window, commands, ExtensionContext, workspace} from 'vscode';
+import { IExtension} from '@sqltools/types';
+import { syncSettings, readEvidenceSettingsFile } from '../settings';
+
+// TODO: move this to a settings file
+const supportedSqlToolsDrivers = [
+    {
+      "id": "mssql",
+      "name": "Microsoft SQL Server",
+      "extensionId": "mtxr.sqltools-driver-mssql"
+    },
+    {
+      "id": "mysql",
+      "name": "MySQL",
+      "extensionId": "mtxr.sqltools-driver-mysql"
+    },
+    {
+      "id": "postgres",
+      "name": "PostgreSQL",
+      "extensionId": "mtxr.sqltools-driver-pg"
+    },
+    {
+      "id": "sqlite",
+      "name": "SQLite",
+      "extensionId": "mtxr.sqltools-driver-sqlite"
+    },
+    {
+      "id": "redshift",
+      "name": "Amazon Redshift",
+      "extensionId": "kj.sqltools-driver-redshift"
+    },
+    {
+      "id": "bigquery",
+      "name": "BigQuery",
+      "extensionId": "Evidence.sqltools-bigquery-driver"
+    },
+    {
+      "id": "snowflake",
+      "name": "Snowflake",
+      "extensionId": "koszti.snowflake-driver-for-sqltools"
+    }
+  ];
+
+export async function activate(context: ExtensionContext) {
+  // load the SQLTools extension
+  const sqltools = extensions.getExtension<IExtension>('mtxr.sqltools');
+  if (sqltools) {
+    if (!sqltools.isActive) {
+      await sqltools.activate();
+    }
+  }
+  // register the command
+  context.subscriptions.push(commands.registerCommand('evidence.runQuery', runQuery));
+}
+
+
+export async function runQuery(name: string, query: string) {
+  const sqltools = extensions.getExtension<IExtension>('mtxr.sqltools');
+  // check if sqltools is installed
+  if (sqltools) {
+    if (!sqltools.isActive) {
+      await sqltools.activate();
+    }
+    // check which database is being used
+    let settings = await readEvidenceSettingsFile();
+    let matchingDB =false;
+    supportedSqlToolsDrivers.forEach(async (db: any) => {
+      if (db.id === settings.database) {
+        //check if the driver is installed
+        matchingDB = true;
+        let driver = extensions.getExtension(db.extensionId);
+        if (driver) {
+          try {
+            // sync settings from evidence to sqltools
+            await syncSettings();
+            // check if it is a chained query by looking for ${...} in the query with regex
+            let isChainedQuery = query.match(/\${.*}/); 
+            if (isChainedQuery) {
+              let compiledQuery = await getCompiledQuery(name);
+              try {
+                commands.executeCommand('sqltools.executeQuery', compiledQuery);
+              } catch (error) {
+                console.error(error);
+              }
+            } else {
+              commands.executeCommand('sqltools.executeQuery', query);
+            }
+          } catch (error) {
+            console.error('Error executing query with SQLTools:', error);
+          }
+        } else {
+          // Prompt the user to install the driver
+          const choice = await window.showInformationMessage(`Install the SQLTools ${db.name} driver?`, 'Install', 'View In Marketplace');
+          if (choice === 'Install') {
+            commands.executeCommand('workbench.extensions.installExtension', db.extensionId);
+          } else if (choice === 'View In Marketplace') {
+            commands.executeCommand('workbench.extensions.action.showExtensionsWithIds', [db.extensionId]);
+          }
+        }
+      }
+    });
+    if (!matchingDB) {
+      window.showErrorMessage(`No supported SQLTools driver found for: ${settings.database}`);
+    }
+  } else {
+    //  Prompt the user to install sqltools
+    const choice = await window.showInformationMessage(`Running queries in VSCode requires the SQLTools extension, install?`, 'Install', 'View In Marketplace');
+    if (choice === 'Install') {
+      commands.executeCommand('workbench.extensions.installExtension', 'mtxr.sqltools');
+    }
+    else if (choice === 'View In Marketplace') {
+      commands.executeCommand('workbench.extensions.action.showExtensionsWithIds', ['mtxr.sqltools']);
+    }
+  }
+}
+
+/* Gets compiled query from evidennce file system. 
+/  This will fail if Evidence has not run since the query was created.
+/  And it will return old queries if the user has not saved the file containing the SQL, or if they are not on the right evidence page.
+/  Maybe this is too broken and we should disable chained queries for now.
+*/
+export async function getCompiledQuery(name: string) {
+  // read in compiled queries from .evidence/template/extracted/queries.json
+  // return the query with the matching name
+  let queryFiles = await workspace.findFiles('.evidence/template/.evidence-queries/extracted/**/queries.json');
+  if(queryFiles.length>0)
+  {
+    let queries = await workspace.fs.readFile(queryFiles[0]);
+    let queryList = JSON.parse(queries.toString());
+    let matchingQuery = queryList.find((q: any) => q.id === name);
+    // return compiled query
+    if (matchingQuery){
+      return matchingQuery.compiledQueryString;
+    } else {
+      window.showErrorMessage(`No compiled query found with name ${name}: Chained queries require Evidence to be running.`);
+      throw new Error(`No chanied query found with name ${name}: Chained queries require Evidence to be running.`);
+    }
+  }
+}

src/extension.ts

@@ -7,6 +7,8 @@ import {
 } from 'vscode';
 
 import { MarkdownSymbolProvider } from './providers/markdownSymbolProvider';
+import { MarkdownCodeLensProvider } from './providers/markdownCodeLensProvider';
+import { SqlCodeLensProvider } from './providers/sqlCodeLensProvider';
 import { setExtensionContext } from './extensionContext';
 import { registerCommands } from './commands/commands';
 import { loadPackageJson, hasDependency } from './utils/jsonUtils';
@@ -30,6 +32,19 @@ export async function activate(context: ExtensionContext) {
   const provider = new MarkdownSymbolProvider();
   // languages.registerDocumentSymbolProvider(markdownLanguage, provider);
 
+
+  // register markdown code lens provider
+  const markdownProvider = new MarkdownCodeLensProvider();
+  context.subscriptions.push(
+    languages.registerCodeLensProvider({ language: 'emd' }, markdownProvider)
+  );  
+
+  // register sql code lens provider
+  const sqlProvider = new SqlCodeLensProvider();
+  context.subscriptions.push(
+    languages.registerCodeLensProvider({ language: 'sql' }, sqlProvider)
+  );
+
   // load package.json
   const workspacePackageJson = await loadPackageJson();
 

src/providers/markdownCodeLensProvider.ts

@@ -0,0 +1,57 @@
+import * as vscode from 'vscode';
+
+export class MarkdownCodeLensProvider implements vscode.CodeLensProvider {
+    provideCodeLenses(document: vscode.TextDocument): vscode.CodeLens[] | Thenable<vscode.CodeLens[]> {
+        const codeLenses: vscode.CodeLens[] = [];
+
+        // Parse the Markdown file to extract code blocks
+        const codeBlocks = parseCodeBlocks(document);
+
+        // Create a CodeLens for each code block
+        for (const codeBlock of codeBlocks) {
+            const range = new vscode.Range(
+                codeBlock.startLine, 0,  // Start position of the code block
+                codeBlock.endLine, 0     // End position of the code block
+            );
+
+            const codeLens = new vscode.CodeLens(range);
+            codeLens.command = {
+                title: '$(play) Run Query',
+                command: 'evidence.runQuery',
+                arguments: [codeBlock.queryName, codeBlock.content]
+            };
+
+            codeLenses.push(codeLens);
+        }
+
+        return codeLenses;
+    }
+
+    resolveCodeLens(codeLens: vscode.CodeLens): vscode.CodeLens | Thenable<vscode.CodeLens> {
+        return codeLens;
+    }
+}
+
+function parseCodeBlocks(document: vscode.TextDocument): CodeBlock[] {
+    const codeBlocks: CodeBlock[] = [];
+    const regex = /```sql\s+(\w+)\n([\s\S]*?)\n```/g;
+    let match;
+
+    while ((match = regex.exec(document.getText()))) {
+        const queryName = match[1];
+        const content = match[2];
+        const startLine = document.positionAt(match.index).line;
+        const endLine = document.positionAt(match.index + match[0].length).line;
+
+        codeBlocks.push({ queryName, content, startLine, endLine });
+    }
+
+    return codeBlocks;
+}
+
+interface CodeBlock {
+    queryName: string;
+    content: string;
+    startLine: number;
+    endLine: number;
+}