Merge pull request #1735 from adityasharad/ast/synthetic-query-pack

Contextual queries: Support running when the library pack is in the package cache

extensions/ql-vscode/src/cli.ts

@@ -970,6 +970,12 @@ export class CodeQLCliServer implements Disposable {
     }
   }
 
+  async packResolveDependencies(dir: string): Promise<{ [pack: string]: string }> {
+    // Uses the default `--mode use-lock`, which creates the lock file if it doesn't exist.
+    const results: { [pack: string]: string } = await this.runJsonCodeQlCliCommand(['pack', 'resolve-dependencies'], [dir], 'Resolving pack dependencies');
+    return results;
+  }
+
   async generateDil(qloFile: string, outFile: string): Promise<void> {
     const extraArgs = await this.cliConstraints.supportsDecompileDil()
       ? ['--kind', 'dil', '-o', outFile, qloFile]

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,69 @@ 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, createdTempLockFile: boolean }> {
+  // 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)));
+  let createdTempLockFile = false;
+  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);
+    createdTempLockFile = true;
+    // 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, createdTempLockFile };
+}
+
+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, createdTempLockFile } = 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 (createdTempLockFile) {
+    await removeTemporaryLockFile(packPath);
+  }
+  return queryResult;
+}

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

@@ -21,18 +21,17 @@ 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';
 
 /**
- * Run templated CodeQL queries to find definitions and references in
+ * Runs templated CodeQL queries to find definitions in
  * source-language files. We may eventually want to find a way to
  * generalize this to other custom queries, e.g. showing dataflow to
  * or from a selected identifier.
  */
-
 export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   private cache: CachedOperation<LocationLink[]>;
 
@@ -77,6 +76,12 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
   }
 }
 
+/**
+ * Runs templated CodeQL queries to find references in
+ * source-language files. We may eventually want to find a way to
+ * generalize this to other custom queries, e.g. showing dataflow to
+ * or from a selected identifier.
+ */
 export class TemplateQueryReferenceProvider implements ReferenceProvider {
   private cache: CachedOperation<FullLocationLink[]>;
 
@@ -131,6 +136,10 @@ type QueryWithDb = {
   dbUri: Uri
 };
 
+/**
+ * Run templated CodeQL queries to produce AST information for
+ * source-language files.
+ */
 export class TemplatePrintAstProvider {
   private cache: CachedOperation<QueryWithDb>;
 
@@ -199,29 +208,18 @@ export class TemplatePrintAstProvider {
         zippedArchive.pathWithinSourceArchive
     };
 
-    const initialInfo = await createInitialQueryInfo(
-      Uri.file(query),
-      {
-        name: db.name,
-        databaseUri: db.databaseUri.toString(),
-      },
-      false
-    );
-
+    const queryResult = await runContextualQuery(query, db, this.queryStorageDir, this.qs, this.cli, progress, token, templates);
     return {
-      query: await this.qs.compileAndRunQueryAgainstDatabase(
-        db,
-        initialInfo,
-        this.queryStorageDir,
-        progress,
-        token,
-        templates
-      ),
+      query: queryResult,
       dbUri: db.databaseUri
     };
   }
 }
 
+/**
+ * Run templated CodeQL queries to produce CFG information for
+ * source-language files.
+ */
 export class TemplatePrintCfgProvider {
   private cache: CachedOperation<[Uri, Record<string, string>] | undefined>;