Add option to pass ContainerClient instead of connection string, cont…

…ainer, and credentials to AzureBlobCacheStorage (#521)
microsoft · Dec 27, 2024 · d823606 · d823606
1 parent 150cc9d
commit d823606
Show file tree

Hide file tree

Showing 4 changed files with 53 additions and 37 deletions.
diff --git a/change/change-2416297f-f036-427e-bc7e-7693f953608c.json b/change/change-2416297f-f036-427e-bc7e-7693f953608c.json
@@ -0,0 +1,18 @@
+{
+  "changes": [
+    {
+      "type": "minor",
+      "comment": "Adding an option to pass ContainerClient instead of connection string, container and credentials to AzureBlobCacheStorage",
+      "packageName": "backfill-cache",
+      "email": "[email protected]",
+      "dependentChangeType": "patch"
+    },
+    {
+      "type": "minor",
+      "comment": "Adding an option to pass ContainerClient instead of connection string, container and credentials to AzureBlobCacheStorage",
+      "packageName": "backfill-config",
+      "email": "[email protected]",
+      "dependentChangeType": "patch"
+    }
+  ]
+}
diff --git a/packages/cache/src/AzureBlobCacheStorage.ts b/packages/cache/src/AzureBlobCacheStorage.ts
@@ -6,7 +6,7 @@ import { Logger } from "backfill-logger";
 import { AzureBlobCacheStorageOptions } from "backfill-config";
 
 import { stat } from "fs-extra";
-import { TokenCredential } from "@azure/core-http";
+import { ContainerClient } from "@azure/storage-blob";
 import { CacheStorage } from "./CacheStorage";
 
 const ONE_MEGABYTE = 1024 * 1024;
@@ -69,42 +69,38 @@ const uploadOptions = {
   maxBuffers: 5,
 };
 
-async function createBlobClient(
-  connectionString: string,
-  containerName: string,
-  blobName: string,
-  credential: TokenCredential | undefined
-) {
-  // This is delay loaded because it's very slow to parse
-  const { BlobServiceClient } = await import("@azure/storage-blob");
-  const blobServiceClient = credential
-    ? new BlobServiceClient(connectionString, credential)
-    : BlobServiceClient.fromConnectionString(connectionString);
-
-  const containerClient = blobServiceClient.getContainerClient(containerName);
-  const blobClient = containerClient.getBlobClient(blobName);
-
-  return blobClient;
-}
-
 export class AzureBlobCacheStorage extends CacheStorage {
+  private readonly getContainerClient: () => Promise<ContainerClient>;
+
   constructor(
     private options: AzureBlobCacheStorageOptions,
     logger: Logger,
     cwd: string,
     incrementalCaching = false
   ) {
     super(logger, cwd, incrementalCaching);
+
+    if ("containerClient" in options) {
+      this.getContainerClient = () => Promise.resolve(options.containerClient);
+    } else {
+      const { connectionString, container, credential } = options;
+      // This is delay loaded because it's very slow to parse
+      this.getContainerClient = () =>
+        import("@azure/storage-blob").then(({ BlobServiceClient }) => {
+          const blobServiceClient = credential
+            ? new BlobServiceClient(connectionString, credential)
+            : BlobServiceClient.fromConnectionString(connectionString);
+
+          const containerClient =
+            blobServiceClient.getContainerClient(container);
+          return containerClient;
+        });
+    }
   }
 
   protected async _fetch(hash: string): Promise<boolean> {
     try {
-      const blobClient = await createBlobClient(
-        this.options.connectionString,
-        this.options.container,
-        hash,
-        this.options.credential
-      );
+      const blobClient = (await this.getContainerClient()).getBlobClient(hash);
 
       // If a maxSize has been specified, make sure to check the properties for the size before transferring
       if (this.options.maxSize) {
@@ -166,12 +162,7 @@ export class AzureBlobCacheStorage extends CacheStorage {
   }
 
   protected async _put(hash: string, filesToCache: string[]): Promise<void> {
-    const blobClient = await createBlobClient(
-      this.options.connectionString,
-      this.options.container,
-      hash,
-      this.options.credential
-    );
+    const blobClient = (await this.getContainerClient()).getBlobClient(hash);
 
     const blockBlobClient = blobClient.getBlockBlobClient();
 

diff --git a/packages/config/package.json b/packages/config/package.json
@@ -17,6 +17,7 @@
   },
   "dependencies": {
     "@azure/core-http": "^3.0.0",
+    "@azure/storage-blob": "^12.15.0",
     "backfill-logger": "^5.2.1",
     "find-up": "^5.0.0",
     "pkg-dir": "^4.2.0"

diff --git a/packages/config/src/cacheConfig.ts b/packages/config/src/cacheConfig.ts
@@ -1,17 +1,23 @@
 import { TokenCredential } from "@azure/core-http";
+import { ContainerClient } from "@azure/storage-blob";
 import { Logger } from "backfill-logger";
 
 export interface ICacheStorage {
   fetch: (hash: string) => Promise<boolean>;
   put: (hash: string, filesToCache: string[]) => Promise<void>;
 }
 
-export type AzureBlobCacheStorageOptions = {
-  connectionString: string;
-  container: string;
-  maxSize?: number;
-  credential?: TokenCredential;
-};
+export type AzureBlobCacheStorageOptions =
+  | {
+      connectionString: string;
+      container: string;
+      maxSize?: number;
+      credential?: TokenCredential;
+    }
+  | {
+      containerClient: ContainerClient;
+      maxSize?: number;
+    };
 
 export type NpmCacheStorageOptions = {
   npmPackageName: string;