feat(redis): add saveRawAsBinary driver option (unjs#559)

cjpearson · Jan 3, 2025 · 5a70fc6 · 5a70fc6
1 parent cc0b0ca
commit 5a70fc6
Show file tree

Hide file tree

Showing 2 changed files with 130 additions and 13 deletions.
diff --git a/src/drivers/redis.ts b/src/drivers/redis.ts
@@ -1,4 +1,5 @@
-import { defineDriver, joinKeys } from "./utils";
+import type { TransactionOptions } from "../types";
+import { defineDriver, joinKeys, createError } from "./utils";
 // TODO: use named import in v2
 import Redis, {
   Cluster,
@@ -32,6 +33,12 @@ export interface RedisOptions extends _RedisOptions {
    * Default TTL for all items in seconds.
    */
   ttl?: number;
+
+  /**
+   * Whether to save Buffer/Uint8Arry as binary data or a base64-encoded string
+   * This option applies to the experimental getItemRaw/setItemRaw methods
+   */
+  saveRawAsBinary?: boolean;
 }
 
 const DRIVER_NAME = "redis";
@@ -56,25 +63,67 @@ export default defineDriver((opts: RedisOptions) => {
   const p = (...keys: string[]) => joinKeys(base, ...keys); // Prefix a key. Uses base for backwards compatibility
   const d = (key: string) => (base ? key.replace(base, "") : key); // Deprefix a key
 
+  const getItem = async (key: string) => {
+    const value = await getRedisClient().get(p(key));
+    return value ?? null;
+  };
+
+  const setItem = async (
+    key: string,
+    value: any,
+    tOptions: TransactionOptions
+  ) => {
+    const ttl = tOptions?.ttl ?? opts.ttl;
+    if (ttl) {
+      await getRedisClient().set(p(key), value, "EX", ttl);
+    } else {
+      await getRedisClient().set(p(key), value);
+    }
+  };
+
+  const getItemRaw = async (key: string) => {
+    const value = await getRedisClient().getBuffer(p(key));
+    return value ?? null;
+  };
+
+  const setItemRaw = async (
+    key: string,
+    value: Uint8Array,
+    tOptions: TransactionOptions
+  ) => {
+    let valueToSave: Buffer;
+    if (value instanceof Uint8Array) {
+      if (value instanceof Buffer) {
+        valueToSave = value;
+      } else {
+        valueToSave = Buffer.copyBytesFrom(
+          value,
+          value.byteOffset,
+          value.byteLength
+        );
+      }
+    } else {
+      throw createError(DRIVER_NAME, "Expected Buffer or Uint8Array");
+    }
+
+    await setItem(key, valueToSave, tOptions);
+  };
+
   return {
     name: DRIVER_NAME,
     options: opts,
     getInstance: getRedisClient,
     async hasItem(key) {
       return Boolean(await getRedisClient().exists(p(key)));
     },
-    async getItem(key) {
-      const value = await getRedisClient().get(p(key));
-      return value ?? null;
-    },
-    async setItem(key, value, tOptions) {
-      const ttl = tOptions?.ttl ?? opts.ttl;
-      if (ttl) {
-        await getRedisClient().set(p(key), value, "EX", ttl);
-      } else {
-        await getRedisClient().set(p(key), value);
-      }
-    },
+    getItem,
+    setItem,
+    ...(opts.saveRawAsBinary
+      ? {
+          getItemRaw,
+          setItemRaw,
+        }
+      : {}),
     async removeItem(key) {
       await getRedisClient().del(p(key));
     },

diff --git a/test/drivers/redis.test.ts b/test/drivers/redis.test.ts
@@ -32,9 +32,77 @@ describe("drivers: redis", () => {
         await client.disconnect();
       });
 
+      it("saves raw data as a base64 string", async () => {
+        const helloBuffer = Buffer.from("Hello, world!", "utf8");
+        const byteArray = new Uint8Array(4);
+        byteArray[0] = 2;
+        byteArray[1] = 0;
+        byteArray[2] = 2;
+        byteArray[3] = 5;
+
+        await ctx.storage.setItemRaw("s4:a", helloBuffer);
+        await ctx.storage.setItemRaw("s5:a", byteArray);
+
+        const client = new ioredis.default("ioredis://localhost:6379/0");
+
+        const bufferValue = await client.get("test:s4:a");
+        expect(bufferValue).toEqual("base64:SGVsbG8sIHdvcmxkIQ==");
+
+        const byteArrayValue = await client.get("test:s5:a");
+        expect(byteArrayValue).toEqual("base64:AgACBQ==");
+
+        await client.disconnect();
+      });
+
       it("exposes instance", () => {
         expect(driver.getInstance?.()).toBeInstanceOf(ioredis.default);
       });
     },
   });
+
+  const binaryDriver = redisDriver({
+    base: "test:",
+    url: "ioredis://localhost:6379/0",
+    lazyConnect: false,
+    saveRawAsBinary: true,
+  });
+
+  testDriver({
+    driver: binaryDriver,
+    additionalTests(ctx) {
+      it("saves raw data as binary", async () => {
+        const helloBuffer = Buffer.from("Hello, world!", "utf8");
+        const byteArray = new Uint8Array(4);
+        byteArray[0] = 2;
+        byteArray[1] = 0;
+        byteArray[2] = 2;
+        byteArray[3] = 5;
+
+        await ctx.storage.setItemRaw("s4:a", helloBuffer);
+        await ctx.storage.setItemRaw("s5:a", byteArray);
+
+        const client = new ioredis.default("ioredis://localhost:6379/0");
+
+        const bufferValue = await client.getBuffer("test:s4:a");
+        expect(bufferValue).toEqual(helloBuffer);
+
+        const byteArrayValue = await client.getBuffer("test:s5:a");
+        expect(byteArrayValue).toEqual(Buffer.from([2, 0, 2, 5]));
+
+        await client.disconnect();
+      });
+
+      it("expects binary data to be sent to setItemRaw", async () => {
+        expect(() =>
+          ctx.storage.setItemRaw("s4:a", "Hello, world!")
+        ).rejects.toThrow("Expected Buffer or Uint8Array");
+        expect(() => ctx.storage.setItemRaw("s5:a", 100)).rejects.toThrow(
+          "Expected Buffer or Uint8Array"
+        );
+        expect(() =>
+          ctx.storage.setItemRaw("s5:a", { foo: "bar" })
+        ).rejects.toThrow("Expected Buffer or Uint8Array");
+      });
+    },
+  });
 });