Use credentials for database download in non-canary mode

extensions/ql-vscode/src/databases/github-database-api.ts

@@ -0,0 +1,110 @@
+import { RequestError } from "@octokit/request-error";
+import { Octokit } from "@octokit/rest";
+import { RestEndpointMethodTypes } from "@octokit/plugin-rest-endpoint-methods";
+import { showNeverAskAgainDialog } from "../common/vscode/dialog";
+import { GitHubDatabaseConfig } from "../config";
+import { Credentials } from "../common/authentication";
+import { AppOctokit } from "../common/octokit";
+
+export type CodeqlDatabase =
+  RestEndpointMethodTypes["codeScanning"]["listCodeqlDatabases"]["response"]["data"][number];
+
+/**
+ * Ask the user if they want to connect to GitHub to download CodeQL databases.
+ * This should be used when the user does not have an access token and should
+ * be followed by an access token prompt.
+ */
+async function askForGitHubConnect(
+  config: GitHubDatabaseConfig,
+): Promise<boolean> {
+  const answer = await showNeverAskAgainDialog(
+    "This repository has an origin (GitHub) that may have one or more CodeQL databases. Connect to GitHub and download any existing databases?",
+    false,
+    "Connect",
+    "Not now",
+    "Never",
+  );
+
+  if (answer === "Not now" || answer === undefined) {
+    return false;
+  }
+
+  if (answer === "Never") {
+    await config.setDownload("never");
+    return false;
+  }
+
+  return true;
+}
+
+export type ListDatabasesResult = {
+  /**
+   * Whether the user has been prompted for credentials. This can be used to determine
+   * follow-up actions based on whether the user has already had any feedback.
+   */
+  promptedForCredentials: boolean;
+  databases: CodeqlDatabase[];
+  octokit: Octokit;
+};
+
+/**
+ * List CodeQL databases for a GitHub repository.
+ *
+ * This will first try to fetch the CodeQL databases for the repository with
+ * existing credentials (or none if there are none). If that fails, it will
+ * prompt the user to connect to GitHub and try again.
+ *
+ * If the user does not want to connect to GitHub, this will return `undefined`.
+ */
+export async function listDatabases(
+  owner: string,
+  repo: string,
+  credentials: Credentials,
+  config: GitHubDatabaseConfig,
+): Promise<ListDatabasesResult | undefined> {
+  const hasAccessToken = !!(await credentials.getExistingAccessToken());
+
+  let octokit = hasAccessToken
+    ? await credentials.getOctokit()
+    : new AppOctokit();
+
+  let promptedForCredentials = false;
+
+  let databases: CodeqlDatabase[];
+  try {
+    const response = await octokit.rest.codeScanning.listCodeqlDatabases({
+      owner,
+      repo,
+    });
+    databases = response.data;
+  } catch (e) {
+    // If we get a 404 when we don't have an access token, it might be because
+    // the repository is private/internal. Therefore, we should ask the user
+    // whether they want to connect to GitHub and try again.
+    if (e instanceof RequestError && e.status === 404 && !hasAccessToken) {
+      // Check whether the user wants to connect to GitHub
+      if (!(await askForGitHubConnect(config))) {
+        return;
+      }
+
+      // Prompt for credentials
+      octokit = await credentials.getOctokit();
+
+      promptedForCredentials = true;
+
+      const response = await octokit.rest.codeScanning.listCodeqlDatabases({
+        owner,
+        repo,
+      });
+      databases = response.data;
+    } else {
+      throw e;
+    }
+  }
+
+  return {
+    promptedForCredentials,
+    databases,
+    octokit,
+  };
+}

extensions/ql-vscode/src/databases/github-database-prompt.ts → extensions/ql-vscode/src/databases/github-database-download.ts

@@ -1,5 +1,4 @@
 import { window } from "vscode";
-import { RestEndpointMethodTypes } from "@octokit/plugin-rest-endpoint-methods";
 import { Octokit } from "@octokit/rest";
 import { showNeverAskAgainDialog } from "../common/vscode/dialog";
 import { getLanguageDisplayName } from "../common/query-language";
@@ -12,65 +11,63 @@ import { DatabaseManager } from "./local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
 import { AppCommandManager } from "../common/commands";
 import { GitHubDatabaseConfig } from "../config";
-
-export type CodeqlDatabase =
-  RestEndpointMethodTypes["codeScanning"]["listCodeqlDatabases"]["response"]["data"][number];
-
-export async function findGitHubDatabasesForRepository(
-  octokit: Octokit,
-  owner: string,
-  repo: string,
-): Promise<CodeqlDatabase[]> {
-  const response = await octokit.rest.codeScanning.listCodeqlDatabases({
-    owner,
-    repo,
-  });
-
-  return response.data;
-}
+import type { CodeqlDatabase } from "./github-database-api";
 
 /**
- * Prompt the user to download a database from GitHub and download that database.
+ * Ask whether the user wants to download a database from GitHub.
+ * @return true if the user wants to download a database, false otherwise.
  */
-export async function promptAndDownloadGitHubDatabase(
-  octokit: Octokit,
-  owner: string,
-  repo: string,
+export async function askForGitHubDatabaseDownload(
   databases: CodeqlDatabase[],
   config: GitHubDatabaseConfig,
-  databaseManager: DatabaseManager,
-  storagePath: string,
-  cliServer: CodeQLCliServer,
-  commandManager: AppCommandManager,
-): Promise<void> {
+): Promise<boolean> {
   const languages = databases.map((database) => database.language);
 
   const message =
     databases.length === 1
       ? `This repository has an origin (GitHub) that has a ${getLanguageDisplayName(
           languages[0],
-        )} CodeQL database. Connect to GitHub and download the existing database?`
+        )} CodeQL database. Download the existing database from GitHub?`
       : `This repository has an origin (GitHub) that has ${joinLanguages(
           languages,
-        )} CodeQL databases. Connect to GitHub and download any existing databases?`;
+        )} CodeQL databases. Download any existing databases from GitHub?`;
 
   const answer = await showNeverAskAgainDialog(
     message,
     false,
-    "Connect",
+    "Download",
     "Not now",
     "Never",
   );
 
   if (answer === "Not now" || answer === undefined) {
-    return;
+    return false;
   }
 
   if (answer === "Never") {
     await config.setDownload("never");
-    return;
+    return false;
   }
 
+  return true;
+}
+
+/**
+ * Download a database from GitHub by asking the user for a language and then
+ * downloading the database for that language.
+ */
+export async function downloadDatabaseFromGitHub(
+  octokit: Octokit,
+  owner: string,
+  repo: string,
+  databases: CodeqlDatabase[],
+  databaseManager: DatabaseManager,
+  storagePath: string,
+  cliServer: CodeQLCliServer,
+  commandManager: AppCommandManager,
+): Promise<void> {
+  const languages = databases.map((database) => database.language);
+
   const language = await promptForLanguage(languages, undefined);
   if (!language) {
     return;

extensions/ql-vscode/src/databases/github-database-module.ts

@@ -1,21 +1,17 @@
+import { window } from "vscode";
 import { DisposableObject } from "../common/disposable-object";
 import { App } from "../common/app";
 import { findGitHubRepositoryForWorkspace } from "./github-repository-finder";
 import { redactableError } from "../common/errors";
 import { asError, getErrorMessage } from "../common/helpers-pure";
 import {
-  CodeqlDatabase,
-  findGitHubDatabasesForRepository,
-  promptAndDownloadGitHubDatabase,
-} from "./github-database-prompt";
-import {
-  GitHubDatabaseConfig,
-  GitHubDatabaseConfigListener,
-  isCanary,
-} from "../config";
-import { AppOctokit } from "../common/octokit";
+  askForGitHubDatabaseDownload,
+  downloadDatabaseFromGitHub,
+} from "./github-database-download";
+import { GitHubDatabaseConfig, GitHubDatabaseConfigListener } from "../config";
 import { DatabaseManager } from "./local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
+import { listDatabases, ListDatabasesResult } from "./github-database-api";
 
 export class GithubDatabaseModule extends DisposableObject {
   private readonly config: GitHubDatabaseConfig;
@@ -93,18 +89,13 @@ export class GithubDatabaseModule extends DisposableObject {
       return;
     }
 
-    const credentials = isCanary() ? this.app.credentials : undefined;
-
-    const octokit = credentials
-      ? await credentials.getOctokit()
-      : new AppOctokit();
-
-    let databases: CodeqlDatabase[];
+    let result: ListDatabasesResult | undefined;
     try {
-      databases = await findGitHubDatabasesForRepository(
-        octokit,
+      result = await listDatabases(
         githubRepository.owner,
         githubRepository.name,
+        this.app.credentials,
+        this.config,
       );
     } catch (e) {
       this.app.telemetry?.sendError(
@@ -120,16 +111,37 @@ export class GithubDatabaseModule extends DisposableObject {
       return;
     }
 
+    // This means the user didn't want to connect, so we can just return.
+    if (result === undefined) {
+      return;
+    }
+
+    const { databases, promptedForCredentials, octokit } = result;
+
     if (databases.length === 0) {
+      // If the user didn't have an access token, they have already been prompted,
+      // so we should give feedback.
+      if (promptedForCredentials) {
+        void window.showInformationMessage(
+          "The GitHub repository does not have any CodeQL databases.",
+        );
+      }
+
       return;
     }
 
-    await promptAndDownloadGitHubDatabase(
+    // If the user already had an access token, first ask if they even want to download the DB.
+    if (!promptedForCredentials) {
+      if (!(await askForGitHubDatabaseDownload(databases, this.config))) {
+        return;
+      }
+    }
+
+    await downloadDatabaseFromGitHub(
       octokit,
       githubRepository.owner,
       githubRepository.name,
       databases,
-      this.config,
       this.databaseManager,
       this.databaseStoragePath,
       this.cliServer,