feat(all): atomic get and set (#14)

* atomic get and set * remove import * types * update retun value * types
blindspot-security · Dec 28, 2023 · b6ce030 · b6ce030
1 parent a96284d
commit b6ce030
Show file tree

Hide file tree

Showing 3 changed files with 40 additions and 5 deletions.
diff --git a/src/redis.ts b/src/redis.ts
@@ -3,10 +3,12 @@ import {
   RedisClientType,
   RedisDefaultModules,
   RedisFunctions,
-  RedisModules, RedisScripts
+  RedisModules,
+  RedisScripts
 } from 'redis';
 import {Milliseconds, RedisStore, TConfig, TMset} from './types';
 import {ScanReply} from '@redis/client/dist/lib/commands/SCAN';
+import {RedisCommandRawReply} from "@redis/client/dist/lib/commands";
 
 export const redisStore = async (config?: TConfig): Promise<RedisStore> => {
   const redisCache: RedisClientType<RedisDefaultModules & RedisModules, RedisFunctions, RedisScripts> = createClient(config);
@@ -110,6 +112,16 @@ class buildRedisStoreWithConfig implements RedisStore {
     return await this.redisCache.scan(cursor, { MATCH: pattern, COUNT: count });
   }
 
+  public async atomicGetAndSet(key: string, updateFunction: (val: any) => any): Promise<RedisCommandRawReply> {
+    await this.redisCache.watch(key);
+    const val = await this.get(key);
+    return await this.redisCache.multi().set(key, updateFunction(val)).get(key).exec();
+  }
+
+  public async flushAll() {
+    await this.redisCache.flushAll();
+  }
+
   public getClient(): RedisClientType<RedisDefaultModules & RedisModules, RedisFunctions, RedisScripts> {
     return this.redisCache;
   }

diff --git a/src/types/RedisStore.interface.ts b/src/types/RedisStore.interface.ts
@@ -8,4 +8,9 @@ export interface RedisStore extends Store {
   getClient(): RedisClientType<RedisDefaultModules & RedisModules, RedisFunctions, RedisScripts>;
 
   scan(pattern: string, cursor? :number, count?: number): Promise<ScanReply>;
+
+  atomicGetAndSet(key: string, updateFunction: (val: any) =>  any): Promise<any>;
+
+  flushAll(): Promise<void>
+
 }
diff --git a/test/redis.test.ts b/test/redis.test.ts
@@ -1,6 +1,6 @@
-import {redisStore} from '../src';
-import {RedisStore} from "../src/types";
-import {describe, beforeEach, it, expect} from "vitest";
+import { redisStore } from '../src';
+import { RedisStore } from "../src/types";
+import { describe, beforeEach, it, expect, afterEach } from "vitest";
 
 let redisClient: RedisStore
 const config = {
@@ -12,10 +12,16 @@ const config = {
   db: 0,
   ttl: 1000 * 60,
 };
+
 beforeEach(async () => {
   redisClient = await redisStore(config);
   await redisClient.reset();
 });
+
+afterEach(async () => {
+  await redisClient.flushAll()
+});
+
 describe('Redis Store', () => {
 
   it('should set and get a value', async () => {
@@ -80,7 +86,7 @@ describe('Redis Store', () => {
     expect(retrievedTtl).toBeLessThanOrEqual(ttl / 1000); // Redis returns TTL in seconds
   });
 
-  it(`should return scan result by pattern`, async () => {
+  it('should return scan result by pattern', async () => {
     const key1 = 'ttl:a:b';
     const key2 = 'ttl1:a:b';
     const key3 = 'ttl:a:b1';
@@ -106,4 +112,16 @@ describe('Redis Store', () => {
     expect(thirdScanWithCount.keys).toEqual([]);
     expect(thirdScanWithCount.cursor).equal(0);
   });
+
+  it('should inc value by one', async () => {
+    await redisClient.set('test', { a: 1 });
+    const res = await redisClient.atomicGetAndSet('test', (obj) => {
+      const parsedVal = obj;
+      parsedVal.a = parsedVal.a + 1;
+      return JSON.stringify(parsedVal);
+    });
+    expect(JSON.parse(res[1])).to.deep.equal({ a: 2 });
+    expect(res[0]).to.deep.equal("OK");
+  })
+
 });