Split up listing of databases to separate function

Author: koesie10

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-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,22 +11,7 @@ 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";
 
 /**
  * Ask whether the user wants to download a database from GitHub.

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

@@ -6,14 +6,12 @@ import { redactableError } from "../common/errors";
 import { asError, getErrorMessage } from "../common/helpers-pure";
 import {
   askForGitHubDatabaseDownload,
-  CodeqlDatabase,
   downloadDatabaseFromGitHub,
-  findGitHubDatabasesForRepository,
 } from "./github-database-download";
 import { GitHubDatabaseConfig, GitHubDatabaseConfigListener } from "../config";
 import { DatabaseManager } from "./local-databases";
 import { CodeQLCliServer } from "../codeql-cli/cli";
-import { showNeverAskAgainDialog } from "../common/vscode/dialog";
+import { listDatabases, ListDatabasesResult } from "./github-database-api";
 
 export class GithubDatabaseModule extends DisposableObject {
   private readonly config: GitHubDatabaseConfig;
@@ -91,38 +89,13 @@ export class GithubDatabaseModule extends DisposableObject {
       return;
     }
 
-    const credentials = this.app.credentials;
-
-    const hasAccessToken = !!(await credentials.getExistingAccessToken());
-
-    // If the user does not have an access token, ask whether they want to connect.
-    if (!hasAccessToken) {
-      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;
-      }
-
-      if (answer === "Never") {
-        await this.config.setDownload("never");
-        return;
-      }
-    }
-
-    const octokit = await credentials.getOctokit();
-
-    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(
@@ -138,10 +111,17 @@ 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 (!hasAccessToken) {
+      if (promptedForCredentials) {
         void window.showInformationMessage(
           "The GitHub repository does not have any CodeQL databases.",
         );
@@ -151,7 +131,7 @@ export class GithubDatabaseModule extends DisposableObject {
     }
 
     // If the user already had an access token, first ask if they even want to download the DB.
-    if (hasAccessToken) {
+    if (!promptedForCredentials) {
       if (!(await askForGitHubDatabaseDownload(databases, this.config))) {
         return;
       }