Add basic "quick query" functionality

Author: jcreedcmu

extensions/ql-vscode/package.json

@@ -142,6 +142,10 @@
         "command": "codeQL.quickEval",
         "title": "CodeQL: Quick Evaluation"
       },
+      {
+        "command": "codeQL.quickQuery",
+        "title": "CodeQL: Quick Query"
+      },
       {
         "command": "codeQL.chooseDatabase",
         "title": "CodeQL: Choose Database",

extensions/ql-vscode/src/extension.ts

@@ -17,6 +17,7 @@ import { QueryHistoryManager } from './query-history';
 import * as qsClient from './queryserver-client';
 import { CodeQLCliServer } from './cli';
 import { assertNever } from './helpers-pure';
+import { displayQuickQuery } from './quick-query';
 
 /**
  * extension.ts
@@ -305,6 +306,7 @@ async function activateWithInstalledDistribution(ctx: ExtensionContext, distribu
 
   ctx.subscriptions.push(commands.registerCommand('codeQL.runQuery', async (uri: Uri | undefined) => await compileAndRunQuery(false, uri)));
   ctx.subscriptions.push(commands.registerCommand('codeQL.quickEval', async (uri: Uri | undefined) => await compileAndRunQuery(true, uri)));
+  ctx.subscriptions.push(commands.registerCommand('codeQL.quickQuery', async () => displayQuickQuery(ctx, cliServer, databaseUI)));
 
   ctx.subscriptions.push(client.start());
 }

extensions/ql-vscode/src/quick-query.ts

@@ -0,0 +1,139 @@
+import * as fs from 'fs-extra';
+import * as glob from 'glob-promise';
+import * as yaml from 'js-yaml';
+import * as path from 'path';
+import { ExtensionContext, window as Window, workspace, Uri } from 'vscode';
+import { ErrorCodes, ResponseError } from 'vscode-languageclient';
+import { CodeQLCliServer } from './cli';
+import { DatabaseUI } from './databases-ui';
+import * as helpers from './helpers';
+import { logger } from './logging';
+import { UserCancellationException } from './queries';
+
+const QUICK_QUERIES_DIR_NAME = 'quick-queries';
+const QUICK_QUERY_QUERY_NAME = 'quick-query.ql';
+
+async function getQlPackFor(cliServer: CodeQLCliServer, dbschemePath: string): Promise<string> {
+  const qlpacks = await cliServer.resolveQlpacks(helpers.getOnDiskWorkspaceFolders());
+  const packs: { packDir: string | undefined, packName: string }[] =
+    Object.entries(qlpacks).map(([packName, dirs]) => {
+      if (dirs.length < 1) {
+        logger.log(`In getQlPackFor ${dbschemePath}, qlpack ${packName} has no directories`);
+        return { packName, packDir: undefined };
+      }
+      if (dirs.length > 1) {
+        logger.log(`In getQlPackFor ${dbschemePath}, qlpack ${packName} has more than one directory; arbitrarily choosing the first`);
+      }
+      return {
+        packName,
+        packDir: dirs[0]
+      }
+    });
+  for (const { packDir, packName } of packs) {
+    if (packDir !== undefined) {
+      const qlpack = yaml.safeLoad(await fs.readFile(path.join(packDir, 'qlpack.yml'), 'utf8'));
+      if (qlpack.dbscheme !== undefined && path.basename(qlpack.dbscheme) === path.basename(dbschemePath)) {
+        return packName;
+      }
+    }
+  }
+  throw new Error(`Could not find qlpack file for dbscheme ${dbschemePath}`);
+}
+
+/**
+ * `getBaseText` heuristically returns an appropriate import
+ * statement prelude based on the filename of the dbscheme file
+ * given. This information might be more appropriately provided by
+ * the qlpack itself.
+ */
+function getBaseText(dbschemeBase: string) {
+  if (dbschemeBase == 'semmlecode.javascript.dbscheme') return 'import javascript\n\nselect ""';
+  if (dbschemeBase == 'semmlecode.cpp.dbscheme') return 'import cpp\n\nselect ""';
+  if (dbschemeBase == 'semmlecode.java.dbscheme') return 'import java\n\nselect ""';
+  if (dbschemeBase == 'semmlecode.python.dbscheme') return 'import python\n\nselect ""';
+  return '';
+}
+
+async function getQuickQueriesDir(ctx: ExtensionContext): Promise<string> {
+  const storagePath = ctx.storagePath;
+  if (storagePath === undefined) {
+    throw new Error('Workspace storage path is undefined');
+  }
+  const queriesPath = path.join(storagePath, QUICK_QUERIES_DIR_NAME);
+  fs.ensureDir(queriesPath, { mode: 0o700 });
+  return queriesPath;
+}
+
+/**
+ * Show a buffer the user can enter a simple query into.
+ */
+export async function displayQuickQuery(ctx: ExtensionContext, cliServer: CodeQLCliServer, databaseUI: DatabaseUI) {
+  try {
+
+    // If there is already a quick query open, don't clobber it, just
+    // show it.
+    const existing = workspace.textDocuments.find(doc => path.basename(doc.uri.fsPath) === QUICK_QUERY_QUERY_NAME);
+    if (existing !== undefined) {
+      Window.showTextDocument(existing);
+      return;
+    }
+
+    const queriesDir = await getQuickQueriesDir(ctx);
+
+    // We need this folder in workspace folders so the language server
+    // knows how to find its qlpack.yml
+    if (workspace.workspaceFolders === undefined
+      || !workspace.workspaceFolders.some(folder => folder.uri.fsPath === queriesDir)) {
+      workspace.updateWorkspaceFolders(
+        (workspace.workspaceFolders || []).length,
+        0,
+        { uri: Uri.file(queriesDir), name: "Quick Queries" }
+      );
+    }
+
+    // We're going to infer which qlpack to use from the current database
+    const dbItem = await databaseUI.getDatabaseItem();
+    if (dbItem === undefined) {
+      throw new Error('Can\'t start quick query without a selected database');
+    }
+
+    const datasetFolder = await dbItem.getDatasetFolder(cliServer);
+    const dbschemes = await glob(path.join(datasetFolder, '*.dbscheme'))
+
+    if (dbschemes.length < 1) {
+      throw new Error(`Can't find dbscheme for current database in ${datasetFolder}`);
+    }
+
+    const dbscheme = dbschemes[0];
+    if (dbschemes.length > 1) {
+      logger.log(`Found multiple dbschemes in ${datasetFolder}; arbitrarily choosing the first, ${dbscheme}`);
+    }
+
+    const qlpack = await getQlPackFor(cliServer, dbscheme);
+    const quickQueryQlpackYaml: any = {
+      name: "quick-query",
+      version: "1.0.0",
+      libraryPathDependencies: [qlpack]
+    };
+
+    const qlFile = path.join(queriesDir, QUICK_QUERY_QUERY_NAME);
+    const qlPackFile = path.join(queriesDir, 'qlpack.yml');
+    await fs.writeFile(qlFile, getBaseText(path.basename(dbscheme)), 'utf8');
+    await fs.writeFile(qlPackFile, yaml.safeDump(quickQueryQlpackYaml), 'utf8');
+    Window.showTextDocument(await workspace.openTextDocument(qlFile));
+  }
+
+  // TODO: clean up error handling for top-level commands like this
+  catch (e) {
+    if (e instanceof UserCancellationException) {
+      logger.log(e.message);
+    }
+    else if (e instanceof ResponseError && e.code == ErrorCodes.RequestCancelled) {
+      logger.log(e.message);
+    }
+    else if (e instanceof Error)
+      helpers.showAndLogErrorMessage(e.message);
+    else
+      throw e;
+  }
+}