mrva/vscode-codeql
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Add support for database registration

Browse Source
This commit is contained in:
Andrew Eisenberg
2020-11-27 15:22:52 -08:00
parent 3c08baf062
commit 2795184e70
7 changed files with 266 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files

31
extensions/ql-vscode/src/databaseFetcher.ts
View File

@@ -27,7 +27,7 @@ export async function promptImportInternetDatabase(
databasesManager: DatabaseManager,
storagePath: string,
progress: ProgressCallback,
_: CancellationToken,
token: CancellationToken,
): Promise<DatabaseItem | undefined> {
const databaseUrl = await window.showInputBox({
prompt: 'Enter URL of zipfile of database to download',
@@ -42,7 +42,8 @@ export async function promptImportInternetDatabase(
databaseUrl,
databasesManager,
storagePath,
progress
progress,
token
);
if (item) {
@@ -65,7 +66,7 @@ export async function promptImportLgtmDatabase(
databasesManager: DatabaseManager,
storagePath: string,
progress: ProgressCallback,
_: CancellationToken
token: CancellationToken
): Promise<DatabaseItem | undefined> {
const lgtmUrl = await window.showInputBox({
prompt:
@@ -82,7 +83,8 @@ export async function promptImportLgtmDatabase(
databaseUrl,
databasesManager,
storagePath,
progress
progress,
token
);
if (item) {
commands.executeCommand('codeQLDatabases.focus');
@@ -108,14 +110,15 @@ export async function importArchiveDatabase(
databasesManager: DatabaseManager,
storagePath: string,
progress: ProgressCallback,
_: CancellationToken,
token: CancellationToken,
): Promise<DatabaseItem | undefined> {
try {
const item = await databaseArchiveFetcher(
databaseUrl,
databasesManager,
storagePath,
progress
progress,
token
);
if (item) {
commands.executeCommand('codeQLDatabases.focus');
@@ -139,15 +142,17 @@ export async function importArchiveDatabase(
* @param databaseUrl URL from which to grab the database
* @param databasesManager the DatabaseManager
* @param storagePath where to store the unzipped database.
* @param progressCallback optional callback to send progress messages to
* @param progress callback to send progress messages to
* @param token cancellation token
*/
async function databaseArchiveFetcher(
databaseUrl: string,
databasesManager: DatabaseManager,
storagePath: string,
progressCallback?: ProgressCallback
progress: ProgressCallback,
token: CancellationToken
): Promise<DatabaseItem> {
progressCallback?.({
progress({
message: 'Getting database',
step: 1,
maxStep: 4,
@@ -161,10 +166,10 @@ async function databaseArchiveFetcher(
if (isFile(databaseUrl)) {
await readAndUnzip(databaseUrl, unzipPath);
} else {
await fetchAndUnzip(databaseUrl, unzipPath, progressCallback);
await fetchAndUnzip(databaseUrl, unzipPath, progress);
}
progressCallback?.({
progress({
message: 'Opening database',
step: 3,
maxStep: 4,
@@ -177,14 +182,14 @@ async function databaseArchiveFetcher(
'codeql-database.yml'
);
if (dbPath) {
progressCallback?.({
progress({
message: 'Validating and fixing source location',
step: 4,
maxStep: 4,
});
await ensureZippedSourceLocation(dbPath);
const item = await databasesManager.openDatabase(Uri.file(dbPath));
const item = await databasesManager.openDatabase(progress, token, Uri.file(dbPath));
databasesManager.setCurrentDatabaseItem(item);
return item;
} else {

26
extensions/ql-vscode/src/databases-ui.ts
View File

@@ -318,9 +318,13 @@ export class DatabaseUI extends DisposableObject {
)
);
this.push(
commandRunner(
commandRunnerWithProgress(
'codeQLDatabases.removeDatabase',
this.handleRemoveDatabase
this.handleRemoveDatabase,
{
title: 'Removing database',
cancellable: false
}
)
);
this.push(
@@ -580,7 +584,7 @@ export class DatabaseUI extends DisposableObject {
token
);
} else {
await this.setCurrentDatabase(uri);
await this.setCurrentDatabase(progress, token, uri);
}
} catch (e) {
// rethrow and let this be handled by default error handling.
@@ -593,15 +597,17 @@ export class DatabaseUI extends DisposableObject {
};
private handleRemoveDatabase = async (
progress: ProgressCallback,
token: CancellationToken,
databaseItem: DatabaseItem,
multiSelect: DatabaseItem[] | undefined
): Promise<void> => {
if (multiSelect?.length) {
multiSelect.forEach((dbItem) =>
this.databaseManager.removeDatabaseItem(dbItem)
);
Promise.all(multiSelect.map((dbItem) =>
this.databaseManager.removeDatabaseItem(progress, token, dbItem)
));
} else {
this.databaseManager.removeDatabaseItem(databaseItem);
await this.databaseManager.removeDatabaseItem(progress, token, databaseItem);
}
};
@@ -651,11 +657,13 @@ export class DatabaseUI extends DisposableObject {
}
private async setCurrentDatabase(
progress: ProgressCallback,
token: CancellationToken,
uri: Uri
): Promise<DatabaseItem | undefined> {
let databaseItem = this.databaseManager.findDatabaseItem(uri);
if (databaseItem === undefined) {
databaseItem = await this.databaseManager.openDatabase(uri);
databaseItem = await this.databaseManager.openDatabase(progress, token, uri);
}
await this.databaseManager.setCurrentDatabaseItem(databaseItem);
@@ -680,7 +688,7 @@ export class DatabaseUI extends DisposableObject {
if (byFolder) {
const fixedUri = await this.fixDbUri(uri);
// we are selecting a database folder
return await this.setCurrentDatabase(fixedUri);
return await this.setCurrentDatabase(progress, token, fixedUri);
} else {
// we are selecting a database archive. Must unzip into a workspace-controlled area
// before importing.

140
extensions/ql-vscode/src/databases.ts
View File

@@ -4,11 +4,12 @@ import * as path from 'path';
import * as vscode from 'vscode';
import * as cli from './cli';
import { ExtensionContext } from 'vscode';
import { showAndLogErrorMessage, showAndLogWarningMessage, showAndLogInformationMessage } from './helpers';
import { showAndLogErrorMessage, showAndLogWarningMessage, showAndLogInformationMessage, ProgressCallback, withProgress } from './helpers';
import { zipArchiveScheme, encodeArchiveBasePath, decodeSourceArchiveUri, encodeSourceArchiveUri } from './archive-filesystem-provider';
import { DisposableObject } from './vscode-utils/disposable-object';
import { QueryServerConfig } from './config';
import { Logger, logger } from './logging';
import { registerDatabases, Dataset, deregisterDatabases } from './pure/messages';
import { QueryServerClient } from './queryserver-client';
/**
* databases.ts
@@ -249,6 +250,11 @@ export interface DatabaseItem {
* Holds if `uri` belongs to this database's source archive.
*/
belongsToSourceArchiveExplorerUri(uri: vscode.Uri): boolean;
/**
* Gets the state of this database, to be persisted in the workspace state.
*/
getPersistedState(): PersistedDatabaseItem;
}
export enum DatabaseEventKind {
@@ -479,13 +485,13 @@ export class DatabaseManager extends DisposableObject {
private readonly _onDidChangeCurrentDatabaseItem = this.push(new vscode.EventEmitter<DatabaseChangedEvent>());
readonly onDidChangeCurrentDatabaseItem = this._onDidChangeCurrentDatabaseItem.event;
private readonly _databaseItems: DatabaseItemImpl[] = [];
private readonly _databaseItems: DatabaseItem[] = [];
private _currentDatabaseItem: DatabaseItem | undefined = undefined;
constructor(
private ctx: ExtensionContext,
public config: QueryServerConfig,
public logger: Logger
private readonly ctx: ExtensionContext,
private readonly qs: QueryServerClient,
public readonly logger: Logger
) {
super();
@@ -493,7 +499,10 @@ export class DatabaseManager extends DisposableObject {
}
public async openDatabase(
uri: vscode.Uri, options?: DatabaseOptions
progress: ProgressCallback,
token: vscode.CancellationToken,
uri: vscode.Uri,
options?: DatabaseOptions
): Promise<DatabaseItem> {
const contents = await resolveDatabaseContents(uri);
@@ -509,7 +518,8 @@ export class DatabaseManager extends DisposableObject {
const databaseItem = new DatabaseItemImpl(uri, contents, fullOptions, (event) => {
this._onDidChangeDatabaseItem.fire(event);
});
await this.addDatabaseItem(databaseItem);
await this.addDatabaseItem(progress, token, databaseItem);
await this.addDatabaseSourceArchiveFolder(databaseItem);
return databaseItem;
@@ -555,6 +565,8 @@ export class DatabaseManager extends DisposableObject {
}
private async createDatabaseItemFromPersistedState(
progress: ProgressCallback,
token: vscode.CancellationToken,
state: PersistedDatabaseItem
): Promise<DatabaseItem> {
@@ -581,33 +593,50 @@ export class DatabaseManager extends DisposableObject {
(event) => {
this._onDidChangeDatabaseItem.fire(event);
});
await this.addDatabaseItem(item);
await this.addDatabaseItem(progress, token, item);
return item;
}
private async loadPersistedState(): Promise<void> {
const currentDatabaseUri = this.ctx.workspaceState.get<string>(CURRENT_DB);
const databases = this.ctx.workspaceState.get<PersistedDatabaseItem[]>(DB_LIST, []);
try {
for (const database of databases) {
const databaseItem = await this.createDatabaseItemFromPersistedState(database);
return withProgress({
location: vscode.ProgressLocation.Notification
},
async (progress, token) => {
const currentDatabaseUri = this.ctx.workspaceState.get<string>(CURRENT_DB);
const databases = this.ctx.workspaceState.get<PersistedDatabaseItem[]>(DB_LIST, []);
let step = 0;
progress({
maxStep: databases.length,
message: 'Loading persisted databases',
step
});
try {
await databaseItem.refresh();
if (currentDatabaseUri === database.uri) {
this.setCurrentDatabaseItem(databaseItem, true);
for (const database of databases) {
progress({
maxStep: databases.length,
message: `Loading ${database.options?.displayName || 'databases'}`,
step: ++step
});
const databaseItem = await this.createDatabaseItemFromPersistedState(progress, token, database);
try {
await databaseItem.refresh();
await this.registerDatabases(progress, token, databaseItem);
if (currentDatabaseUri === database.uri) {
this.setCurrentDatabaseItem(databaseItem, true);
}
}
catch (e) {
// When loading from persisted state, leave invalid databases in the list. They will be
// marked as invalid, and cannot be set as the current database.
}
}
} catch (e) {
// database list had an unexpected type - nothing to be done?
showAndLogErrorMessage(`Database list loading failed: ${e.message}`);
}
catch (e) {
// When loading from persisted state, leave invalid databases in the list. They will be
// marked as invalid, and cannot be set as the current database.
}
}
} catch (e) {
// database list had an unexpected type - nothing to be done?
showAndLogErrorMessage(`Database list loading failed: ${e.message}`);
}
});
}
public get databaseItems(): readonly DatabaseItem[] {
@@ -618,8 +647,10 @@ export class DatabaseManager extends DisposableObject {
return this._currentDatabaseItem;
}
public async setCurrentDatabaseItem(item: DatabaseItem | undefined,
skipRefresh = false): Promise<void> {
public async setCurrentDatabaseItem(
item: DatabaseItem | undefined,
skipRefresh = false
): Promise<void> {
if (!skipRefresh && (item !== undefined)) {
await item.refresh(); // Will throw on invalid database.
@@ -627,6 +658,7 @@ export class DatabaseManager extends DisposableObject {
if (this._currentDatabaseItem !== item) {
this._currentDatabaseItem = item;
this.updatePersistedCurrentDatabaseItem();
this._onDidChangeCurrentDatabaseItem.fire({
item,
kind: DatabaseEventKind.Change
@@ -653,9 +685,20 @@ export class DatabaseManager extends DisposableObject {
return this._databaseItems.find(item => item.sourceArchive && item.sourceArchive.toString(true) === uriString);
}
private async addDatabaseItem(item: DatabaseItemImpl) {
private async addDatabaseItem(
progress: ProgressCallback,
token: vscode.CancellationToken,
item: DatabaseItem
) {
this._databaseItems.push(item);
this.updatePersistedDatabaseList();
// Add this database item to the allow-list
// Database items reconstituted from persisted state
// will not have their contents yet.
if (item.contents?.datasetUri) {
await this.registerDatabases(progress, token, item);
}
// note that we use undefined as the item in order to reset the entire tree
this._onDidChangeDatabaseItem.fire({
item: undefined,
@@ -673,7 +716,11 @@ export class DatabaseManager extends DisposableObject {
});
}
public removeDatabaseItem(item: DatabaseItem) {
public async removeDatabaseItem(
progress: ProgressCallback,
token: vscode.CancellationToken,
item: DatabaseItem
) {
if (this._currentDatabaseItem == item) {
this._currentDatabaseItem = undefined;
}
@@ -700,6 +747,9 @@ export class DatabaseManager extends DisposableObject {
e => logger.log(`Failed to delete '${item.databaseUri.fsPath}'. Reason: ${e.message}`));
}
// Remove this database item from the allow-list
await this.deregisterDatabases(progress, token, item);
// note that we use undefined as the item in order to reset the entire tree
this._onDidChangeDatabaseItem.fire({
item: undefined,
@@ -707,6 +757,34 @@ export class DatabaseManager extends DisposableObject {
});
}
private async deregisterDatabases(
progress: ProgressCallback,
token: vscode.CancellationToken,
dbItem: DatabaseItem,
) {
if (dbItem.contents) {
const databases: Dataset[] = [{
dbDir: dbItem.contents.datasetUri.fsPath,
workingSet: 'default'
}];
await this.qs.sendRequest(deregisterDatabases, { databases }, token, progress);
}
}
private async registerDatabases(
progress: ProgressCallback,
token: vscode.CancellationToken,
dbItem: DatabaseItem,
) {
if (dbItem.contents) {
const databases: Dataset[] = [{
dbDir: dbItem.contents.datasetUri.fsPath,
workingSet: 'default'
}];
await this.qs.sendRequest(registerDatabases, { databases }, token, progress);
}
}
private updatePersistedCurrentDatabaseItem(): void {
this.ctx.workspaceState.update(CURRENT_DB, this._currentDatabaseItem ?
this._currentDatabaseItem.databaseUri.toString(true) : undefined);

2
extensions/ql-vscode/src/extension.ts
View File

@@ -339,7 +339,7 @@ async function activateWithInstalledDistribution(
await qs.startQueryServer();
logger.log('Initializing database manager.');
const dbm = new DatabaseManager(ctx, qlConfigurationListener, logger);
const dbm = new DatabaseManager(ctx, qs, logger);
ctx.subscriptions.push(dbm);
logger.log('Initializing database panel.');
const databaseUI = new DatabaseUI(

30
extensions/ql-vscode/src/pure/messages.ts
View File

@@ -842,7 +842,6 @@ export interface RunUpgradeParams {
toRun: CompiledUpgrades;
}
/**
* The result of running an upgrade
*/
@@ -862,6 +861,21 @@ export interface RunUpgradeResult {
finalSha: string;
}
export interface RegisterDatabaseParams {
databases: Dataset[];
}
export interface DeregisterDatabaseParams {
databases: Dataset[];
}
export type RegisterDatabaseResult = {
registeredDatabases: Dataset[];
};
export type DeregisterDatabaseResult = {
registeredDatabases: Dataset[];
};
/**
* Type for any action that could have progress messages.
@@ -939,6 +953,20 @@ export const runQueries = new rpc.RequestType<WithProgressId<EvaluateQueriesPara
*/
export const runUpgrade = new rpc.RequestType<WithProgressId<RunUpgradeParams>, RunUpgradeResult, void, void>('evaluation/runUpgrade');
export const registerDatabases = new rpc.RequestType<
WithProgressId<RegisterDatabaseParams>,
RegisterDatabaseResult,
void,
void
>('evaluation/registerDatabases');
export const deregisterDatabases = new rpc.RequestType<
WithProgressId<DeregisterDatabaseParams>,
DeregisterDatabaseResult,
void,
void
>('evaluation/deregisterDatabases');
/**
* Request returned to the client to notify completion of a query.
* The full runQueries job is completed when all queries are acknowledged.

4
extensions/ql-vscode/src/queryserver-client.ts
View File

@@ -104,6 +104,10 @@ export class QueryServerClient extends DisposableObject {
private async startQueryServerImpl(progressReporter: ProgressReporter): Promise<void> {
const ramArgs = await this.cliServer.resolveRam(this.config.queryMemoryMb, progressReporter);
const args = ['--threads', this.config.numThreads.toString()].concat(ramArgs);
// TODO: This should only be added after an appropriate version check
args.push('--require-db-registration');
if (this.config.debug) {
args.push('--debug', '--tuple-counting');
}

103
extensions/ql-vscode/src/vscode-tests/minimal-workspace/databases.test.ts
View File

@@ -5,7 +5,7 @@ import * as tmp from 'tmp';
import * as fs from 'fs-extra';
import * as path from 'path';
import { expect } from 'chai';
import { ExtensionContext, Uri, workspace } from 'vscode';
import { CancellationToken, ExtensionContext, Uri, workspace } from 'vscode';
import {
DatabaseEventKind,
@@ -15,9 +15,11 @@ import {
isLikelyDbLanguageFolder,
FullDatabaseOptions
} from '../../databases';
import { QueryServerConfig } from '../../config';
import { Logger } from '../../logging';
import { encodeArchiveBasePath, encodeSourceArchiveUri } from '../../archive-filesystem-provider';
import { QueryServerClient } from '../../queryserver-client';
import { ProgressCallback } from '../../helpers';
import { registerDatabases } from '../../pure/messages';
describe('databases', () => {
@@ -30,6 +32,7 @@ describe('databases', () => {
let updateSpy: sinon.SinonSpy;
let getSpy: sinon.SinonStub;
let dbChangedHandler: sinon.SinonSpy;
let sendRequestSpy: sinon.SinonSpy;
let sandbox: sinon.SinonSandbox;
let dir: tmp.DirResult;
@@ -42,6 +45,8 @@ describe('databases', () => {
sandbox = sinon.createSandbox();
updateSpy = sandbox.spy();
getSpy = sandbox.stub();
getSpy.returns([]);
sendRequestSpy = sandbox.stub();
dbChangedHandler = sandbox.spy();
databaseManager = new DatabaseManager(
{
@@ -53,7 +58,9 @@ describe('databases', () => {
// so that they are deleted upon removal
storagePath: dir.name
} as unknown as ExtensionContext,
{} as QueryServerConfig,
{
sendRequest: sendRequestSpy
} as unknown as QueryServerClient,
{} as Logger,
);
@@ -69,11 +76,15 @@ describe('databases', () => {
sandbox.restore();
});
it('should fire events when adding and removing a db item', () => {
it('should fire events when adding and removing a db item', async () => {
const mockDbItem = createMockDB();
const spy = sinon.spy();
databaseManager.onDidChangeDatabaseItem(spy);
(databaseManager as any).addDatabaseItem(mockDbItem);
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
expect((databaseManager as any)._databaseItems).to.deep.eq([mockDbItem]);
expect(updateSpy).to.have.been.calledWith('databaseList', [{
@@ -88,7 +99,11 @@ describe('databases', () => {
sinon.reset();
// now remove the item
databaseManager.removeDatabaseItem(mockDbItem);
await databaseManager.removeDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem,
);
expect((databaseManager as any)._databaseItems).to.deep.eq([]);
expect(updateSpy).to.have.been.calledWith('databaseList', []);
expect(spy).to.have.been.calledWith({
@@ -97,14 +112,19 @@ describe('databases', () => {
});
});
it('should rename a db item and emit an event', () => {
it('should rename a db item and emit an event', async () => {
const mockDbItem = createMockDB();
const spy = sinon.spy();
databaseManager.onDidChangeDatabaseItem(spy);
(databaseManager as any).addDatabaseItem(mockDbItem);
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
sinon.restore();
databaseManager.renameDatabaseItem(mockDbItem, 'new name');
await databaseManager.renameDatabaseItem(mockDbItem, 'new name');
expect(mockDbItem.name).to.eq('new name');
expect(updateSpy).to.have.been.calledWith('databaseList', [{
@@ -124,7 +144,11 @@ describe('databases', () => {
databaseManager.onDidChangeDatabaseItem(spy);
const mockDbItem = createMockDB();
await (databaseManager as any).addDatabaseItem(mockDbItem);
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
expect(databaseManager.databaseItems).to.deep.eq([mockDbItem]);
expect(updateSpy).to.have.been.calledWith('databaseList', [{
@@ -161,10 +185,19 @@ describe('databases', () => {
// pretend that this item is the first workspace folder in the list
sandbox.stub(mockDbItem, 'belongsToSourceArchiveExplorerUri').returns(true);
await (databaseManager as any).addDatabaseItem(mockDbItem);
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
updateSpy.resetHistory();
await databaseManager.removeDatabaseItem(mockDbItem);
await databaseManager.removeDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
expect(databaseManager.databaseItems).to.deep.eq([]);
expect(updateSpy).to.have.been.calledWith('databaseList', []);
@@ -182,13 +215,21 @@ describe('databases', () => {
// pretend that this item is the first workspace folder in the list
sandbox.stub(mockDbItem, 'belongsToSourceArchiveExplorerUri').returns(true);
await (databaseManager as any).addDatabaseItem(mockDbItem);
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
updateSpy.resetHistory();
// pretend that the database location is not controlled by the extension
(databaseManager as any).ctx.storagePath = 'hucairz';
await databaseManager.removeDatabaseItem(mockDbItem);
await databaseManager.removeDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
expect(databaseManager.databaseItems).to.deep.eq([]);
expect(updateSpy).to.have.been.calledWith('databaseList', []);
@@ -198,6 +239,37 @@ describe('databases', () => {
// should NOT delete the db contents
expect(fs.remove).not.to.have.been.called;
});
it('should register and deregister a database when adding and removing it', async () => {
// similar test as above, but also check the call to sendRequestSpy to make sure they send the
// registration messages.
const mockDbItem = createMockDB();
const registration = {
databases: [{
dbDir: mockDbItem.contents!.datasetUri.fsPath,
workingSet: 'default'
}]
};
sandbox.stub(fs, 'remove').resolves();
await (databaseManager as any).addDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
// Should have registered this database
expect(sendRequestSpy).to.have.been.calledWith(registerDatabases, registration, {}, {});
await databaseManager.removeDatabaseItem(
{} as ProgressCallback,
{} as CancellationToken,
mockDbItem
);
// Should have deregistered this database
expect(sendRequestSpy).to.have.been.calledWith(registerDatabases, registration, {}, {});
});
});
describe('resolveSourceFile', () => {
@@ -292,7 +364,8 @@ describe('databases', () => {
return new DatabaseItemImpl(
databaseUri,
{
sourceArchiveUri
sourceArchiveUri,
datasetUri: databaseUri
} as DatabaseContents,
MOCK_DB_OPTIONS,
dbChangedHandler