Use new feature flag architecture when uploading databases

2025-12-27 09:40:17 +08:00 · 2021-12-14 20:48:32 +00:00
parent 04671efa1d
commit d6499fad61
18 changed files with 158 additions and 177 deletions
--- a/src/analyze-action-env.test.ts
+++ b/src/analyze-action-env.test.ts
@@ -25,6 +25,8 @@ test("analyze action with RAM & threads from environment variables", async (t) =
      .resolves({} as actionsUtil.StatusReportBase);
    sinon.stub(actionsUtil, "sendStatusReport").resolves(true);
    sinon.stub(configUtils, "getConfig").resolves({
+      // Use GHES so we don't try to call the feature flags API endpoint
+      gitHubVersion: { type: util.GitHubVariant.GHES, version: "3.0.0" },
      languages: [],
    } as unknown as configUtils.Config);
    const requiredInputStub = sinon.stub(actionsUtil, "getRequiredInput");
--- a/src/analyze-action-input.test.ts
+++ b/src/analyze-action-input.test.ts
@@ -25,6 +25,8 @@ test("analyze action with RAM & threads from action inputs", async (t) => {
      .resolves({} as actionsUtil.StatusReportBase);
    sinon.stub(actionsUtil, "sendStatusReport").resolves(true);
    sinon.stub(configUtils, "getConfig").resolves({
+      // Use GHES so we don't try to call the feature flags API endpoint
+      gitHubVersion: { type: util.GitHubVariant.GHES, version: "3.0.0" },
      languages: [],
    } as unknown as configUtils.Config);
    const requiredInputStub = sinon.stub(actionsUtil, "getRequiredInput");
--- a/src/analyze-action.ts
+++ b/src/analyze-action.ts
@@ -15,6 +15,7 @@ import {
 import { CODEQL_VERSION_NEW_TRACING, getCodeQL } from "./codeql";
 import { Config, getConfig } from "./config-utils";
 import { uploadDatabases } from "./database-upload";
+import { GitHubFeatureFlags } from "./feature-flags";
 import { getActionsLogger } from "./logging";
 import { parseRepositoryNwo } from "./repository";
 import * as upload_lib from "./upload-lib";
@@ -99,6 +100,14 @@ async function run() {
    const memory = util.getMemoryFlag(
      actionsUtil.getOptionalInput("ram") || process.env["CODEQL_RAM"]
    );
+
+    const featureFlags = new GitHubFeatureFlags(
+      config.gitHubVersion,
+      apiDetails,
+      logger
+    );
+    void featureFlags.preloadFeatureFlags();
+
    await runFinalize(outputDir, threads, memory, config, logger);
    if (actionsUtil.getRequiredInput("skip-queries") !== "true") {
      runStats = await runQueries(
@@ -176,7 +185,14 @@ async function run() {
    const repositoryNwo = parseRepositoryNwo(
      util.getRequiredEnvParam("GITHUB_REPOSITORY")
    );
-    await uploadDatabases(repositoryNwo, config, apiDetails, logger); // Possibly upload the database bundles for remote queries
+    // Possibly upload the database bundles for remote queries
+    await uploadDatabases(
+      repositoryNwo,
+      config,
+      featureFlags,
+      apiDetails,
+      logger
+    );

    if (
      uploadResult !== undefined &&
--- a/src/database-upload.test.ts
+++ b/src/database-upload.test.ts
@@ -10,6 +10,7 @@ import * as apiClient from "./api-client";
 import { setCodeQL } from "./codeql";
 import { Config } from "./config-utils";
 import { uploadDatabases } from "./database-upload";
+import { createFeatureFlags, FeatureFlags } from "./feature-flags";
 import { Language } from "./languages";
 import { RepositoryNwo } from "./repository";
 import {
@@ -32,6 +33,11 @@ test.beforeEach(() => {
  initializeEnvironment(Mode.actions, "1.2.3");
 });

+const uploadToUploadsDomainFlags = createFeatureFlags([
+  "database_uploads_enabled",
+  "uploads_domain_enabled",
+]);
+
 const testRepoName: RepositoryNwo = { owner: "github", repo: "example" };
 const testApiDetails: GitHubApiDetails = {
  auth: "1234",
@@ -55,44 +61,25 @@ function getTestConfig(tmpDir: string): Config {
  };
 }

-function mockHttpRequests(
-  optInStatusCode: number,
-  useUploadDomain?: boolean,
-  databaseUploadStatusCode?: number
+async function mockHttpRequests(
+  featureFlags: FeatureFlags,
+  databaseUploadStatusCode: number
 ) {
  // Passing an auth token is required, so we just use a dummy value
  const client = github.getOctokit("123");

  const requestSpy = sinon.stub(client, "request");

-  const optInSpy = requestSpy.withArgs(
-    "GET /repos/:owner/:repo/code-scanning/codeql/databases"
-  );
-  if (optInStatusCode < 300) {
-    optInSpy.resolves({
-      status: optInStatusCode,
-      data: {
-        useUploadDomain,
-      },
-      headers: {},
-      url: "GET /repos/:owner/:repo/code-scanning/codeql/databases",
-    });
+  const url = (await featureFlags.getUploadsDomainEnabled())
+    ? "POST https://uploads.github.com/repos/:owner/:repo/code-scanning/codeql/databases/:language?name=:name"
+    : "PUT /repos/:owner/:repo/code-scanning/codeql/databases/:language";
+  const databaseUploadSpy = requestSpy.withArgs(url);
+  if (databaseUploadStatusCode < 300) {
+    databaseUploadSpy.resolves(undefined);
  } else {
-    optInSpy.throws(new HTTPError("some error message", optInStatusCode));
-  }
-
-  if (databaseUploadStatusCode !== undefined) {
-    const url = useUploadDomain
-      ? "POST https://uploads.github.com/repos/:owner/:repo/code-scanning/codeql/databases/:language?name=:name"
-      : "PUT /repos/:owner/:repo/code-scanning/codeql/databases/:language";
-    const databaseUploadSpy = requestSpy.withArgs(url);
-    if (databaseUploadStatusCode < 300) {
-      databaseUploadSpy.resolves(undefined);
-    } else {
-      databaseUploadSpy.throws(
-        new HTTPError("some error message", databaseUploadStatusCode)
-      );
-    }
+    databaseUploadSpy.throws(
+      new HTTPError("some error message", databaseUploadStatusCode)
+    );
  }

  sinon.stub(apiClient, "getApiClient").value(() => client);
@@ -111,6 +98,7 @@ test("Abort database upload if 'upload-database' input set to false", async (t)
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      uploadToUploadsDomainFlags,
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -140,6 +128,7 @@ test("Abort database upload if running against GHES", async (t) => {
    await uploadDatabases(
      testRepoName,
      config,
+      createFeatureFlags([]),
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -169,6 +158,7 @@ test("Abort database upload if running against GHAE", async (t) => {
    await uploadDatabases(
      testRepoName,
      config,
+      createFeatureFlags([]),
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -195,6 +185,7 @@ test("Abort database upload if not analyzing default branch", async (t) => {
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      uploadToUploadsDomainFlags,
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -208,7 +199,7 @@ test("Abort database upload if not analyzing default branch", async (t) => {
  });
 });

-test("Abort database upload if opt-in request returns 404", async (t) => {
+test("Abort database upload if feature flag is disabled", async (t) => {
  await withTmpDir(async (tmpDir) => {
    setupActionsVars(tmpDir, tmpDir);
    sinon
@@ -217,8 +208,6 @@ test("Abort database upload if opt-in request returns 404", async (t) => {
      .returns("true");
    sinon.stub(actionsUtil, "isAnalyzingDefaultBranch").resolves(true);

-    mockHttpRequests(404);
-
    setCodeQL({
      async databaseBundle() {
        return;
@@ -229,6 +218,7 @@ test("Abort database upload if opt-in request returns 404", async (t) => {
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      createFeatureFlags(["uploads_domain_enabled"]),
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -243,41 +233,6 @@ test("Abort database upload if opt-in request returns 404", async (t) => {
  });
 });

-test("Abort database upload if opt-in request fails with something other than 404", async (t) => {
-  await withTmpDir(async (tmpDir) => {
-    setupActionsVars(tmpDir, tmpDir);
-    sinon
-      .stub(actionsUtil, "getRequiredInput")
-      .withArgs("upload-database")
-      .returns("true");
-    sinon.stub(actionsUtil, "isAnalyzingDefaultBranch").resolves(true);
-
-    mockHttpRequests(500);
-
-    setCodeQL({
-      async databaseBundle() {
-        return;
-      },
-    });
-
-    const loggedMessages = [] as LoggedMessage[];
-    await uploadDatabases(
-      testRepoName,
-      getTestConfig(tmpDir),
-      testApiDetails,
-      getRecordingLogger(loggedMessages)
-    );
-    t.assert(
-      loggedMessages.find(
-        (v) =>
-          v.type === "info" &&
-          v.message ===
-            "Skipping database upload due to unknown error: Error: some error message"
-      ) !== undefined
-    );
-  });
-});
-
 test("Don't crash if uploading a database fails", async (t) => {
  await withTmpDir(async (tmpDir) => {
    setupActionsVars(tmpDir, tmpDir);
@@ -287,7 +242,9 @@ test("Don't crash if uploading a database fails", async (t) => {
      .returns("true");
    sinon.stub(actionsUtil, "isAnalyzingDefaultBranch").resolves(true);

-    mockHttpRequests(200, false, 500);
+    const featureFlags = createFeatureFlags(["database_uploads_enabled"]);
+
+    await mockHttpRequests(featureFlags, 500);

    setCodeQL({
      async databaseBundle(_: string, outputFilePath: string) {
@@ -299,9 +256,13 @@ test("Don't crash if uploading a database fails", async (t) => {
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      featureFlags,
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
+
+    console.log(loggedMessages);
+
    t.assert(
      loggedMessages.find(
        (v) =>
@@ -322,7 +283,7 @@ test("Successfully uploading a database to api.github.com", async (t) => {
      .returns("true");
    sinon.stub(actionsUtil, "isAnalyzingDefaultBranch").resolves(true);

-    mockHttpRequests(200, false, 201);
+    await mockHttpRequests(uploadToUploadsDomainFlags, 201);

    setCodeQL({
      async databaseBundle(_: string, outputFilePath: string) {
@@ -334,6 +295,7 @@ test("Successfully uploading a database to api.github.com", async (t) => {
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      uploadToUploadsDomainFlags,
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
@@ -356,7 +318,7 @@ test("Successfully uploading a database to uploads.github.com", async (t) => {
      .returns("true");
    sinon.stub(actionsUtil, "isAnalyzingDefaultBranch").resolves(true);

-    mockHttpRequests(200, true, 201);
+    await mockHttpRequests(uploadToUploadsDomainFlags, 201);

    setCodeQL({
      async databaseBundle(_: string, outputFilePath: string) {
@@ -368,6 +330,7 @@ test("Successfully uploading a database to uploads.github.com", async (t) => {
    await uploadDatabases(
      testRepoName,
      getTestConfig(tmpDir),
+      uploadToUploadsDomainFlags,
      testApiDetails,
      getRecordingLogger(loggedMessages)
    );
--- a/src/database-upload.ts
+++ b/src/database-upload.ts
@@ -4,6 +4,7 @@ import * as actionsUtil from "./actions-util";
 import { getApiClient, GitHubApiDetails } from "./api-client";
 import { getCodeQL } from "./codeql";
 import { Config } from "./config-utils";
+import { FeatureFlags } from "./feature-flags";
 import { Logger } from "./logging";
 import { RepositoryNwo } from "./repository";
 import * as util from "./util";
@@ -12,6 +13,7 @@ import { bundleDb } from "./util";
 export async function uploadDatabases(
  repositoryNwo: RepositoryNwo,
  config: Config,
+  featureFlags: FeatureFlags,
  apiDetails: GitHubApiDetails,
  logger: Logger
 ): Promise<void> {
@@ -32,30 +34,17 @@ export async function uploadDatabases(
    return;
  }

-  const client = getApiClient(apiDetails);
-  let useUploadDomain: boolean;
-  try {
-    const response = await client.request(
-      "GET /repos/:owner/:repo/code-scanning/codeql/databases",
-      {
-        owner: repositoryNwo.owner,
-        repo: repositoryNwo.repo,
-      }
+  if (!(await featureFlags.getDatabaseUploadsEnabled())) {
+    logger.debug(
+      "Repository is not opted in to database uploads. Skipping upload."
    );
-    useUploadDomain = response.data["uploads_domain_enabled"];
-  } catch (e) {
-    if (util.isHTTPError(e) && e.status === 404) {
-      logger.debug(
-        "Repository is not opted in to database uploads. Skipping upload."
-      );
-    } else {
-      console.log(e);
-      logger.info(`Skipping database upload due to unknown error: ${e}`);
-    }
    return;
  }

+  const client = getApiClient(apiDetails);
  const codeql = await getCodeQL(config.codeQLCmd);
+  const useUploadDomain = await featureFlags.getUploadsDomainEnabled();
+
  for (const language of config.languages) {
    // Upload the database bundle.
    // Although we are uploading arbitrary file contents to the API, it's worth
--- a/src/feature-flags.ts
+++ b/src/feature-flags.ts
@@ -3,7 +3,7 @@ import { Logger } from "./logging";
 import { parseRepositoryNwo } from "./repository";
 import * as util from "./util";

-interface FeatureFlags {
+export interface FeatureFlags {
  getDatabaseUploadsEnabled(): Promise<boolean>;
  getMlPoweredQueriesEnabled(): Promise<boolean>;
  getUploadsDomainEnabled(): Promise<boolean>;
@@ -88,3 +88,29 @@ export class GitHubFeatureFlags implements FeatureFlags {
    return apiResponse;
  }
 }
+
+type FeatureFlagName =
+  | "database_uploads_enabled"
+  | "ml_powered_queries_enabled"
+  | "uploads_domain_enabled";
+
+/**
+ * Create a feature flags instance with the specified set of enabled flags.
+ *
+ * This should be only used within tests.
+ */
+export function createFeatureFlags(
+  enabledFlags: FeatureFlagName[]
+): FeatureFlags {
+  return {
+    getDatabaseUploadsEnabled: async () => {
+      return enabledFlags.includes("database_uploads_enabled");
+    },
+    getMlPoweredQueriesEnabled: async () => {
+      return enabledFlags.includes("ml_powered_queries_enabled");
+    },
+    getUploadsDomainEnabled: async () => {
+      return enabledFlags.includes("uploads_domain_enabled");
+    },
+  };
+}