From 0fb50de032d78a129276b2f262c071c146710c1d Mon Sep 17 00:00:00 2001 From: Josef Svenningsson Date: Wed, 14 Oct 2026 12:52:26 +0100 Subject: [PATCH] Refresh AST and contextual query caches when reimporting test databases Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> --- .../contextual/cached-operation.ts | 4 + .../contextual/template-provider.ts | 62 +++- .../common/cached-operation.test.ts | 29 ++ .../contextual/template-provider.test.ts | 275 ++++++++++++++++++ 4 files changed, 364 insertions(+), 6 deletions(-) create mode 100644 extensions/ql-vscode/test/unit-tests/common/cached-operation.test.ts create mode 100644 extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/template-provider.test.ts diff --git a/extensions/ql-vscode/src/language-support/contextual/cached-operation.ts b/extensions/ql-vscode/src/language-support/contextual/cached-operation.ts index 3c6e845d740..69d292619f1 100644 --- a/extensions/ql-vscode/src/language-support/contextual/cached-operation.ts +++ b/extensions/ql-vscode/src/language-support/contextual/cached-operation.ts @@ -77,6 +77,10 @@ export class CachedOperation { } } + has(t: string): boolean { + return this.cached.has(t); + } + reset() { this.cached.clear(); this.lru.length = 0; diff --git a/extensions/ql-vscode/src/language-support/contextual/template-provider.ts b/extensions/ql-vscode/src/language-support/contextual/template-provider.ts index 8d9cf6327ec..733e06617c0 100644 --- a/extensions/ql-vscode/src/language-support/contextual/template-provider.ts +++ b/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,6 +60,11 @@ 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( @@ -66,9 +72,13 @@ export class TemplateQueryDefinitionProvider implements DefinitionProvider { position: Position, token: CancellationToken, ): Promise { + 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 { + 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 { + 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); + } +} diff --git a/extensions/ql-vscode/test/unit-tests/common/cached-operation.test.ts b/extensions/ql-vscode/test/unit-tests/common/cached-operation.test.ts new file mode 100644 index 00000000000..0fb890fba38 --- /dev/null +++ b/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); + }); +}); diff --git a/extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/template-provider.test.ts b/extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/template-provider.test.ts new file mode 100644 index 00000000000..99feda9d170 --- /dev/null +++ b/extensions/ql-vscode/test/vscode-tests/no-workspace/language-support/contextual/template-provider.test.ts @@ -0,0 +1,275 @@ +import { CancellationTokenSource, EventEmitter, Position, Uri } from "vscode"; +import type { TextDocument } from "vscode"; +import { encodeSourceArchiveUri } from "../../../../../src/common/vscode/archive-filesystem-provider"; +import type { CodeQLCliServer } from "../../../../../src/codeql-cli/cli"; +import type { + DatabaseChangedEvent, + DatabaseItem, + DatabaseManager, +} from "../../../../../src/databases/local-databases"; +import { DatabaseEventKind } from "../../../../../src/databases/local-databases"; +import { + TemplatePrintAstProvider, + TemplatePrintCfgProvider, + TemplateQueryDefinitionProvider, + TemplateQueryReferenceProvider, +} from "../../../../../src/language-support"; +import * as queryResolver from "../../../../../src/language-support/contextual/query-resolver"; +import * as locationFinder from "../../../../../src/language-support/contextual/location-finder"; +import type { + CoreCompletedQuery, + QueryRunner, +} from "../../../../../src/query-server"; +import { QueryResultType } from "../../../../../src/query-server/messages"; +import { QueryOutputDir } from "../../../../../src/local-queries/query-output-dir"; +import { mockDatabaseItem, mockedObject } from "../../../utils/mocking.helpers"; + +describe("template-provider", () => { + let dbChangedEmitter: EventEmitter; + let mockDb: DatabaseItem; + let mockDbm: DatabaseManager; + let maybeReimportTestDatabase: jest.MockedFunction< + typeof DatabaseManager.prototype.maybeReimportTestDatabase + >; + let mockCli: CodeQLCliServer; + let mockQs: QueryRunner; + let fileUri: Uri; + + beforeEach(() => { + dbChangedEmitter = new EventEmitter(); + mockDb = mockDatabaseItem({ + databaseUri: Uri.file("/path/to/db"), + }); + maybeReimportTestDatabase = jest.fn().mockResolvedValue(undefined); + mockDbm = mockedObject({ + onDidChangeDatabaseItem: dbChangedEmitter.event, + findDatabaseItemBySourceArchive: jest.fn().mockReturnValue(mockDb), + findDatabaseItem: jest.fn().mockReturnValue(mockDb), + maybeReimportTestDatabase, + }); + mockCli = mockedObject({}); + mockQs = mockedObject({}); + fileUri = encodeSourceArchiveUri({ + sourceArchiveZipPath: "/path/to/db/src.zip", + pathWithinSourceArchive: "/src/test.cs", + }); + }); + + afterEach(() => { + dbChangedEmitter.dispose(); + }); + + describe("TemplatePrintAstProvider", () => { + let runContextualQuerySpy: jest.SpiedFunction< + typeof queryResolver.runContextualQuery + >; + + beforeEach(() => { + jest + .spyOn(queryResolver, "resolveContextualQlPacksForDatabase") + .mockResolvedValue({ + dbschemePack: "my-qlpack", + dbschemePackIsLibraryPack: false, + }); + jest + .spyOn(queryResolver, "resolveContextualQueries") + .mockResolvedValue(["/path/to/PrintAst.ql"]); + runContextualQuerySpy = jest + .spyOn(queryResolver, "runContextualQuery") + .mockImplementation(async () => + mockedObject({ + outputDir: new QueryOutputDir("/path/to/output"), + dbPath: "/path/to/db", + results: new Map([ + [ + "/path/to/PrintAst.ql", + { + resultType: QueryResultType.SUCCESS, + message: undefined, + outputBaseName: "results", + evaluationTime: 10, + }, + ], + ]), + }), + ); + }); + + it("caches AST results and invalidates cache when database is removed", async () => { + const provider = new TemplatePrintAstProvider( + mockCli, + mockQs, + mockDbm, + "/path/to/storage", + ); + const progress = jest.fn(); + const token = new CancellationTokenSource().token; + + await provider.provideAst(progress, token, fileUri); + expect(runContextualQuerySpy).toHaveBeenCalledTimes(1); + expect(maybeReimportTestDatabase).not.toHaveBeenCalled(); + + // Second call uses cache, but checks whether test database needs reimporting first + await provider.provideAst(progress, token, fileUri); + expect(maybeReimportTestDatabase).toHaveBeenCalledTimes(1); + expect(maybeReimportTestDatabase).toHaveBeenCalledWith( + mockDb.databaseUri, + ); + expect(runContextualQuerySpy).toHaveBeenCalledTimes(1); + + // Simulate database removal (e.g. during test database reimport) + dbChangedEmitter.fire({ + item: mockDb, + kind: DatabaseEventKind.Remove, + fullRefresh: true, + }); + + await provider.provideAst(progress, token, fileUri); + expect(runContextualQuerySpy).toHaveBeenCalledTimes(2); + }); + + it("re-evaluates AST when maybeReimportTestDatabase reimports the database on a cached file", async () => { + const provider = new TemplatePrintAstProvider( + mockCli, + mockQs, + mockDbm, + "/path/to/storage", + ); + const progress = jest.fn(); + const token = new CancellationTokenSource().token; + + await provider.provideAst(progress, token, fileUri); + expect(runContextualQuerySpy).toHaveBeenCalledTimes(1); + + // Next time maybeReimportTestDatabase is called, simulate reimporting the database + maybeReimportTestDatabase.mockImplementationOnce(async () => { + dbChangedEmitter.fire({ + item: mockDb, + kind: DatabaseEventKind.Remove, + fullRefresh: true, + }); + }); + + await provider.provideAst(progress, token, fileUri); + expect(maybeReimportTestDatabase).toHaveBeenCalledTimes(1); + expect(runContextualQuerySpy).toHaveBeenCalledTimes(2); + }); + }); + + describe("TemplateQueryDefinitionProvider", () => { + it("resets cache on database removal and checks for outdated test database on cache hit", async () => { + const getLocationsSpy = jest + .spyOn(locationFinder, "getLocationsForUriString") + .mockResolvedValue([]); + + const provider = new TemplateQueryDefinitionProvider( + mockCli, + mockQs, + mockDbm, + "/path/to/storage", + ); + const doc = mockedObject({ uri: fileUri }); + const pos = new Position(0, 0); + const token = new CancellationTokenSource().token; + + await provider.provideDefinition(doc, pos, token); + expect(getLocationsSpy).toHaveBeenCalledTimes(1); + expect(maybeReimportTestDatabase).not.toHaveBeenCalled(); + + await provider.provideDefinition(doc, pos, token); + expect(maybeReimportTestDatabase).toHaveBeenCalledTimes(1); + expect(getLocationsSpy).toHaveBeenCalledTimes(1); + + dbChangedEmitter.fire({ + item: mockDb, + kind: DatabaseEventKind.Remove, + fullRefresh: true, + }); + + await provider.provideDefinition(doc, pos, token); + expect(getLocationsSpy).toHaveBeenCalledTimes(2); + }); + }); + + describe("TemplateQueryReferenceProvider", () => { + it("resets cache on database removal and checks for outdated test database on cache hit", async () => { + const getLocationsSpy = jest + .spyOn(locationFinder, "getLocationsForUriString") + .mockResolvedValue([]); + + const provider = new TemplateQueryReferenceProvider( + mockCli, + mockQs, + mockDbm, + "/path/to/storage", + ); + const doc = mockedObject({ uri: fileUri }); + const pos = new Position(0, 0); + const token = new CancellationTokenSource().token; + + await provider.provideReferences( + doc, + pos, + { includeDeclaration: true }, + token, + ); + expect(getLocationsSpy).toHaveBeenCalledTimes(1); + expect(maybeReimportTestDatabase).not.toHaveBeenCalled(); + + await provider.provideReferences( + doc, + pos, + { includeDeclaration: true }, + token, + ); + expect(maybeReimportTestDatabase).toHaveBeenCalledTimes(1); + expect(getLocationsSpy).toHaveBeenCalledTimes(1); + + dbChangedEmitter.fire({ + item: mockDb, + kind: DatabaseEventKind.Remove, + fullRefresh: true, + }); + + await provider.provideReferences( + doc, + pos, + { includeDeclaration: true }, + token, + ); + expect(getLocationsSpy).toHaveBeenCalledTimes(2); + }); + }); + + describe("TemplatePrintCfgProvider", () => { + it("resets cache on database removal", async () => { + const resolveQlPacksSpy = jest + .spyOn(queryResolver, "resolveContextualQlPacksForDatabase") + .mockResolvedValue({ + dbschemePack: "my-qlpack", + dbschemePackIsLibraryPack: false, + }); + jest + .spyOn(queryResolver, "resolveContextualQueries") + .mockResolvedValue(["/path/to/PrintCfg.ql"]); + + const provider = new TemplatePrintCfgProvider(mockCli, mockDbm); + const doc = mockedObject({ uri: fileUri }); + + await provider.provideCfgUri(doc, 1, 1); + expect(resolveQlPacksSpy).toHaveBeenCalledTimes(1); + + await provider.provideCfgUri(doc, 1, 1); + expect(resolveQlPacksSpy).toHaveBeenCalledTimes(1); + + dbChangedEmitter.fire({ + item: mockDb, + kind: DatabaseEventKind.Remove, + fullRefresh: true, + }); + + await provider.provideCfgUri(doc, 1, 1); + expect(resolveQlPacksSpy).toHaveBeenCalledTimes(2); + }); + }); +});