Augments the add database command to handle zip files

The add database command can now add databases by zip file. When a
file is selected, the zip file is attempted to be extracted into a
directory managed by the extension.

Once extracted, a database is searched for, by looking for a .dbinfo
file.

Crucially, we are using the same infrastructure to download a database
as we are to add a database by zip file.

Author: aeisenberg

extensions/ql-vscode/src/databaseFetcher.ts

@@ -3,7 +3,7 @@ import * as unzipper from "unzipper";
 import { Uri, ProgressOptions, ProgressLocation, commands, window } from "vscode";
 import * as fs from "fs-extra";
 import * as path from "path";
-import { DatabaseManager } from "./databases";
+import { DatabaseManager, DatabaseItem } from "./databases";
 import { ProgressCallback, showAndLogErrorMessage, withProgress } from "./helpers";
 
 export default async function promptFetchDatabase(dbm: DatabaseManager, storagePath: string) {
@@ -28,13 +28,13 @@ export default async function promptFetchDatabase(dbm: DatabaseManager, storageP
   }
 }
 
-async function databaseFetcher(
+export async function databaseFetcher(
   databaseUrl: string,
   databasesManager: DatabaseManager,
   storagePath: string,
-  progressCallback: ProgressCallback
-): Promise<void> {
-  progressCallback({
+  progressCallback?: ProgressCallback
+): Promise<DatabaseItem> {
+  progressCallback?.({
     maxStep: 3,
     message: 'Downloading database',
     step: 1
@@ -45,36 +45,27 @@ async function databaseFetcher(
   await fs.ensureDir(storagePath);
   const unzipPath = await getStorageFolder(storagePath, databaseUrl);
 
-  const response = await fetch.default(databaseUrl);
-  const unzipStream = unzipper.Extract({
-    path: unzipPath
-  });
-  progressCallback({
-    maxStep: 3,
-    message: 'Unzipping database',
-    step: 2
-  });
-  await new Promise((resolve, reject) => {
-    response.body.on('error', reject);
-    unzipStream.on('error', reject);
-    unzipStream.on('close', resolve);
-    response.body.pipe(unzipStream);
-  });
-  progressCallback({
+  if (isFile(databaseUrl)) {
+    await readAndUnzip(databaseUrl, unzipPath);
+  } else {
+    await fetchAndUnzip(databaseUrl, unzipPath, progressCallback);
+  }
+
+  progressCallback?.({
     maxStep: 3,
     message: 'Opening database',
     step: 3
   });
 
-  // if there is a single directory inside, then assume that's what we want to import
-  const dirs = await fs.readdir(unzipPath);
-  const dbPath = dirs?.length === 1 && (await fs.stat(path.join(unzipPath, dirs[0]))).isDirectory
-    ? path.join(unzipPath, dirs[0])
-    : unzipPath;
-
-  // might need to upgrade before importing...
-  const item = await databasesManager.openDatabase(Uri.parse(dbPath));
-  databasesManager.setCurrentDatabaseItem(item);
+  const dbPath = await findDirWithFile(unzipPath, '.dbinfo');
+  if (dbPath) {
+    // might need to upgrade before importing...
+    const item = await databasesManager.openDatabase(Uri.parse(dbPath));
+    databasesManager.setCurrentDatabaseItem(item);
+    return item;
+  } else {
+    throw new Error('Database not found in archive.');
+  }
 }
 
 async function getStorageFolder(storagePath: string, urlStr: string) {
@@ -110,3 +101,68 @@ function validateUrl(databaseUrl: string) {
     throw new Error('Must use https for downloading a database.');
   }
 }
+
+async function readAndUnzip(databaseUrl: string, unzipPath: string) {
+  const unzipStream = unzipper.Extract({
+    path: unzipPath,
+    verbose: true
+  });
+
+  await new Promise((resolve, reject) => {
+    // we already know this is a file scheme
+    const databaseFile = Uri.parse(databaseUrl).path;
+    const stream = fs.createReadStream(databaseFile);
+    stream.on('error', reject);
+    unzipStream.on('error', reject);
+    unzipStream.on('close', resolve);
+    stream.pipe(unzipStream);
+  });
+}
+
+async function fetchAndUnzip(databaseUrl: string, unzipPath: string, progressCallback?: ProgressCallback) {
+  const response = await fetch.default(databaseUrl);
+  const unzipStream = unzipper.Extract({
+    path: unzipPath
+  });
+  progressCallback?.({
+    maxStep: 3,
+    message: 'Unzipping database',
+    step: 2
+  });
+  await new Promise((resolve, reject) => {
+    response.body.on('error', reject);
+    unzipStream.on('error', reject);
+    unzipStream.on('close', resolve);
+    response.body.pipe(unzipStream);
+  });
+}
+
+function isFile(databaseUrl: string) {
+  return Uri.parse(databaseUrl).scheme === 'file';
+}
+
+/**
+ * Recursively looks for a file in a directory. If the file exists, then returns the directory containing the file.
+ *
+ * @param dir The directory to search
+ * @param toFind The file to recursively look for in this directory
+ *
+ * @returns the directory containing the file, or undefined if not found.
+ */
+async function findDirWithFile(dir: string, toFind: string): Promise<string | undefined> {
+  if (!(await fs.stat(dir)).isDirectory()) {
+    return;
+  }
+  const files = await fs.readdir(dir);
+  if (files.includes(toFind)) {
+    return dir;
+  }
+  for (const file of files) {
+    const newPath = path.join(dir, file);
+    const result = await findDirWithFile(newPath, toFind);
+    if (result) {
+      return result;
+    }
+  }
+  return;
+}

extensions/ql-vscode/src/databases-ui.ts

@@ -8,6 +8,7 @@ import { logger } from './logging';
 import { clearCacheInDatabase, UserCancellationException } from './run-queries';
 import * as qsClient from './queryserver-client';
 import { upgradeDatabase } from './upgrades';
+import { databaseFetcher } from './databaseFetcher';
 
 type ThemableIconPath = { light: string; dark: string } | string;
 
@@ -149,10 +150,11 @@ function getFirst(list: Uri[] | undefined): Uri | undefined {
  */
 async function chooseDatabaseDir(): Promise<Uri | undefined> {
   const chosen = await window.showOpenDialog({
-    openLabel: 'Choose Database',
+    openLabel: 'Choose Database folder or archive',
     canSelectFiles: true,
     canSelectFolders: true,
-    canSelectMany: false
+    canSelectMany: false,
+
   });
   return getFirst(chosen);
 }
@@ -164,7 +166,8 @@ export class DatabaseUI extends DisposableObject {
     ctx: ExtensionContext,
     private cliserver: cli.CodeQLCliServer,
     private databaseManager: DatabaseManager,
-    private readonly queryServer: qsClient.QueryServerClient | undefined
+    private readonly queryServer: qsClient.QueryServerClient | undefined,
+    private readonly storagePath: string
   ) {
     super();
 
@@ -189,7 +192,12 @@ export class DatabaseUI extends DisposableObject {
   }
 
   private handleChooseDatabase = async (): Promise<DatabaseItem | undefined> => {
-    return await this.chooseAndSetDatabase();
+    try {
+      return await this.chooseAndSetDatabase();
+    } catch (e) {
+      showAndLogErrorMessage(e.message);
+      return undefined;
+    }
   }
 
   private handleSortByName = async () => {
@@ -326,7 +334,11 @@ export class DatabaseUI extends DisposableObject {
    */
   private async chooseAndSetDatabase(): Promise<DatabaseItem | undefined> {
     const uri = await chooseDatabaseDir();
-    if (uri !== undefined) {
+
+    if (uri?.path.endsWith('.zip')) {
+      return await databaseFetcher(uri.toString(), this.databaseManager, this.storagePath);
+    }
+    else if (uri !== undefined) {
       return await this.setCurrentDatabase(uri);
     }
     else {

extensions/ql-vscode/src/extension.ts

@@ -265,7 +265,7 @@ async function activateWithInstalledDistribution(ctx: ExtensionContext, distribu
 
   const dbm = new DatabaseManager(ctx, qlConfigurationListener, logger);
   ctx.subscriptions.push(dbm);
-  const databaseUI = new DatabaseUI(ctx, cliServer, dbm, qs);
+  const databaseUI = new DatabaseUI(ctx, cliServer, dbm, qs, getContextStoragePath(ctx));
   ctx.subscriptions.push(databaseUI);
 
   const queryHistoryConfigurationListener = new QueryHistoryConfigListener();