Add GT and LT options to ZADD (#781)

* Add GT and LT commands to zadd * Remove unused import * Remove unused codes and format code * Simplify types * Remove unused imports
upstash · Dec 13, 2023 · affbf52 · affbf52
1 parent 9a31dd1
commit affbf52
Show file tree

Hide file tree

Showing 5 changed files with 316 additions and 127 deletions.
diff --git a/pkg/commands/types.ts b/pkg/commands/types.ts
@@ -127,17 +127,15 @@ export { Type, type TypeCommand } from "./type";
 export { type UnlinkCommand } from "./unlink";
 export { type XAddCommand } from "./xadd";
 export { type XRangeCommand } from "./xrange";
-export {
-  ScoreMember,
-  ZAddCommandOptions,
-  ZAddCommandOptionsWithIncr,
-  type ZAddCommand,
-} from "./zadd";
+export { ScoreMember, ZAddCommandOptions, type ZAddCommand } from "./zadd";
 export { type ZCardCommand } from "./zcard";
 export { type ZCountCommand } from "./zcount";
 export { type ZDiffStoreCommand } from "./zdiffstore";
 export { type ZIncrByCommand } from "./zincrby";
-export { type ZInterStoreCommand, ZInterStoreCommandOptions } from "./zinterstore";
+export {
+  type ZInterStoreCommand,
+  ZInterStoreCommandOptions,
+} from "./zinterstore";
 export { type ZLexCountCommand } from "./zlexcount";
 export { type ZMScoreCommand } from "./zmscore";
 export { type ZPopMaxCommand } from "./zpopmax";
@@ -152,4 +150,7 @@ export { type ZRevRankCommand } from "./zrevrank";
 export { type ZScanCommand } from "./zscan";
 export { type ZScoreCommand } from "./zscore";
 export { type ZUnionCommand, ZUnionCommandOptions } from "./zunion";
-export { type ZUnionStoreCommand, ZUnionStoreCommandOptions } from "./zunionstore";
+export {
+  type ZUnionStoreCommand,
+  ZUnionStoreCommandOptions,
+} from "./zunionstore";
diff --git a/pkg/commands/zadd.test.ts b/pkg/commands/zadd.test.ts
@@ -2,6 +2,7 @@ import { keygen, newHttpClient, randomID } from "../test-utils";
 
 import { afterAll, describe, expect, test } from "bun:test";
 import { ZAddCommand } from "./zadd";
+import { ZRangeCommand } from "./zrange";
 import { ZScoreCommand } from "./zscore";
 
 const client = newHttpClient();
@@ -12,60 +13,68 @@ afterAll(cleanup);
 describe("command format", () => {
   describe("without options", () => {
     test("build the correct command", () => {
-      expect(new ZAddCommand(["key", { score: 0, member: "member" }]).command).toEqual([
-        "zadd",
-        "key",
-        0,
-        "member",
-      ]);
+      expect(
+        new ZAddCommand(["key", { score: 0, member: "member" }]).command
+      ).toEqual(["zadd", "key", 0, "member"]);
     });
   });
 
   describe("with nx", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { nx: true }, { score: 0, member: "member" }]).command,
+        new ZAddCommand(["key", { nx: true }, { score: 0, member: "member" }])
+          .command
       ).toEqual(["zadd", "key", "nx", 0, "member"]);
     });
   });
 
   describe("with xx", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { xx: true }, { score: 0, member: "member" }]).command,
+        new ZAddCommand(["key", { xx: true }, { score: 0, member: "member" }])
+          .command
       ).toEqual(["zadd", "key", "xx", 0, "member"]);
     });
   });
 
   describe("with ch", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { ch: true }, { score: 0, member: "member" }]).command,
+        new ZAddCommand(["key", { ch: true }, { score: 0, member: "member" }])
+          .command
       ).toEqual(["zadd", "key", "ch", 0, "member"]);
     });
   });
 
   describe("with incr", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { incr: true }, { score: 0, member: "member" }]).command,
+        new ZAddCommand(["key", { incr: true }, { score: 0, member: "member" }])
+          .command
       ).toEqual(["zadd", "key", "incr", 0, "member"]);
     });
   });
 
   describe("with nx and ch", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { nx: true, ch: true }, { score: 0, member: "member" }]).command,
+        new ZAddCommand([
+          "key",
+          { nx: true, ch: true },
+          { score: 0, member: "member" },
+        ]).command
       ).toEqual(["zadd", "key", "nx", "ch", 0, "member"]);
     });
   });
 
   describe("with nx,ch and incr", () => {
     test("build the correct command", () => {
       expect(
-        new ZAddCommand(["key", { nx: true, ch: true, incr: true }, { score: 0, member: "member" }])
-          .command,
+        new ZAddCommand([
+          "key",
+          { nx: true, ch: true, incr: true },
+          { score: 0, member: "member" },
+        ]).command
       ).toEqual(["zadd", "key", "nx", "ch", "incr", 0, "member"]);
     });
   });
@@ -78,7 +87,7 @@ describe("command format", () => {
           { nx: true },
           { score: 0, member: "member" },
           { score: 1, member: "member1" },
-        ]).command,
+        ]).command
       ).toEqual(["zadd", "key", "nx", 0, "member", 1, "member1"]);
     });
   });
@@ -102,9 +111,11 @@ describe("xx", () => {
       const score = Math.floor(Math.random() * 10);
       await new ZAddCommand([key, { score, member }]).exec(client);
       const newScore = score + 1;
-      const res = await new ZAddCommand([key, { xx: true }, { score: newScore, member }]).exec(
-        client,
-      );
+      const res = await new ZAddCommand([
+        key,
+        { xx: true },
+        { score: newScore, member },
+      ]).exec(client);
       expect(res).toEqual(0);
 
       const res2 = await new ZScoreCommand([key, member]).exec(client);
@@ -118,9 +129,11 @@ describe("xx", () => {
       const score = Math.floor(Math.random() * 10);
       await new ZAddCommand([key, { score, member }]).exec(client);
       const newScore = score + 1;
-      const res = await new ZAddCommand([key, { xx: true }, { score: newScore, member }]).exec(
-        client,
-      );
+      const res = await new ZAddCommand([
+        key,
+        { xx: true },
+        { score: newScore, member },
+      ]).exec(client);
       expect(res).toEqual(0);
     });
   });
@@ -134,9 +147,11 @@ describe("nx", () => {
       const score = Math.floor(Math.random() * 10);
       await new ZAddCommand([key, { score, member }]).exec(client);
       const newScore = score + 1;
-      const res = await new ZAddCommand([key, { nx: true }, { score: newScore, member }]).exec(
-        client,
-      );
+      const res = await new ZAddCommand([
+        key,
+        { nx: true },
+        { score: newScore, member },
+      ]).exec(client);
       expect(res).toEqual(0);
 
       const res2 = await new ZScoreCommand([key, member]).exec(client);
@@ -148,7 +163,11 @@ describe("nx", () => {
       const key = newKey();
       const member = randomID();
       const score = Math.floor(Math.random() * 10);
-      const res = await new ZAddCommand([key, { nx: true }, { score, member }]).exec(client);
+      const res = await new ZAddCommand([
+        key,
+        { nx: true },
+        { score, member },
+      ]).exec(client);
       expect(res).toEqual(1);
     });
   });
@@ -161,9 +180,11 @@ describe("ch", () => {
     const score = Math.floor(Math.random() * 10);
     await new ZAddCommand([key, { score, member }]).exec(client);
     const newScore = score + 1;
-    const res = await new ZAddCommand([key, { ch: true }, { score: newScore, member }]).exec(
-      client,
-    );
+    const res = await new ZAddCommand([
+      key,
+      { ch: true },
+      { score: newScore, member },
+    ]).exec(client);
     expect(res).toEqual(1);
   });
 });
@@ -174,8 +195,108 @@ describe("incr", () => {
     const member = randomID();
     const score = Math.floor(Math.random() * 10);
     await new ZAddCommand([key, { score, member }]).exec(client);
-    const res = await new ZAddCommand([key, { incr: true }, { score: 1, member }]).exec(client);
+    const res = await new ZAddCommand([
+      key,
+      { incr: true },
+      { score: 1, member },
+    ]).exec(client);
     expect(typeof res).toBe("number");
     expect(res).toEqual(score + 1);
   });
 });
+
+describe("LT and GT", () => {
+  describe("GT", () => {
+    test("should replace successfully if greater than", async () => {
+      const key = newKey();
+
+      await new ZAddCommand([key, { score: 1, member: "one" }]).exec(client);
+      await new ZAddCommand([key, { score: 2, member: "two" }]).exec(client);
+      await new ZAddCommand([key, { score: 3, member: "three" }]).exec(client);
+
+      await new ZAddCommand([
+        key,
+        { gt: true },
+        { score: 4, member: "two" },
+      ]).exec(client);
+
+      const res2 = await new ZRangeCommand([
+        key,
+        0,
+        -1,
+        { withScores: true },
+      ]).exec(client);
+
+      expect(res2).toEqual(["one", 1, "three", 3, "two", 4]);
+    });
+
+    test("should fail to replace if its not greater than", async () => {
+      const key = newKey();
+
+      await new ZAddCommand([key, { score: 1, member: "one" }]).exec(client);
+      await new ZAddCommand([key, { score: 2, member: "two" }]).exec(client);
+      await new ZAddCommand([key, { score: 3, member: "three" }]).exec(client);
+
+      await new ZAddCommand([
+        key,
+        { gt: true },
+        { score: 1, member: "two" },
+      ]).exec(client);
+
+      const res2 = await new ZRangeCommand([
+        key,
+        0,
+        -1,
+        { withScores: true },
+      ]).exec(client);
+
+      expect(res2).toEqual(["one", 1, "two", 2, "three", 3]);
+    });
+  });
+
+  describe("LT", () => {
+    test("should replace successfully if less than", async () => {
+      const key = newKey();
+
+      await new ZAddCommand([key, { score: 1, member: "one" }]).exec(client);
+      await new ZAddCommand([key, { score: 2, member: "two" }]).exec(client);
+      await new ZAddCommand([key, { score: 3, member: "three" }]).exec(client);
+      await new ZAddCommand([
+        key,
+        { lt: true },
+        { score: 2, member: "three" },
+      ]).exec(client);
+
+      const res2 = await new ZRangeCommand([
+        key,
+        0,
+        -1,
+        { withScores: true },
+      ]).exec(client);
+      expect(res2).toEqual(["one", 1, "three", 2, "two", 2]);
+    });
+
+    test("should fail to replace if its not less than", async () => {
+      const key = newKey();
+
+      await new ZAddCommand([key, { score: 1, member: "one" }]).exec(client);
+      await new ZAddCommand([key, { score: 2, member: "two" }]).exec(client);
+      await new ZAddCommand([key, { score: 3, member: "three" }]).exec(client);
+
+      await new ZAddCommand([
+        key,
+        { lt: true },
+        { score: 6, member: "two" },
+      ]).exec(client);
+
+      const res2 = await new ZRangeCommand([
+        key,
+        0,
+        -1,
+        { withScores: true },
+      ]).exec(client);
+
+      expect(res2).toEqual(["one", 1, "two", 2, "three", 3]);
+    });
+  });
+});
diff --git a/pkg/commands/zadd.ts b/pkg/commands/zadd.ts
@@ -1,42 +1,32 @@
 import { Command, CommandOptions } from "./command";
 
-export type ZAddCommandOptions = (
+type NXAndXXOptions =
   | { nx: true; xx?: never }
   | { nx?: never; xx: true }
-  | { nx?: never; xx?: never }
-) & { ch?: true };
+  | { nx?: never; xx?: never };
 
-export type ZAddCommandOptionsWithIncr = ZAddCommandOptions & { incr: true };
+type LTAndGTOptions =
+  | { lt: true; gt?: never }
+  | { lt?: never; gt: true }
+  | { lt?: never; gt?: never };
 
-/**
- * This type is defiend up here because otherwise it would be automatically formatted into
- * multiple lines by Deno. As a result of that, Deno will add a comma to the end and then
- * complain about the comma being there...
- */
-type Arg2<TData> = ScoreMember<TData> | ZAddCommandOptions | ZAddCommandOptionsWithIncr;
+export type ZAddCommandOptions = NXAndXXOptions &
+  LTAndGTOptions & { ch?: true } & { incr?: true };
+
+type Arg2<TData> = ScoreMember<TData> | ZAddCommandOptions;
 export type ScoreMember<TData> = { score: number; member: TData };
 /**
  * @see https://redis.io/commands/zadd
  */
-export class ZAddCommand<TData = string> extends Command<number | null, number | null> {
-  constructor(
-    cmd: [key: string, scoreMember: ScoreMember<TData>, ...scoreMemberPairs: ScoreMember<TData>[]],
-    opts?: CommandOptions<number | null, number | null>,
-  );
-  constructor(
-    cmd: [
-      key: string,
-      opts: ZAddCommandOptions | ZAddCommandOptionsWithIncr,
-      ...scoreMemberPairs: ScoreMember<TData>[],
-    ],
-    opts?: CommandOptions<number | null, number | null>,
-  );
+export class ZAddCommand<TData = string> extends Command<
+  number | null,
+  number | null
+> {
   constructor(
     [key, arg1, ...arg2]: [string, Arg2<TData>, ...ScoreMember<TData>[]],
-    opts?: CommandOptions<number | null, number | null>,
+    opts?: CommandOptions<number | null, number | null>
   ) {
     const command: unknown[] = ["zadd", key];
-
     if ("nx" in arg1 && arg1.nx) {
       command.push("nx");
     } else if ("xx" in arg1 && arg1.xx) {
@@ -49,6 +39,12 @@ export class ZAddCommand<TData = string> extends Command<number | null, number |
       command.push("incr");
     }
 
+    if ("lt" in arg1 && arg1.lt) {
+      command.push("lt");
+    } else if ("gt" in arg1 && arg1.gt) {
+      command.push("gt");
+    }
+
     if ("score" in arg1 && "member" in arg1) {
       command.push(arg1.score, arg1.member);
     }