Download databases from GitHub

extensions/ql-vscode/package.json

@@ -45,6 +45,7 @@
     "onCommand:codeQLDatabases.chooseDatabaseFolder",
     "onCommand:codeQLDatabases.chooseDatabaseArchive",
     "onCommand:codeQLDatabases.chooseDatabaseInternet",
+    "onCommand:codeQLDatabases.chooseDatabaseGithub",
     "onCommand:codeQLDatabases.chooseDatabaseLgtm",
     "onCommand:codeQL.setCurrentDatabase",
     "onCommand:codeQL.viewAst",
@@ -54,6 +55,7 @@
     "onCommand:codeQL.chooseDatabaseFolder",
     "onCommand:codeQL.chooseDatabaseArchive",
     "onCommand:codeQL.chooseDatabaseInternet",
+    "onCommand:codeQL.chooseDatabaseGithub",
     "onCommand:codeQL.chooseDatabaseLgtm",
     "onCommand:codeQLDatabases.chooseDatabase",
     "onCommand:codeQLDatabases.setCurrentDatabase",
@@ -356,6 +358,14 @@
           "dark": "media/dark/cloud-download.svg"
         }
       },
+      {
+        "command": "codeQLDatabases.chooseDatabaseGithub",
+        "title": "Download Database from GitHub",
+        "icon": {
+          "light": "media/light/github.svg",
+          "dark": "media/dark/github.svg"
+        }
+      },
       {
         "command": "codeQLDatabases.chooseDatabaseLgtm",
         "title": "Download from LGTM",
@@ -428,6 +438,10 @@
         "command": "codeQL.chooseDatabaseInternet",
         "title": "CodeQL: Download Database"
       },
+      {
+        "command": "codeQL.chooseDatabaseGithub",
+        "title": "CodeQL: Download Database from GitHub"
+      },
       {
         "command": "codeQL.chooseDatabaseLgtm",
         "title": "CodeQL: Download Database from LGTM"
@@ -604,6 +618,11 @@
           "when": "view == codeQLDatabases",
           "group": "navigation"
         },
+        {
+          "command": "codeQLDatabases.chooseDatabaseGithub",
+          "when": "config.codeQL.canary && view == codeQLDatabases",
+          "group": "navigation"
+        },
         {
           "command": "codeQLDatabases.chooseDatabaseLgtm",
           "when": "view == codeQLDatabases",
@@ -829,6 +848,10 @@
           "command": "codeQL.viewCfg",
           "when": "resourceScheme == codeql-zip-archive && config.codeQL.canary"
         },
+        {
+          "command": "codeQL.chooseDatabaseGithub",
+          "when": "config.codeQL.canary"
+        },
         {
           "command": "codeQLDatabases.setCurrentDatabase",
           "when": "false"
@@ -873,6 +896,10 @@
           "command": "codeQLDatabases.chooseDatabaseInternet",
           "when": "false"
         },
+        {
+          "command": "codeQLDatabases.chooseDatabaseGithub",
+          "when": "false"
+        },
         {
           "command": "codeQLDatabases.chooseDatabaseLgtm",
           "when": "false"

extensions/ql-vscode/src/databaseFetcher.ts

@@ -21,6 +21,8 @@ import {
 } from './commandRunner';
 import { logger } from './logging';
 import { tmpDir } from './helpers';
+import { Credentials } from './authentication';
+import { REPO_REGEX } from './pure/helpers-pure';
 
 /**
  * Prompts a user to fetch a database from a remote location. Database is assumed to be an archive file.
@@ -46,6 +48,7 @@ export async function promptImportInternetDatabase(
 
   const item = await databaseArchiveFetcher(
     databaseUrl,
+    {},
     databaseManager,
     storagePath,
     progress,
@@ -61,6 +64,79 @@ export async function promptImportInternetDatabase(
 
 }
 
+/**
+ * Prompts a user to fetch a database from GitHub.
+ * User enters a GitHub repository and then the user is asked which language
+ * to download (if there is more than one)
+ *
+ * @param databaseManager the DatabaseManager
+ * @param storagePath where to store the unzipped database.
+ */
+export async function promptImportGithubDatabase(
+  databaseManager: DatabaseManager,
+  storagePath: string,
+  credentials: Credentials,
+  progress: ProgressCallback,
+  token: CancellationToken,
+  cli?: CodeQLCliServer
+): Promise<DatabaseItem | undefined> {
+  progress({
+    message: 'Choose repository',
+    step: 1,
+    maxStep: 2
+  });
+  const githubRepo = await window.showInputBox({
+    title: 'Enter a GitHub repository in the format <owner>/<repo> (e.g. github/codeql)',
+    placeHolder: '<owner>/<repo>',
+    ignoreFocusOut: true,
+  });
+  if (!githubRepo) {
+    return;
+  }
+
+  if (!REPO_REGEX.test(githubRepo)) {
+    throw new Error(`Invalid GitHub repository: ${githubRepo}`);
+  }
+
+  const databaseUrl = await convertGithubNwoToDatabaseUrl(githubRepo, credentials, progress);
+  if (!databaseUrl) {
+    return;
+  }
+
+  const octokit = await credentials.getOctokit();
+  /**
+   * The 'token' property of the token object returned by `octokit.auth()`.
+   * The object is undocumented, but looks something like this:
+   * {
+   *   token: 'xxxx',
+   *   tokenType: 'oauth',
+   *   type: 'token',
+   * }
+   * We only need the actual token string.
+   */
+  const octokitToken = (await octokit.auth() as { token: string })?.token;
+  if (!octokitToken) {
+    // Just print a generic error message for now. Ideally we could show more debugging info, like the
+    // octokit object, but that would expose a user token.
+    throw new Error('Unable to get GitHub token.');
+  }
+  const item = await databaseArchiveFetcher(
+    databaseUrl,
+    { 'Accept': 'application/zip', 'Authorization': `Bearer ${octokitToken}` },
+    databaseManager,
+    storagePath,
+    progress,
+    token,
+    cli
+  );
+  if (item) {
+    await commands.executeCommand('codeQLDatabases.focus');
+    void showAndLogInformationMessage('Database downloaded and imported successfully.');
+    return item;
+  }
+  return;
+}
+
 /**
  * Prompts a user to fetch a database from lgtm.
  * User enters a project url and then the user is asked which language
@@ -94,6 +170,7 @@ export async function promptImportLgtmDatabase(
     if (databaseUrl) {
       const item = await databaseArchiveFetcher(
         databaseUrl,
+        {},
         databaseManager,
         storagePath,
         progress,
@@ -140,6 +217,7 @@ export async function importArchiveDatabase(
   try {
     const item = await databaseArchiveFetcher(
       databaseUrl,
+      {},
       databaseManager,
       storagePath,
       progress,
@@ -166,13 +244,15 @@ export async function importArchiveDatabase(
  * or in the local filesystem.
  *
  * @param databaseUrl URL from which to grab the database
+ * @param requestHeaders Headers to send with the request
  * @param databaseManager the DatabaseManager
  * @param storagePath where to store the unzipped database.
  * @param progress callback to send progress messages to
  * @param token cancellation token
  */
 async function databaseArchiveFetcher(
   databaseUrl: string,
+  requestHeaders: { [key: string]: string },
   databaseManager: DatabaseManager,
   storagePath: string,
   progress: ProgressCallback,
@@ -193,7 +273,7 @@ async function databaseArchiveFetcher(
   if (isFile(databaseUrl)) {
     await readAndUnzip(databaseUrl, unzipPath, cli, progress);
   } else {
-    await fetchAndUnzip(databaseUrl, unzipPath, cli, progress);
+    await fetchAndUnzip(databaseUrl, requestHeaders, unzipPath, cli, progress);
   }
 
   progress({
@@ -292,6 +372,7 @@ async function readAndUnzip(
 
 async function fetchAndUnzip(
   databaseUrl: string,
+  requestHeaders: { [key: string]: string },
   unzipPath: string,
   cli?: CodeQLCliServer,
   progress?: ProgressCallback
@@ -310,7 +391,10 @@ async function fetchAndUnzip(
     step: 1,
   });
 
-  const response = await checkForFailingResponse(await fetch(databaseUrl), 'Error downloading database');
+  const response = await checkForFailingResponse(
+    await fetch(databaseUrl, { headers: requestHeaders }),
+    'Error downloading database'
+  );
   const archiveFileStream = fs.createWriteStream(archivePath);
 
   const contentLength = response.headers.get('content-length');
@@ -381,6 +465,37 @@ export async function findDirWithFile(
   return;
 }
 
+export async function convertGithubNwoToDatabaseUrl(
+  githubRepo: string,
+  credentials: Credentials,
+  progress: ProgressCallback): Promise<string | undefined> {
+  try {
+    // TODO: In future, we could accept GitHub URLs in addition to NWOs.
+    // Similar to "looksLikeLgtmUrl".
+    if (!REPO_REGEX.test(githubRepo)) {
+      throw new Error('Invalid repository format. Must be in the format <owner>/<repo> (e.g. github/codeql)');
+    }
+
+    const [owner, repo] = githubRepo.split('/');
+
+    const octokit = await credentials.getOctokit();
+    const response = await octokit.request('GET /repos/:owner/:repo/code-scanning/codeql/databases', { owner, repo });
+
+    const languages = response.data.map((db: any) => db.language);
+
+    const language = await promptForLanguage(languages, progress);
+    if (!language) {
+      return;
+    }
+
+    return `https://api.github.com/repos/${owner}/${repo}/code-scanning/codeql/databases/${language}`;
+
+  } catch (e) {
+    void logger.log(`Error: ${e.message}`);
+    throw new Error(`Unable to get database for '${githubRepo}'`);
+  }
+}
+
 /**
  * The URL pattern is https://lgtm.com/projects/{provider}/{org}/{name}/{irrelevant-subpages}.
  * There are several possibilities for the provider: in addition to GitHub.com (g),
@@ -506,7 +621,7 @@ async function promptForLanguage(
     maxStep: 2
   });
   if (!languages.length) {
-    return;
+    throw new Error('No databases found');
   }
   if (languages.length === 1) {
     return languages[0];

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

@@ -33,11 +33,13 @@ import * as qsClient from './queryserver-client';
 import { upgradeDatabaseExplicit } from './upgrades';
 import {
   importArchiveDatabase,
+  promptImportGithubDatabase,
   promptImportInternetDatabase,
   promptImportLgtmDatabase,
 } from './databaseFetcher';
 import { CancellationToken } from 'vscode';
 import { asyncFilter } from './pure/helpers-pure';
+import { Credentials } from './authentication';
 
 type ThemableIconPath = { light: string; dark: string } | string;
 
@@ -219,7 +221,8 @@ export class DatabaseUI extends DisposableObject {
     private databaseManager: DatabaseManager,
     private readonly queryServer: qsClient.QueryServerClient | undefined,
     private readonly storagePath: string,
-    readonly extensionPath: string
+    readonly extensionPath: string,
+    private readonly getCredentials: () => Promise<Credentials>
   ) {
     super();
 
@@ -291,6 +294,20 @@ export class DatabaseUI extends DisposableObject {
         }
       )
     );
+    this.push(
+      commandRunnerWithProgress(
+        'codeQLDatabases.chooseDatabaseGithub',
+        async (
+          progress: ProgressCallback,
+          token: CancellationToken
+        ) => {
+          const credentials = await this.getCredentials();
+          await this.handleChooseDatabaseGithub(credentials, progress, token);
+        },
+        {
+          title: 'Adding database from GitHub',
+        })
+    );
     this.push(
       commandRunnerWithProgress(
         'codeQLDatabases.chooseDatabaseLgtm',
@@ -462,6 +479,21 @@ export class DatabaseUI extends DisposableObject {
     );
   };
 
+  handleChooseDatabaseGithub = async (
+    credentials: Credentials,
+    progress: ProgressCallback,
+    token: CancellationToken
+  ): Promise<DatabaseItem | undefined> => {
+    return await promptImportGithubDatabase(
+      this.databaseManager,
+      this.storagePath,
+      credentials,
+      progress,
+      token,
+      this.queryServer?.cliServer
+    );
+  };
+
   handleChooseDatabaseLgtm = async (
     progress: ProgressCallback,
     token: CancellationToken

extensions/ql-vscode/src/extension.ts

@@ -433,7 +433,8 @@ async function activateWithInstalledDistribution(
     dbm,
     qs,
     getContextStoragePath(ctx),
-    ctx.extensionPath
+    ctx.extensionPath,
+    () => Credentials.initialize(ctx),
   );
   databaseUI.init();
   ctx.subscriptions.push(databaseUI);
@@ -931,6 +932,18 @@ async function activateWithInstalledDistribution(
       title: 'Choose a Database from an Archive'
     })
   );
+  ctx.subscriptions.push(
+    commandRunnerWithProgress('codeQL.chooseDatabaseGithub', async (
+      progress: ProgressCallback,
+      token: CancellationToken
+    ) => {
+      const credentials = await Credentials.initialize(ctx);
+      await databaseUI.handleChooseDatabaseGithub(credentials, progress, token);
+    },
+      {
+        title: 'Adding database from GitHub',
+      })
+  );
   ctx.subscriptions.push(
     commandRunnerWithProgress('codeQL.chooseDatabaseLgtm', (
       progress: ProgressCallback,