feat: add custom prefix for session key (#665)

src/convenience/session.ts

@@ -118,6 +118,15 @@ export interface SessionOptions<S, C extends Context = Context> {
      * session data.
      */
     initial?: () => S;
+    /**
+     * An optional prefix to prepend to the session key after it was generated.
+     *
+     * This makes it easier to store session data under a namespace. You can
+     * technically achieve the same functionality by returning an already
+     * prefixed key from `getSessionKey`. This option is merely more convenient,
+     * as it does not require you to think about session key generation.
+     */
+    prefix?: string;
     /**
      * This option lets you generate your own session keys per context object.
      * The session key determines how to map the different session objects to
@@ -424,15 +433,28 @@ class PropertySession<O extends {}, P extends keyof O> {
 }
 
 function fillDefaults<S, C extends Context>(opts: SessionOptions<S, C> = {}) {
-    let { getSessionKey = defaultGetSessionKey, initial, storage } = opts;
+    let {
+        prefix = "",
+        getSessionKey = defaultGetSessionKey,
+        initial,
+        storage,
+    } = opts;
     if (storage == null) {
         debug(
             "Storing session data in memory, all data will be lost when the bot restarts.",
         );
         storage = new MemorySessionStorage<S>();
     }
     const custom = getSessionKey !== defaultGetSessionKey;
-    return { initial, storage, getSessionKey, custom };
+    return {
+        initial,
+        storage,
+        getSessionKey: async (ctx: C) => {
+            const key = await getSessionKey(ctx);
+            return key === undefined ? undefined : prefix + key;
+        },
+        custom,
+    };
 }
 
 /** Stores session data per chat by default */

test/convenience/session.test.ts

@@ -126,6 +126,50 @@ describe("session", () => {
         assertEquals(storage.delete.calls[0].args, ["42"]);
     });
 
+    it("should work with custom session keys", async () => {
+        let val: number | undefined = 0;
+        const storage = {
+            read: spy((_key: string) => val),
+            write: spy((_key: string, value: number) => {
+                val = value;
+            }),
+            delete: spy((_key: string) => {
+                val = undefined;
+            }),
+        };
+        type C = Context & SessionFlavor<number>;
+        const composer = new Composer<C>();
+        let ctx = { chatId: 42 } as C;
+        composer.use(
+            session({
+                prefix: "xyz-",
+                getSessionKey: (ctx) =>
+                    ctx.chatId ? (ctx.chatId ** 2).toString() : undefined,
+                storage,
+            }),
+        )
+            .use((ctx) => {
+                if (ctx.session === 0) ctx.session = 1;
+                else ctx.session = null;
+            });
+
+        await composer.middleware()(ctx, next);
+        ctx = { chatId: 42 } as C;
+        await composer.middleware()(ctx, next);
+
+        assertEquals(storage.read.calls.length, 2);
+        assertEquals(storage.read.calls[0].args, ["xyz-1764"]);
+        assertEquals(storage.read.calls[1].args, ["xyz-1764"]);
+        assertEquals(storage.read.calls[0].returned, 0);
+        assertEquals(storage.read.calls[1].returned, 1);
+
+        assertEquals(storage.write.calls.length, 1);
+        assertEquals(storage.write.calls[0].args, ["xyz-1764", 1]);
+
+        assertEquals(storage.delete.calls.length, 1);
+        assertEquals(storage.delete.calls[0].args, ["xyz-1764"]);
+    });
+
     it("should do IO with objects", async () => {
         let val: Record<string, number> | undefined;
         const storage = {