Integrated CLI database unbundle command for archive download

Co-authored by: Marc Jaramillo mnj.webdeveloper@gmail.comm
Co-authored by: Musab Guma'a mgsium@github.com

Author: marcnjaramillo

extensions/ql-vscode/CHANGELOG.md

@@ -2,7 +2,8 @@
 
 ## [UNRELEASED]
 
-- Replace certain control codes (`U+0000` - `U+001F`) with their corresponding control labels (`U+2400` - `U+241F`)  in the results view. [#963](https://github.com/github/vscode-codeql/pull/963)
+- Use the CodeQL CLI command `database unbundle` for the archive download option, which ensures even archives over 4GB in size can be downloaded. [#971](https://github.com/github/vscode-codeql/pull/971)
+- Fix a bug with importing large databases. Databases over 4GB can now be imported directly from LGTM or from a zip file. This functionality is only available when using CodeQL CLI version 2.6.0 or later. [#963](https://github.com/github/vscode-codeql/pull/963)
 
 ## 1.5.6 - 07 October 2021
 

extensions/ql-vscode/src/cli.ts

@@ -600,6 +600,15 @@ export class CodeQLCliServer implements Disposable {
     return await this.runJsonCodeQlCliCommand<BQRSInfo>(['bqrs', 'info'], subcommandArgs, 'Reading bqrs header');
   }
 
+  async databaseUnbundle(archivePath: string, target: string, name?: string): Promise<string> {
+    const subcommandArgs = [];
+    if (target) subcommandArgs.push('--target', target);
+    if (name) subcommandArgs.push('--name', name);
+    subcommandArgs.push(archivePath);
+ 
+    return await this.runCodeQlCliCommand(['database', 'unbundle'], subcommandArgs, `Extracting ${archivePath} to directory ${target}`);
+  }
+
   /**
   * Gets the results from a bqrs.
   * @param bqrsPath The path to the bqrs.
@@ -1074,6 +1083,11 @@ export class CliVersionConstraint {
    */
   public static CLI_VERSION_WITH_ALLOW_LIBRARY_PACKS_IN_RESOLVE_QUERIES = new SemVer('2.6.1');
 
+  /**
+   * CLI version where the `database unbundle` subcommand was introduced.
+   */
+  public static CLI_VERSION_WITH_DATABASE_UNBUNDLE = new SemVer('2.6.0');
+
   constructor(private readonly cli: CodeQLCliServer) {
     /**/
   }
@@ -1105,4 +1119,9 @@ export class CliVersionConstraint {
   async supportsDatabaseRegistration() {
     return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_DB_REGISTRATION);
   }
+
+  async supportsDatabaseUnbundle() {
+    return this.isVersionAtLeast(CliVersionConstraint.CLI_VERSION_WITH_DATABASE_UNBUNDLE);
+  }
+
 }

extensions/ql-vscode/src/databaseFetcher.ts

@@ -1,12 +1,13 @@
 import fetch, { Response } from 'node-fetch';
-import * as unzipper from 'unzipper';
 import { zip } from 'zip-a-folder';
+import * as unzipper from 'unzipper';
 import {
   Uri,
   CancellationToken,
   commands,
   window,
 } from 'vscode';
+import { CodeQLCliServer } from './cli';
 import * as fs from 'fs-extra';
 import * as path from 'path';
 
@@ -32,6 +33,7 @@ export async function promptImportInternetDatabase(
   storagePath: string,
   progress: ProgressCallback,
   token: CancellationToken,
+  cli?: CodeQLCliServer
 ): Promise<DatabaseItem | undefined> {
   const databaseUrl = await window.showInputBox({
     prompt: 'Enter URL of zipfile of database to download',
@@ -47,7 +49,8 @@ export async function promptImportInternetDatabase(
     databaseManager,
     storagePath,
     progress,
-    token
+    token,
+    cli
   );
 
   if (item) {
@@ -70,7 +73,8 @@ export async function promptImportLgtmDatabase(
   databaseManager: DatabaseManager,
   storagePath: string,
   progress: ProgressCallback,
-  token: CancellationToken
+  token: CancellationToken,
+  cli?: CodeQLCliServer
 ): Promise<DatabaseItem | undefined> {
   progress({
     message: 'Choose project',
@@ -93,7 +97,8 @@ export async function promptImportLgtmDatabase(
         databaseManager,
         storagePath,
         progress,
-        token
+        token,
+        cli
       );
       if (item) {
         await commands.executeCommand('codeQLDatabases.focus');
@@ -120,14 +125,16 @@ export async function importArchiveDatabase(
   storagePath: string,
   progress: ProgressCallback,
   token: CancellationToken,
+  cli?: CodeQLCliServer,
 ): Promise<DatabaseItem | undefined> {
   try {
     const item = await databaseArchiveFetcher(
       databaseUrl,
       databaseManager,
       storagePath,
       progress,
-      token
+      token,
+      cli
     );
     if (item) {
       await commands.executeCommand('codeQLDatabases.focus');
@@ -159,7 +166,8 @@ async function databaseArchiveFetcher(
   databaseManager: DatabaseManager,
   storagePath: string,
   progress: ProgressCallback,
-  token: CancellationToken
+  token: CancellationToken,
+  cli?: CodeQLCliServer,
 ): Promise<DatabaseItem> {
   progress({
     message: 'Getting database',
@@ -173,9 +181,9 @@ async function databaseArchiveFetcher(
   const unzipPath = await getStorageFolder(storagePath, databaseUrl);
 
   if (isFile(databaseUrl)) {
-    await readAndUnzip(databaseUrl, unzipPath, progress);
+    await readAndUnzip(databaseUrl, unzipPath, cli, progress);
   } else {
-    await fetchAndUnzip(databaseUrl, unzipPath, progress);
+    await fetchAndUnzip(databaseUrl, unzipPath, cli, progress);
   }
 
   progress({
@@ -249,6 +257,7 @@ function validateHttpsUrl(databaseUrl: string) {
 async function readAndUnzip(
   zipUrl: string,
   unzipPath: string,
+  cli?: CodeQLCliServer,
   progress?: ProgressCallback
 ) {
   // TODO: Providing progress as the file is unzipped is currently blocked
@@ -259,16 +268,22 @@ async function readAndUnzip(
     step: 9,
     message: `Unzipping into ${path.basename(unzipPath)}`
   });
-  // Must get the zip central directory since streaming the
-  // zip contents may not have correct local file headers.
-  // Instead, we can only rely on the central directory.
-  const directory = await unzipper.Open.file(zipFile);
-  await directory.extract({ path: unzipPath });
+  if (cli && await cli.cliConstraints.supportsDatabaseUnbundle()) {
+    // Use the `database unbundle` command if the installed cli version supports it
+    await cli.databaseUnbundle(zipFile, unzipPath);
+  } else {
+    // Must get the zip central directory since streaming the
+    // zip contents may not have correct local file headers.
+    // Instead, we can only rely on the central directory.
+    const directory = await unzipper.Open.file(zipFile);
+    await directory.extract({ path: unzipPath });
+  }
 }
 
 async function fetchAndUnzip(
   databaseUrl: string,
   unzipPath: string,
+  cli?: CodeQLCliServer,
   progress?: ProgressCallback
 ) {
   // Although it is possible to download and stream directly to an unzipped directory,
@@ -298,7 +313,8 @@ async function fetchAndUnzip(
       .on('error', reject)
   );
 
-  await readAndUnzip(Uri.file(archivePath).toString(true), unzipPath, progress);
+  await readAndUnzip(Uri.file(archivePath).toString(true), unzipPath, cli, progress);
+
 
   // remove archivePath eagerly since these archives can be large.
   await fs.remove(archivePath);

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

@@ -451,27 +451,46 @@ export class DatabaseUI extends DisposableObject {
   handleChooseDatabaseInternet = async (
     progress: ProgressCallback,
     token: CancellationToken
-  ): Promise<
-    DatabaseItem | undefined
-  > => {
-    return await promptImportInternetDatabase(
-      this.databaseManager,
-      this.storagePath,
-      progress,
-      token
-    );
+  ): Promise<DatabaseItem | undefined> => {
+    if (this.queryServer && await this.queryServer.cliServer.cliConstraints.supportsDatabaseUnbundle()) {
+      return await promptImportInternetDatabase(
+        this.databaseManager,
+        this.storagePath,
+        progress,
+        token,
+        this.queryServer.cliServer
+      );
+    } else {
+      return await promptImportInternetDatabase(        
+        this.databaseManager,
+        this.storagePath,
+        progress,
+        token
+      );
+    }
   };
 
   handleChooseDatabaseLgtm = async (
     progress: ProgressCallback,
     token: CancellationToken
   ): Promise<DatabaseItem | undefined> => {
-    return await promptImportLgtmDatabase(
-      this.databaseManager,
-      this.storagePath,
-      progress,
-      token
-    );
+
+    if (this.queryServer && await this.queryServer.cliServer.cliConstraints.supportsDatabaseUnbundle()) {
+      return await promptImportLgtmDatabase(
+        this.databaseManager,
+        this.storagePath,
+        progress,
+        token,
+        this.queryServer.cliServer
+      );
+    } else {
+      return await promptImportLgtmDatabase(
+        this.databaseManager,
+        this.storagePath,
+        progress,
+        token
+      );
+    }
   };
 
   async tryUpgradeCurrentDatabase(
@@ -572,6 +591,10 @@ export class DatabaseUI extends DisposableObject {
     token: CancellationToken,
     uri: Uri,
   ): Promise<void> => {
+    if (!this.queryServer) {
+      throw new Error('Query server not started');
+    }
+
     try {
       // Assume user has selected an archive if the file has a .zip extension
       if (uri.path.endsWith('.zip')) {
@@ -580,7 +603,8 @@ export class DatabaseUI extends DisposableObject {
           this.databaseManager,
           this.storagePath,
           progress,
-          token
+          token,
+          this.queryServer.cliServer
         );
       } else {
         await this.setCurrentDatabase(progress, token, uri);
@@ -695,8 +719,11 @@ export class DatabaseUI extends DisposableObject {
     progress: ProgressCallback,
     token: CancellationToken,
   ): Promise<DatabaseItem | undefined> {
-    const uri = await chooseDatabaseDir(byFolder);
+    if (!this.queryServer) {
+      throw new Error('Query server not started');
+    }
 
+    const uri = await chooseDatabaseDir(byFolder);
     if (!uri) {
       return undefined;
     }
@@ -713,7 +740,8 @@ export class DatabaseUI extends DisposableObject {
         this.databaseManager,
         this.storagePath,
         progress,
-        token
+        token,
+        this.queryServer.cliServer
       );
     }
   }

extensions/ql-vscode/src/vscode-tests/cli-integration/databases.test.ts

@@ -5,6 +5,7 @@ import { expect } from 'chai';
 import { extensions, CancellationToken, Uri, window } from 'vscode';
 
 import { CodeQLExtensionInterface } from '../../extension';
+import { CodeQLCliServer } from '../../cli';
 import { DatabaseManager } from '../../databases';
 import { promptImportLgtmDatabase, importArchiveDatabase, promptImportInternetDatabase } from '../../databaseFetcher';
 import { ProgressCallback } from '../../commandRunner';
@@ -17,10 +18,11 @@ describe('Databases', function() {
   this.timeout(60000);
 
   const LGTM_URL = 'https://lgtm.com/projects/g/aeisenberg/angular-bind-notifier/';
-
+  
   let databaseManager: DatabaseManager;
   let sandbox: sinon.SinonSandbox;
   let inputBoxStub: sinon.SinonStub;
+  let cli: CodeQLCliServer;
   let progressCallback: ProgressCallback;
 
   beforeEach(async () => {
@@ -53,7 +55,7 @@ describe('Databases', function() {
   it('should add a database from a folder', async () => {
     const progressCallback = sandbox.spy() as ProgressCallback;
     const uri = Uri.file(dbLoc);
-    let dbItem = await importArchiveDatabase(uri.toString(true), databaseManager, storagePath, progressCallback, {} as CancellationToken);
+    let dbItem = await importArchiveDatabase(uri.toString(true), databaseManager, storagePath, progressCallback, {} as CancellationToken, cli);
     expect(dbItem).to.be.eq(databaseManager.currentDatabaseItem);
     expect(dbItem).to.be.eq(databaseManager.databaseItems[0]);
     expect(dbItem).not.to.be.undefined;
@@ -64,7 +66,7 @@ describe('Databases', function() {
 
   it('should add a database from lgtm with only one language', async () => {
     inputBoxStub.resolves(LGTM_URL);
-    let dbItem = await promptImportLgtmDatabase(databaseManager, storagePath, progressCallback, {} as CancellationToken);
+    let dbItem = await promptImportLgtmDatabase(databaseManager, storagePath, progressCallback, {} as CancellationToken, cli);
     expect(dbItem).not.to.be.undefined;
     dbItem = dbItem!;
     expect(dbItem.name).to.eq('aeisenberg_angular-bind-notifier_106179a');
@@ -74,7 +76,7 @@ describe('Databases', function() {
   it('should add a database from a url', async () => {
     inputBoxStub.resolves(DB_URL);
 
-    let dbItem = await promptImportInternetDatabase(databaseManager, storagePath, progressCallback, {} as CancellationToken);
+    let dbItem = await promptImportInternetDatabase(databaseManager, storagePath, progressCallback, {} as CancellationToken, cli);
     expect(dbItem).not.to.be.undefined;
     dbItem = dbItem!;
     expect(dbItem.name).to.eq('db');

extensions/ql-vscode/src/vscode-tests/cli-integration/queries.test.ts

@@ -68,7 +68,8 @@ describe('Queries', function() {
         databaseManager,
         storagePath,
         progress,
-        token
+        token,
+        cli
       );
 
       if (!maybeDbItem) {