Fix AST viewer and contextual query caches when refreshing test databases

extensions/ql-vscode/src/language-support/contextual/cached-operation.ts

@@ -77,6 +77,10 @@ export class CachedOperation<S extends unknown[], U> {
     }
   }
 
+  has(t: string): boolean {
+    return this.cached.has(t);
+  }
+
   reset() {
     this.cached.clear();
     this.lru.length = 0;

extensions/ql-vscode/src/language-support/contextual/template-provider.ts

@@ -17,6 +17,7 @@ import {
 } from "../../common/vscode/archive-filesystem-provider";
 import type { CodeQLCliServer } from "../../codeql-cli/cli";
 import type { DatabaseManager } from "../../databases/local-databases";
+import { DatabaseEventKind } from "../../databases/local-databases";
 import { CachedOperation } from "./cached-operation";
 import type { ProgressCallback } from "../../common/vscode/progress";
 import { withProgress } from "../../common/vscode/progress";
@@ -59,16 +60,25 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider {
     private queryStorageDir: string,
   ) {
     this.cache = new CachedOperation(this.getDefinitions.bind(this));
+    this.dbm.onDidChangeDatabaseItem(({ kind }) => {
+      if (kind === DatabaseEventKind.Remove) {
+        this.cache.reset();
+      }
+    });
   }
 
   async provideDefinition(
     document: TextDocument,
     position: Position,
     token: CancellationToken,
   ): Promise<LocationLink[]> {
+    const uriString = document.uri.toString();
+    if (this.shouldUseCache() && this.cache.has(uriString)) {
+      await maybeReimportTestDatabase(this.dbm, document.uri);
+    }
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString(), token)
-      : await this.getDefinitions(document.uri.toString(), token);
+      ? await this.cache.get(uriString, token)
+      : await this.getDefinitions(uriString, token);
 
     const locLinks: LocationLink[] = [];
     for (const link of fileLinks) {
@@ -119,6 +129,11 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     private queryStorageDir: string,
   ) {
     this.cache = new CachedOperation(this.getReferences.bind(this));
+    this.dbm.onDidChangeDatabaseItem(({ kind }) => {
+      if (kind === DatabaseEventKind.Remove) {
+        this.cache.reset();
+      }
+    });
   }
 
   async provideReferences(
@@ -127,9 +142,13 @@ export class TemplateQueryReferenceProvider implements ReferenceProvider {
     _context: ReferenceContext,
     token: CancellationToken,
   ): Promise<Location[]> {
+    const uriString = document.uri.toString();
+    if (this.shouldUseCache() && this.cache.has(uriString)) {
+      await maybeReimportTestDatabase(this.dbm, document.uri);
+    }
     const fileLinks = this.shouldUseCache()
-      ? await this.cache.get(document.uri.toString(), token)
-      : await this.getReferences(document.uri.toString(), token);
+      ? await this.cache.get(uriString, token)
+      : await this.getReferences(uriString, token);
 
     const locLinks: Location[] = [];
     for (const link of fileLinks) {
@@ -193,6 +212,11 @@ export class TemplatePrintAstProvider {
     private queryStorageDir: string,
   ) {
     this.cache = new CachedOperation(this.getAst.bind(this));
+    this.dbm.onDidChangeDatabaseItem(({ kind }) => {
+      if (kind === DatabaseEventKind.Remove) {
+        this.cache.reset();
+      }
+    });
   }
 
   async provideAst(
@@ -205,9 +229,13 @@ export class TemplatePrintAstProvider {
         "Cannot view the AST. Please select a valid source file inside a CodeQL database.",
       );
     }
+    const uriString = fileUri.toString();
+    if (this.shouldUseCache() && this.cache.has(uriString)) {
+      await maybeReimportTestDatabase(this.dbm, fileUri);
+    }
     const completedQuery = this.shouldUseCache()
-      ? await this.cache.get(fileUri.toString(), progress, token)
-      : await this.getAst(fileUri.toString(), progress, token);
+      ? await this.cache.get(uriString, progress, token)
+      : await this.getAst(uriString, progress, token);
 
     const queryResults = Array.from(completedQuery.results.values());
     if (queryResults.length !== 1) {
@@ -296,6 +324,11 @@ export class TemplatePrintCfgProvider {
     private dbm: DatabaseManager,
   ) {
     this.cache = new CachedOperation(this.getCfgUri.bind(this));
+    this.dbm.onDidChangeDatabaseItem(({ kind }) => {
+      if (kind === DatabaseEventKind.Remove) {
+        this.cache.reset();
+      }
+    });
   }
 
   async provideCfgUri(
@@ -369,3 +402,20 @@ export class TemplatePrintCfgProvider {
     return [queryUri, templates];
   }
 }
+
+async function maybeReimportTestDatabase(
+  dbm: DatabaseManager,
+  uri: Uri,
+): Promise<void> {
+  if (uri.scheme !== zipArchiveScheme) {
+    return;
+  }
+  const zippedArchive = decodeSourceArchiveUri(uri);
+  const sourceArchiveUri = encodeArchiveBasePath(
+    zippedArchive.sourceArchiveZipPath,
+  );
+  const db = dbm.findDatabaseItemBySourceArchive(sourceArchiveUri);
+  if (db) {
+    await dbm.maybeReimportTestDatabase(db.databaseUri);
+  }
+}

extensions/ql-vscode/test/unit-tests/common/cached-operation.test.ts

@@ -0,0 +1,29 @@
+import { CachedOperation } from "../../../src/language-support/contextual/cached-operation";
+
+describe("CachedOperation", () => {
+  it("caches results and tracks cached keys with has()", async () => {
+    const operation = jest.fn(async (key: string, suffix: string) => {
+      return `${key}:${suffix}`;
+    });
+    const cache = new CachedOperation<[string], string>(operation);
+
+    expect(cache.has("a")).toBe(false);
+
+    const res1 = await cache.get("a", "1");
+    expect(res1).toBe("a:1");
+    expect(cache.has("a")).toBe(true);
+    expect(operation).toHaveBeenCalledTimes(1);
+
+    const res2 = await cache.get("a", "2");
+    expect(res2).toBe("a:1");
+    expect(operation).toHaveBeenCalledTimes(1);
+
+    cache.reset();
+    expect(cache.has("a")).toBe(false);
+
+    const res3 = await cache.get("a", "2");
+    expect(res3).toBe("a:2");
+    expect(cache.has("a")).toBe(true);
+    expect(operation).toHaveBeenCalledTimes(2);
+  });
+});