From 2fc04c80cc3220d7f4758050cc5ef4cd11024e41 Mon Sep 17 00:00:00 2001
From: Chuan-kai Lin <cklin@github.com>
Date: Tue, 17 Jun 2025 08:09:39 -0700
Subject: [PATCH] Download overlay-base database from actions cache

---
 src/init-action.ts            | 33 +++++++++++++-
 src/overlay-database-utils.ts | 84 +++++++++++++++++++++++++++++++++++
 2 files changed, 116 insertions(+), 1 deletion(-)

diff --git a/src/init-action.ts b/src/init-action.ts
index a3837cbd2..79eb73ba4 100644
--- a/src/init-action.ts
+++ b/src/init-action.ts
@@ -42,7 +42,10 @@ import {
 } from "./init";
 import { Language } from "./languages";
 import { getActionsLogger, Logger } from "./logging";
-import { OverlayDatabaseMode } from "./overlay-database-utils";
+import {
+  downloadOverlayBaseDatabaseFromCache,
+  OverlayDatabaseMode,
+} from "./overlay-database-utils";
 import { getRepositoryNwo } from "./repository";
 import { ToolsSource } from "./setup-codeql";
 import {
@@ -398,6 +401,34 @@ async function run() {
   }
 
   try {
+    if (
+      config.augmentationProperties.overlayDatabaseMode ===
+        OverlayDatabaseMode.Overlay &&
+      config.augmentationProperties.useOverlayDatabaseCaching
+    ) {
+      // OverlayDatabaseMode.Overlay comes in two flavors: with database
+      // caching, or without. The flavor with database caching is intended to be
+      // an "automatic control" mode, which is supposed to be fail-safe. If we
+      // cannot download an overlay-base database, we revert to
+      // OverlayDatabaseMode.None so that the workflow can continue to run.
+      //
+      // The flavor without database caching is intended to be a "manual
+      // control" mode, where the workflow is supposed to make all the
+      // necessary preparations. So, in that mode, we would assume that
+      // everything is in order and let the analysis fail if that turns out not
+      // to be the case.
+      const overlayDatabaseDownloaded =
+        await downloadOverlayBaseDatabaseFromCache(codeql, config, logger);
+      if (!overlayDatabaseDownloaded) {
+        config.augmentationProperties.overlayDatabaseMode =
+          OverlayDatabaseMode.None;
+        logger.info(
+          "No overlay-base database found in cache, " +
+            `reverting overlay database mode to ${OverlayDatabaseMode.None}.`,
+        );
+      }
+    }
+
     if (
       config.augmentationProperties.overlayDatabaseMode !==
       OverlayDatabaseMode.Overlay
diff --git a/src/overlay-database-utils.ts b/src/overlay-database-utils.ts
index deead9835..c88cea261 100644
--- a/src/overlay-database-utils.ts
+++ b/src/overlay-database-utils.ts
@@ -235,6 +235,90 @@ export async function uploadOverlayBaseDatabaseToCache(
   return true;
 }
 
+/**
+ * Downloads the overlay-base database from the GitHub Actions cache. If conditions
+ * for downloading are not met, the function does nothing and returns false.
+ *
+ * @param codeql The CodeQL instance
+ * @param config The configuration object
+ * @param logger The logger instance
+ * @returns A promise that resolves to true if the download was performed and
+ * successfully completed, or false otherwise
+ */
+export async function downloadOverlayBaseDatabaseFromCache(
+  codeql: CodeQL,
+  config: Config,
+  logger: Logger,
+): Promise<boolean> {
+  const overlayDatabaseMode = config.augmentationProperties.overlayDatabaseMode;
+  if (overlayDatabaseMode !== OverlayDatabaseMode.Overlay) {
+    logger.debug(
+      `Overlay database mode is ${overlayDatabaseMode}. ` +
+        "Skip downloading overlay-base database from cache.",
+    );
+    return false;
+  }
+  if (!config.augmentationProperties.useOverlayDatabaseCaching) {
+    logger.debug(
+      "Overlay database caching is disabled. " +
+        "Skip downloading overlay-base database from cache.",
+    );
+    return false;
+  }
+  if (isInTestMode()) {
+    logger.debug(
+      "In test mode. Skip downloading overlay-base database from cache.",
+    );
+    return false;
+  }
+
+  const dbLocation = config.dbLocation;
+  const codeQlVersion = (await codeql.getVersion()).version;
+  const restoreKey = getCacheRestoreKey(config, codeQlVersion);
+
+  logger.info(
+    `Looking in Actions cache for overlay-base database with restore key ${restoreKey}`,
+  );
+
+  try {
+    const foundKey = await withTimeout(
+      MAX_CACHE_OPERATION_MS,
+      actionsCache.restoreCache([dbLocation], restoreKey),
+      () => {
+        logger.info("Timed out downloading overlay-base database from cache");
+      },
+    );
+
+    if (foundKey === undefined) {
+      logger.info("No overlay-base database found in Actions cache");
+      return false;
+    }
+
+    logger.info(
+      `Downloaded overlay-base database in cache with key ${foundKey}`,
+    );
+  } catch (error) {
+    logger.warning(
+      "Failed to download overlay-base database from cache: " +
+        `${error instanceof Error ? error.message : String(error)}`,
+    );
+    return false;
+  }
+
+  const databaseIsValid = checkOverlayBaseDatabase(
+    config,
+    logger,
+    "Downloaded overlay-base database is invalid",
+  );
+  if (!databaseIsValid) {
+    logger.warning("Downloaded overlay-base database failed validation");
+    return false;
+  }
+
+  logger.info(`Successfully downloaded overlay-base database to ${dbLocation}`);
+  return true;
+}
+
 async function generateCacheKey(
   config: Config,
   codeQlVersion: string,