Move Azure cache implementation to plugin

docs/docs/guides/remote-cache.md

@@ -12,17 +12,23 @@ The theory is that when the CI job runs, it'll produce a "last known good" cache
 
 ## Setting up remote cache - Azure Blob Storage
 
-Follow these steps to set up a remote cache.
+Azure Blob Storage cache is available as a plugin: `@lage-run/azure-blob-cache-storage`. This plugin must be installed separately.
 
-### 1. Upgrade to latest `lage`
+### 1. Install the plugin
+
+```
+yarn add @lage-run/azure-blob-cache-storage
+```
+
+### 2. Upgrade to latest `lage`
 
 See the [migration guide](../cookbook/migration.mdx) for more details.
 
 ```
 yarn upgrade lage
 ```
 
-### 2. Create `.env` and add to `.gitignore`
+### 3. Create `.env` and add to `.gitignore`
 
 Create the file:
 
@@ -39,7 +45,7 @@ lib
 dist
 ```
 
-### 3. Generate auth tokens from Azure storage account
+### 4. Generate auth tokens from Azure storage account
 
 Prerequisite is to have a working Storage Account with Blob Storage Container created. Note that container name, it'll be needed for Step 5.
 
@@ -51,7 +57,7 @@ Prerequisite is to have a working Storage Account with Blob Storage Container cr
 6. Click "show keys"
 7. Save the "connection string" - this is your **read-write** connection string (alternatively, you can create a read-write SAS connection string)
 
-### 4. Modify the `.env` file with the remote cache connection information
+### 5. Modify the `.env` file with the remote cache connection information
 
 ```txt title=".env"
 ## This is required as of right now
@@ -61,7 +67,24 @@ BACKFILL_CACHE_PROVIDER="azure-blob"
 BACKFILL_CACHE_PROVIDER_OPTIONS={"connectionString":"the **read-only** connection string","container":"CONTAINER NAME"}
 ```
 
-### 5. Create a "secret" in the CI system for a Read/Write token
+Alternatively, you can configure it directly in `lage.config.js`:
+
+```js title="lage.config.js"
+module.exports = {
+  cacheOptions: {
+    cacheStorageConfig: {
+      provider: "custom",
+      plugin: "@lage-run/azure-blob-cache-storage",
+      options: {
+        connectionString: "...",
+        container: "...",
+      },
+    },
+  },
+};
+```
+
+### 6. Create a "secret" in the CI system for a Read/Write token
 
 Here's an example snippet of Github Action with the correct environment variable set:
 
@@ -81,7 +104,7 @@ Create a secret named "BACKFILL_CACHE_PROVIDER_OPTIONS":
 
 `process.env.BACKFILL_CACHE_PROVIDER_OPTIONS`is evaluated via backfill (see [`getEnvConfig()`](https://github.com/microsoft/lage/blob/master/packages/backfill-config/src/envConfig.ts#L82) in `backfill-config`).
 
-For "azure-blob" cache provider with a non-sas/key-based `connectionString`(storage account endpoint) requiring azure identity authentication do not use `BACKFILL_CACHE_PROVIDER_OPTIONS`, instead populate the required env variables according to the desired identity/environment. (See [Azure Idenity SDK](https://learn.microsoft.com/en-us/javascript/api/overview/azure/identity-readme)) and set `credentialName` property in the `lage.config.js` under `cacheOptions.cacheStorageConfig.options.credentialName` or via env var `AZURE_IDENTITY_CREDENTIAL_NAME` Supported options are:
+For the Azure Blob cache provider with a non-sas/key-based `connectionString` (storage account endpoint) requiring Azure Identity authentication, you can pass a `credentialName` option in the plugin config or via the `AZURE_IDENTITY_CREDENTIAL_NAME` environment variable. Supported options are:
 
 - `"azure-cli"`
 - `"managed-identity"`

docs/docs/reference/config.md

@@ -69,9 +69,10 @@ const config = {
   cacheOptions: {
     /** @see https://www.npmjs.com/package/backfill#configuration */
     cacheStorageConfig: {
-      // use this to specify a remote cache provider such as "azure-blob",
-      provider: "azure-blob",
-      // there are specific options here for each cache provider
+      // use this to specify a remote cache plugin such as "@lage-run/azure-blob-cache-storage",
+      provider: "custom",
+      plugin: "@lage-run/azure-blob-cache-storage",
+      // there are specific options here for each cache plugin
       options: {}
     },
 

packages/azure-blob-cache-storage/.npmignore

@@ -0,0 +1,5 @@
+**/*
+!lib/**/*
+lib/**/__tests__/*
+lib/**/*.d.ts.map
+!bin/**/*

packages/azure-blob-cache-storage/package.json

@@ -0,0 +1,41 @@
+{
+  "name": "@lage-run/azure-blob-cache-storage",
+  "version": "0.1.0",
+  "description": "Azure Blob Storage cache plugin for backfill/lage",
+  "license": "MIT",
+  "repository": {
+    "type": "git",
+    "url": "https://github.com/microsoft/lage"
+  },
+  "homepage": "https://microsoft.github.io/lage/",
+  "main": "lib/index.js",
+  "types": "lib/index.d.ts",
+  "scripts": {
+    "build": "yarn types && yarn transpile",
+    "transpile": "monorepo-scripts transpile",
+    "types": "yarn run -T tsc",
+    "lint": "monorepo-scripts lint"
+  },
+  "dependencies": {
+    "@azure/core-auth": "1.9.0",
+    "@azure/identity": "4.9.1",
+    "@azure/storage-blob": "12.27.0",
+    "@lage-run/globby": "workspace:^",
+    "backfill-config": "workspace:^",
+    "backfill-logger": "workspace:^",
+    "fs-extra": "8.1.0",
+    "tar-fs": "2.1.4"
+  },
+  "devDependencies": {
+    "@lage-run/monorepo-scripts": "workspace:^",
+    "@types/fs-extra": "^8.0.0",
+    "@types/tar-fs": "^2.0.1"
+  },
+  "engines": {
+    "node": ">=14"
+  },
+  "files": [
+    "lib/!(__*)",
+    "lib/!(__*)/**"
+  ]
+}

packages/azure-blob-cache-storage/src/CacheStorage.ts

@@ -0,0 +1,79 @@
+import { globAsync } from "@lage-run/globby";
+
+import type { Logger } from "backfill-logger";
+import type { ICacheStorage } from "backfill-config";
+import { getFileHash } from "./hashFile.js";
+
+// First key is the hash, second key is the file relative path
+const savedHashes: Map<string, Map<string, string>> = new Map();
+
+// contract: cwd should be absolute
+// The return keys are relative path with posix file separators
+async function getHashesFor(cwd: string): Promise<Map<string, string>> {
+  const result = new Map<string, string>();
+
+  const allFiles = await globAsync(["**/*", "!node_modules"], { cwd });
+  //globby returns relative path with posix file separator
+  await Promise.all(
+    allFiles.map(async (f) => {
+      const hash = await getFileHash(cwd, f);
+      result.set(f, hash);
+    })
+  );
+
+  return result;
+}
+
+export type { ICacheStorage };
+
+export abstract class CacheStorage implements ICacheStorage {
+  public constructor(
+    protected logger: Logger,
+    protected cwd: string,
+    private incrementalCaching = false
+  ) {}
+  public async fetch(hash: string): Promise<boolean> {
+    const tracer = this.logger.setTime("fetchTime");
+
+    const result = await this._fetch(hash);
+
+    tracer.stop();
+
+    this.logger.setHit(result);
+
+    if (!result && this.incrementalCaching) {
+      savedHashes.set(hash, await getHashesFor(this.cwd));
+    }
+
+    return result;
+  }
+
+  public async put(hash: string, outputGlob: string[]): Promise<void> {
+    const tracer = this.logger.setTime("putTime");
+
+    const filesMatchingOutputGlob = await globAsync(outputGlob, {
+      cwd: this.cwd,
+    });
+
+    let filesToCache = filesMatchingOutputGlob;
+    if (this.incrementalCaching) {
+      // Get the list of files that have not changed so we don't need to cache them.
+      const hashesNow = await getHashesFor(this.cwd);
+      const hashesThen =
+        (await savedHashes.get(hash)) || new Map<string, string>();
+      const unchangedFiles = [...hashesThen.keys()].filter(
+        (s) => hashesThen.get(s) === hashesNow.get(s)
+      );
+      filesToCache = filesMatchingOutputGlob.filter(
+        (f) => !unchangedFiles.includes(f)
+      );
+    }
+
+    await this._put(hash, filesToCache);
+    tracer.stop();
+  }
+
+  protected abstract _fetch(hash: string): Promise<boolean>;
+
+  protected abstract _put(hash: string, filesToCache: string[]): Promise<void>;
+}

packages/cache/src/CredentialCache.ts → packages/azure-blob-cache-storage/src/CredentialCache.ts

@@ -6,11 +6,10 @@ import {
   EnvironmentCredential,
   WorkloadIdentityCredential,
 } from "@azure/identity";
-import type { AzureCredentialName } from "@lage-run/config";
-/**
- * Exhaustive credential factory map keyed by AzureCredentialName.
- * This enforces compile-time alignment with the AzureCredentialName union and provides a single source of truth.
- */
+
+/** Allowed credential names matching camelCase of @azure/identity credential class names */
+export type AzureCredentialName = "environment" | "workload-identity" | "managed-identity" | "visual-studio-code" | "azure-cli";
+
 type CredentialFactoryMap = { [K in AzureCredentialName]: () => TokenCredential };
 const CREDENTIAL_FACTORY: CredentialFactoryMap = {
   environment: () => new EnvironmentCredential(),
@@ -23,13 +22,8 @@ const CREDENTIAL_FACTORY: CredentialFactoryMap = {
 export class CredentialCache {
   private static cache: Map<AzureCredentialName, TokenCredential> = new Map();
 
-  // Expose the list for runtime validation elsewhere (derived from the exhaustive factory above)
   public static readonly credentialNames: readonly AzureCredentialName[] = Object.keys(CREDENTIAL_FACTORY) as AzureCredentialName[];
 
-  /**
-   * Returns a credential instance based on the provided name. Results are cached per name.
-   * If no name is provided, EnvironmentCredential is used by default.
-   */
   public static getInstance(credentialName?: AzureCredentialName): TokenCredential {
     const key = (credentialName ?? "environment") as AzureCredentialName;
     const existing = this.cache.get(key);

packages/azure-blob-cache-storage/src/hashFile.ts

@@ -0,0 +1,56 @@
+import * as path from "path";
+import { promises as fs } from "fs";
+import * as crypto from "crypto";
+import pLimit from "p-limit";
+
+let MAX_FILE_OPERATIONS = 5000;
+
+try {
+  const maxFileOpEnv = process.env["BACKFILL_MAX_FILE_OP"];
+  if (maxFileOpEnv) {
+    MAX_FILE_OPERATIONS = parseInt(maxFileOpEnv);
+  }
+} catch (_) {
+  /* The env variable is not an integer, this is fine.*/
+}
+
+const diskLimit = pLimit(MAX_FILE_OPERATIONS);
+
+// The first key is the file path, the second key is mtime
+const memo = new Map<string, Map<number, string>>();
+
+async function computeHash(filePath: string): Promise<string> {
+  const fileBuffer = await diskLimit(() => {
+    return fs.readFile(filePath);
+  });
+  // We use sha1 for perf reason and because the hashing is not used for security reason.
+  const hashSum = crypto.createHash("sha1");
+  hashSum.update(fileBuffer);
+  const hash = hashSum.digest("hex");
+  return hash;
+}
+
+/*
+ * Get the hash of a file.
+ * This function memoizes the hash for files and mtimes.
+ */
+export async function getFileHash(
+  cwd: string,
+  filePath: string
+): Promise<string> {
+  const fileAbsPath = path.join(cwd, filePath);
+  const stat = await fs.stat(fileAbsPath);
+
+  let memoForFile = memo.get(fileAbsPath);
+  if (!memoForFile) {
+    memoForFile = new Map<number, string>();
+    memo.set(fileAbsPath, memoForFile);
+  }
+
+  let hash = memoForFile.get(stat.mtimeMs);
+  if (!hash) {
+    hash = await computeHash(fileAbsPath);
+    memoForFile.set(stat.mtimeMs, hash);
+  }
+  return hash;
+}

packages/azure-blob-cache-storage/src/index.ts

@@ -0,0 +1,51 @@
+import type { Logger } from "backfill-logger";
+import type {
+  ICacheStorage,
+  CustomCacheStoragePlugin,
+  AzureBlobCacheStorageOptions,
+  AzureBlobCacheStorageConnectionStringOptions,
+} from "backfill-config";
+
+import { AzureBlobCacheStorage } from "./AzureBlobCacheStorage.js";
+import { CredentialCache, type AzureCredentialName } from "./CredentialCache.js";
+
+export type { AzureCredentialName } from "./CredentialCache.js";
+export { CredentialCache } from "./CredentialCache.js";
+
+export type AzureBlobPluginOptions = AzureBlobCacheStorageOptions & {
+  /** Optional credential name for Azure Identity authentication. */
+  credentialName?: AzureCredentialName;
+};
+
+function isTokenConnectionString(connectionString: string) {
+  return connectionString.includes("SharedAccessSignature") || connectionString.includes("AccountKey");
+}
+
+const plugin: CustomCacheStoragePlugin<AzureBlobPluginOptions> = {
+  name: "azure-blob",
+  getProvider(logger: Logger, cwd: string, options: AzureBlobPluginOptions): ICacheStorage {
+    // Handle credential injection for connection-string-based options
+    if ("connectionString" in options && !isTokenConnectionString(options.connectionString)) {
+      const connStringOptions = options as AzureBlobCacheStorageConnectionStringOptions & { credentialName?: AzureCredentialName };
+      if (!connStringOptions.credential) {
+        const credName = connStringOptions.credentialName
+          ?? (process.env.AZURE_IDENTITY_CREDENTIAL_NAME || undefined);
+
+        if (credName != null) {
+          if (!CredentialCache.credentialNames.includes(credName as AzureCredentialName)) {
+            throw new Error(
+              `Invalid credentialName: "${credName}". Allowed values: ${CredentialCache.credentialNames.join(", ")}`
+            );
+          }
+          connStringOptions.credential = CredentialCache.getInstance(credName as AzureCredentialName);
+        } else {
+          connStringOptions.credential = CredentialCache.getInstance();
+        }
+      }
+    }
+
+    return new AzureBlobCacheStorage(options, logger, cwd);
+  },
+};
+
+export default plugin;

packages/azure-blob-cache-storage/tsconfig.json

@@ -0,0 +1,7 @@
+{
+  "extends": "@lage-run/monorepo-scripts/config/tsconfig.base.json",
+  "compilerOptions": {
+    "outDir": "lib"
+  },
+  "include": ["src"]
+}

packages/backfill-cache/package.json

@@ -18,7 +18,6 @@
     "test": "yarn run -T jest"
   },
   "dependencies": {
-    "@azure/storage-blob": "12.27.0",
     "@lage-run/globby": "workspace:^",
     "backfill-config": "workspace:^",
     "backfill-logger": "workspace:^",