From 90ae63199138b86232fc249ebc86b481ca39dc31 Mon Sep 17 00:00:00 2001
From: Pooya Parsa <pooya@pi0.io>
Date: Tue, 10 Dec 2024 15:06:45 +0100
Subject: [PATCH] feat(cloudflare-r2-binding): allow specify raw type (#519)

---
 docs/2.drivers/cloudflare.md               |  9 ++++++
 src/drivers/cloudflare-r2-binding.ts       | 32 ++++++++++++++++++++--
 test/drivers/cloudflare-r2-binding.test.ts |  5 ++--
 3 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/docs/2.drivers/cloudflare.md b/docs/2.drivers/cloudflare.md
index f1e53846..6ccce461 100644
--- a/docs/2.drivers/cloudflare.md
+++ b/docs/2.drivers/cloudflare.md
@@ -157,3 +157,12 @@ const storage = createStorage({
 
 - `binding`: Bucket binding or name. Default is `BUCKET`.
 - `base`: Prefix all keys with base.
+
+**Transaction options:**
+
+- `getItemRaw(key, { type: "..." })`
+  - `type: "object"`: Return the [R2 object body](https://developers.cloudflare.com/r2/api/workers/workers-api-reference/#r2objectbody-definition).
+  - `type: "stream"`: Return body stream.
+  - `type: "blob"`: Return a `Blob`.
+  - `type: "bytes"`: Return an `Uint8Array`.
+  - `type: "arrayBuffer"`: Return an `ArrayBuffer` (default)
diff --git a/src/drivers/cloudflare-r2-binding.ts b/src/drivers/cloudflare-r2-binding.ts
index 62ff6f89..484b25d9 100644
--- a/src/drivers/cloudflare-r2-binding.ts
+++ b/src/drivers/cloudflare-r2-binding.ts
@@ -47,10 +47,11 @@ export default defineDriver((opts: CloudflareR2Options = {}) => {
       const binding = getR2Binding(opts.binding);
       return binding.get(key, topts).then((r) => r?.text() ?? null);
     },
-    getItemRaw(key, topts) {
+    async getItemRaw(key, topts) {
       key = r(key);
       const binding = getR2Binding(opts.binding);
-      return binding.get(key, topts).then((r) => r?.arrayBuffer());
+      const object = await binding.get(key, topts);
+      return object ? getObjBody(object, topts?.type) : null;
     },
     async setItem(key, value, topts) {
       key = r(key);
@@ -79,3 +80,30 @@ export default defineDriver((opts: CloudflareR2Options = {}) => {
     },
   };
 });
+
+function getObjBody(
+  object: R2ObjectBody,
+  type: "object" | "stream" | "blob" | "arrayBuffer" | "bytes"
+) {
+  switch (type) {
+    case "object": {
+      return object;
+    }
+    case "stream": {
+      return object.body;
+    }
+    case "blob": {
+      return object.blob();
+    }
+    case "arrayBuffer": {
+      return object.arrayBuffer();
+    }
+    case "bytes": {
+      return object.arrayBuffer().then((buffer) => new Uint8Array(buffer));
+    }
+    // TODO: Default to bytes in v2
+    default: {
+      return object.arrayBuffer();
+    }
+  }
+}
diff --git a/test/drivers/cloudflare-r2-binding.test.ts b/test/drivers/cloudflare-r2-binding.test.ts
index 1e9a5d40..be235631 100644
--- a/test/drivers/cloudflare-r2-binding.test.ts
+++ b/test/drivers/cloudflare-r2-binding.test.ts
@@ -24,12 +24,13 @@ describe("drivers: cloudflare-r2-binding", async () => {
         const storageSnapshot = await snapshot(storage, "");
 
         storageSnapshot["base:data:raw.bin"] = (await storage.getItemRaw(
-          "base:data:raw.bin"
+          "base:data:raw.bin",
+          { type: "bytes" }
         )) as any;
 
         expect(storageSnapshot).toMatchInlineSnapshot(`
           {
-            "base:data:raw.bin": ArrayBuffer [
+            "base:data:raw.bin": Uint8Array [
               1,
               2,
               3,