Contextual queries: Factor common code for resolving and running

Shared by the AST viewer, jump to def, and find references
contextual queries.

This allows contextual queries to have their dependencies
resolved and be run whether the library pack is in the
workspace or in the package cache.

extensions/ql-vscode/src/contextual/locationFinder.ts

@@ -5,9 +5,9 @@ import { DatabaseManager, DatabaseItem } from '../databases';
 import fileRangeFromURI from './fileRangeFromURI';
 import { ProgressCallback } from '../commandRunner';
 import { KeyType } from './keyType';
-import { qlpackOfDatabase, resolveQueries } from './queryResolver';
+import { qlpackOfDatabase, resolveQueries, runContextualQuery } from './queryResolver';
 import { CancellationToken, LocationLink, Uri } from 'vscode';
-import { createInitialQueryInfo, QueryWithResults } from '../run-queries-shared';
+import { QueryWithResults } from '../run-queries-shared';
 import { QueryRunner } from '../queryRunner';
 
 export const SELECT_QUERY_NAME = '#select';
@@ -56,15 +56,7 @@ export async function getLocationsForUriString(
 
   const links: FullLocationLink[] = [];
   for (const query of await resolveQueries(cli, qlpack, keyType)) {
-    const initialInfo = await createInitialQueryInfo(
-      Uri.file(query),
-      {
-        name: db.name,
-        databaseUri: db.databaseUri.toString(),
-      },
-      false
-    );
-    const results = await qs.compileAndRunQueryAgainstDatabase(db, initialInfo, queryStorageDir, progress, token, templates);
+    const results = await runContextualQuery(query, db, queryStorageDir, qs, cli, progress, token, templates);
     if (results.successful) {
       links.push(...await getLinksFromResults(results, cli, db, filter));
     }

extensions/ql-vscode/src/contextual/queryResolver.ts

@@ -1,6 +1,7 @@
 import * as fs from 'fs-extra';
 import * as yaml from 'js-yaml';
 import * as tmp from 'tmp-promise';
+import * as path from 'path';
 
 import * as helpers from '../helpers';
 import {
@@ -12,6 +13,11 @@ import {
 import { CodeQLCliServer } from '../cli';
 import { DatabaseItem } from '../databases';
 import { QlPacksForLanguage } from '../helpers';
+import { logger } from '../logging';
+import { createInitialQueryInfo } from '../run-queries-shared';
+import { CancellationToken, Uri } from 'vscode';
+import { ProgressCallback } from '../commandRunner';
+import { QueryRunner } from '../queryRunner';
 
 export async function qlpackOfDatabase(cli: CodeQLCliServer, db: DatabaseItem): Promise<QlPacksForLanguage> {
   if (db.contents === undefined) {
@@ -104,3 +110,67 @@ export async function resolveQueries(cli: CodeQLCliServer, qlpacks: QlPacksForLa
   void helpers.showAndLogErrorMessage(errorMessage);
   throw new Error(`Couldn't find any queries tagged ${tagOfKeyType(keyType)} in any of the following packs: ${packsToSearch.join(', ')}.`);
 }
+
+async function resolveContextualQuery(cli: CodeQLCliServer, query: string): Promise<{ packPath: string, lockFilePath?: string }> {
+  // Contextual queries now live within the standard library packs.
+  // This simplifies distribution (you don't need the standard query pack to use the AST viewer),
+  // but if the library pack doesn't have a lockfile, we won't be able to find
+  // other pack dependencies of the library pack.
+
+  // Work out the enclosing pack.
+  const packContents = await cli.packPacklist(query, false);
+  const packFilePath = packContents.find((p) => ['codeql-pack.yml', 'qlpack.yml'].includes(path.basename(p)));
+  if (packFilePath === undefined) {
+    // Should not happen; we already resolved this query.
+    throw new Error(`Could not find a CodeQL pack file for the pack enclosing the contextual query ${query}`);
+  }
+  const packPath = path.dirname(packFilePath);
+  const lockFilePath = packContents.find((p) => ['codeql-pack.lock.yml', 'qlpack.lock.yml'].includes(path.basename(p)));
+  if (!lockFilePath) {
+    // No lock file, likely because this library pack is in the package cache.
+    // Create a lock file so that we can resolve dependencies and library path
+    // for the contextual query.
+    void logger.log(`Library pack ${packPath} is missing a lock file; creating a temporary lock file`);
+    await cli.packResolveDependencies(packPath);
+    // Clear CLI server pack cache before installing dependencies,
+    // so that it picks up the new lock file, not the previously cached pack.
+    void logger.log('Clearing the CodeQL CLI server\'s pack cache');
+    await cli.clearCache();
+    // Install dependencies.
+    void logger.log(`Installing package dependencies for library pack ${packPath}`);
+    await cli.packInstall(packPath);
+  }
+  return { packPath, lockFilePath };
+}
+
+async function removeTemporaryLockFile(packPath: string) {
+  const tempLockFilePath = path.resolve(packPath, 'codeql-pack.lock.yml');
+  void logger.log(`Deleting temporary package lock file at ${tempLockFilePath}`);
+  // It's fine if the file doesn't exist.
+  await fs.promises.rm(path.resolve(packPath, 'codeql-pack.lock.yml'), { force: true });
+}
+
+export async function runContextualQuery(query: string, db: DatabaseItem, queryStorageDir: string, qs: QueryRunner, cli: CodeQLCliServer, progress: ProgressCallback, token: CancellationToken, templates: Record<string, string>) {
+  const { packPath, lockFilePath } = await resolveContextualQuery(cli, query);
+  const initialInfo = await createInitialQueryInfo(
+    Uri.file(query),
+    {
+      name: db.name,
+      databaseUri: db.databaseUri.toString(),
+    },
+    false
+  );
+  void logger.log(`Running contextual query ${query}; results will be stored in ${queryStorageDir}`);
+  const queryResult = await qs.compileAndRunQueryAgainstDatabase(
+    db,
+    initialInfo,
+    queryStorageDir,
+    progress,
+    token,
+    templates
+  );
+  if (!lockFilePath) {
+    await removeTemporaryLockFile(packPath);
+  }
+  return queryResult;
+}

extensions/ql-vscode/src/contextual/templateProvider.ts

@@ -1,5 +1,3 @@
-import * as fs from 'fs-extra';
-import * as path from 'path';
 import {
   CancellationToken,
   DefinitionProvider,
@@ -16,7 +14,6 @@ import {
 import { decodeSourceArchiveUri, encodeArchiveBasePath, zipArchiveScheme } from '../archive-filesystem-provider';
 import { CodeQLCliServer } from '../cli';
 import { DatabaseManager } from '../databases';
-import { logger } from '../logging';
 import { CachedOperation } from '../helpers';
 import { ProgressCallback, withProgress } from '../commandRunner';
 import AstBuilder from './astBuilder';
@@ -24,9 +21,9 @@ import {
   KeyType,
 } from './keyType';
 import { FullLocationLink, getLocationsForUriString, TEMPLATE_NAME } from './locationFinder';
-import { qlpackOfDatabase, resolveQueries } from './queryResolver';
+import { qlpackOfDatabase, resolveQueries, runContextualQuery } from './queryResolver';
 import { isCanary, NO_CACHE_AST_VIEWER } from '../config';
-import { createInitialQueryInfo, QueryWithResults } from '../run-queries-shared';
+import { QueryWithResults } from '../run-queries-shared';
 import { QueryRunner } from '../queryRunner';
 
 /**
@@ -211,61 +208,7 @@ export class TemplatePrintAstProvider {
         zippedArchive.pathWithinSourceArchive
     };
 
-    // The AST viewer queries now live within the standard library packs.
-    // This simplifies distribution (you don't need the standard query pack to use the AST viewer),
-    // but if the library pack doesn't have a lockfile, we won't be able to find
-    // other pack dependencies of the library pack.
-
-    // Work out the enclosing pack.
-    const packContents = await this.cli.packPacklist(query, false);
-    const packFilePath = packContents.find((p) => ['codeql-pack.yml', 'qlpack.yml'].includes(path.basename(p)));
-    if (packFilePath === undefined) {
-      // Should not happen; we already resolved this query.
-      throw new Error(`Could not find a CodeQL pack file for the pack enclosing the query ${query}`);
-    }
-    const packPath = path.dirname(packFilePath);
-    const lockFilePath = packContents.find((p) => ['codeql-pack.lock.yml', 'qlpack.lock.yml'].includes(path.basename(p)));
-    if (!lockFilePath) {
-      // No lock file, likely because this library pack is in the package cache.
-      // Create a lock file so that we can resolve dependencies and library path
-      // for the AST query.
-      void logger.log(`Library pack ${packPath} is missing a lock file; creating a temporary lock file`);
-      await this.cli.packResolveDependencies(packPath);
-      // Clear CLI server pack cache before installing dependencies,
-      // so that it picks up the new lock file, not the previously cached pack.
-      void logger.log('Clearing the CodeQL CLI server\'s pack cache');
-      await this.cli.clearCache();
-      // Install dependencies.
-      void logger.log(`Installing package dependencies for library pack ${packPath}`);
-      await this.cli.packInstall(packPath);
-    }
-
-    const initialInfo = await createInitialQueryInfo(
-      Uri.file(query),
-      {
-        name: db.name,
-        databaseUri: db.databaseUri.toString(),
-      },
-      false
-    );
-
-    void logger.log(`Running AST query ${query}; results will be stored in ${this.queryStorageDir}`);
-    const queryResult = await this.qs.compileAndRunQueryAgainstDatabase(
-      db,
-      initialInfo,
-      this.queryStorageDir,
-      progress,
-      token,
-      templates
-    );
-
-    if (!lockFilePath) {
-      // Clean up the temporary lock file we created.
-      const tempLockFilePath = path.resolve(packPath, 'codeql-pack.lock.yml');
-      void logger.log(`Deleting temporary package lock file at ${tempLockFilePath}`);
-      // It's fine if the file doesn't exist.
-      await fs.promises.rm(path.resolve(packPath, 'codeql-pack.lock.yml'), { force: true });
-    }
+    const queryResult = await runContextualQuery(query, db, this.queryStorageDir, this.qs, this.cli, progress, token, templates);
     return {
       query: queryResult,
       dbUri: db.databaseUri