Added zcount command in node. (valkey-io#909)

Bit-Quill · Feb 14, 2024 · 6ea7701 · 6ea7701
1 parent 811bcc6
commit 6ea7701
Show file tree

Hide file tree

Showing 6 changed files with 150 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -11,6 +11,7 @@
 * Python, Node: When recieving LPOP/RPOP with count, convert result to Array. ([#811](https://github.com/aws/glide-for-redis/pull/811))
 * Python: Added TYPE command ([#945](https://github.com/aws/glide-for-redis/pull/945))
 * Python: Added HLEN command ([#944](https://github.com/aws/glide-for-redis/pull/944))
+* Node: Added ZCOUNT command ([#909](https://github.com/aws/glide-for-redis/pull/909))
 
 #### Features
 * Python, Node: Added support in Lua Scripts ([#775](https://github.com/aws/glide-for-redis/pull/775), [#860](https://github.com/aws/glide-for-redis/pull/860))

diff --git a/node/src/BaseClient.ts b/node/src/BaseClient.ts
@@ -12,6 +12,7 @@ import * as net from "net";
 import { Buffer, BufferWriter, Reader, Writer } from "protobufjs";
 import {
     ExpireOptions,
+    ScoreLimit,
     SetOptions,
     ZaddOptions,
     createDecr,
@@ -53,6 +54,7 @@ import {
     createUnlink,
     createZadd,
     createZcard,
+    createZcount,
     createZrem,
     createZscore,
 } from "./Commands";
@@ -1077,6 +1079,25 @@ export class BaseClient {
         return this.createWritePromise(createZscore(key, member));
     }
 
+    /** Returns the number of members in the sorted set stored at `key` with scores between `minScore` and `maxScore`.
+     * See https://redis.io/commands/zcount/ for more details.
+     *
+     * @param key - The key of the sorted set.
+     * @param minScore - The minimum score to count from. Can be positive/negative infinity, or specific score and inclusivity.
+     * @param maxScore - The maximum score to count up to. Can be positive/negative infinity, or specific score and inclusivity.
+     * @returns The number of members in the specified score range.
+     * If `key` does not exist, it is treated as an empty sorted set, and the command returns 0.
+     * If `minScore` is greater than `maxScore`, 0 is returned.
+     * If `key` holds a value that is not a sorted set, an error is returned.
+     */
+    public zcount(
+        key: string,
+        minScore: ScoreLimit,
+        maxScore: ScoreLimit
+    ): Promise<number> {
+        return this.createWritePromise(createZcount(key, minScore, maxScore));
+    }
+
     private readonly MAP_READ_FROM_STRATEGY: Record<
         ReadFrom,
         connection_request.ReadFrom

diff --git a/node/src/Commands.ts b/node/src/Commands.ts
@@ -767,3 +767,52 @@ export function createZcard(key: string): redis_request.Command {
 export function createZscore(key: string, member: string): redis_request.Command {
     return createCommand(RequestType.ZScore, [key, member]);
 }
+
+export type ScoreLimit =
+    | `positiveInfinity`
+    | `negativeInfinity`
+    | {
+          bound: number;
+          isInclusive?: boolean;
+      };
+
+const positiveInfinityArg = "+inf";
+const negativeInfinityArg = "-inf";
+const isInclusiveArg = "(";
+
+/**
+ * @internal
+ */
+export function createZcount(
+    key: string,
+    minScore: ScoreLimit,
+    maxScore: ScoreLimit
+): redis_request.Command {
+    const args = [key];
+
+    if (minScore == "positiveInfinity") {
+        args.push(positiveInfinityArg);
+    } else if (minScore == "negativeInfinity") {
+        args.push(negativeInfinityArg);
+    } else {
+        const value =
+            minScore.isInclusive == false
+                ? isInclusiveArg + minScore.bound.toString()
+                : minScore.bound.toString();
+        args.push(value);
+    }
+
+    if (maxScore == "positiveInfinity") {
+        args.push(positiveInfinityArg);
+    } else if (maxScore == "negativeInfinity") {
+        args.push(negativeInfinityArg);
+    } else {
+        const value =
+            maxScore.isInclusive == false
+                ? isInclusiveArg + maxScore.bound.toString()
+                : maxScore.bound.toString();
+        args.push(value);
+    }
+
+    return createCommand(RequestType.Zcount, args);
+}    
diff --git a/node/src/Transaction.ts b/node/src/Transaction.ts
@@ -5,6 +5,7 @@
 import {
     ExpireOptions,
     InfoOptions,
+    ScoreLimit,
     SetOptions,
     ZaddOptions,
     createClientGetName,
@@ -56,6 +57,7 @@ import {
     createUnlink,
     createZadd,
     createZcard,
+    createZcount,
     createZrem,
     createZscore,
 } from "./Commands";
@@ -833,8 +835,24 @@ export class BaseTransaction<T extends BaseTransaction<T>> {
      * If `key` does not exist, null is returned.
      * If `key` holds a value that is not a sorted set, an error is returned.
      */
-    public zscore(key: string, member: string) {
-        this.commands.push(createZscore(key, member));
+    public zscore(key: string, member: string): T {
+        return this.addAndReturn(createZscore(key, member));
+    }
+
+    /** Returns the number of members in the sorted set stored at `key` with scores between `minScore` and `maxScore`.
+     * See https://redis.io/commands/zcount/ for more details.
+     *
+     * @param key - The key of the sorted set.
+     * @param minScore - The minimum score to count from. Can be positive/negative infinity, or specific score and inclusivity.
+     * @param maxScore - The maximum score to count up to. Can be positive/negative infinity, or specific score and inclusivity.
+     *
+     * Command Response - The number of members in the specified score range.
+     * If `key` does not exist, it is treated as an empty sorted set, and the command returns 0.
+     * If `minScore` is greater than `maxScore`, 0 is returned.
+     * If `key` holds a value that is not a sorted set, an error is returned.
+     */
+    public zcount(key: string, minScore: ScoreLimit, maxScore: ScoreLimit): T {
+        return this.addAndReturn(createZcount(key, minScore, maxScore));
     }
 
     /** Executes a single command, without checking inputs. Every part of the command, including subcommands,

diff --git a/node/tests/SharedTests.ts b/node/tests/SharedTests.ts
@@ -1413,6 +1413,62 @@ export function runBaseTests<Context>(config: {
         },
         config.timeout
     );
+
+    it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])(
+        `zcount test_%p`,
+        async (protocol) => {
+            await runTest(async (client: BaseClient) => {
+                const key1 = uuidv4();
+                const key2 = uuidv4();
+                const membersScores = { one: 1, two: 2, three: 3 };
+                expect(await client.zadd(key1, membersScores)).toEqual(3);
+                expect(
+                    await client.zcount(
+                        key1,
+                        "negativeInfinity",
+                        "positiveInfinity"
+                    )
+                ).toEqual(3);
+                expect(
+                    await client.zcount(
+                        key1,
+                        { bound: 1, isInclusive: false },
+                        { bound: 3, isInclusive: false }
+                    )
+                ).toEqual(1);
+                expect(
+                    await client.zcount(
+                        key1,
+                        { bound: 1, isInclusive: false },
+                        { bound: 3 }
+                    )
+                ).toEqual(2);
+                expect(
+                    await client.zcount(key1, "negativeInfinity", {
+                        bound: 3,
+                    })
+                ).toEqual(3);
+                expect(
+                    await client.zcount(key1, "positiveInfinity", {
+                        bound: 3,
+                    })
+                ).toEqual(0);
+                expect(
+                    await client.zcount(
+                        "nonExistingKey",
+                        "negativeInfinity",
+                        "positiveInfinity"
+                    )
+                ).toEqual(0);
+
+                expect(await client.set(key2, "foo")).toEqual("OK");
+                await expect(
+                    client.zcount(key2, "negativeInfinity", "positiveInfinity")
+                ).rejects.toThrow();
+            }, protocol);
+        },
+        config.timeout
+    );
 }
 
 export function runCommonTests<Context>(config: {

diff --git a/node/tests/TestUtilities.ts b/node/tests/TestUtilities.ts
@@ -96,7 +96,8 @@ export function transactionTest(
         .zaddIncr(key8, "member2", 1)
         .zrem(key8, ["member1"])
         .zcard(key8)
-        .zscore(key8, "member2");
+        .zscore(key8, "member2")
+        .zcount(key8, { bound: 2 }, "positiveInfinity");
     return [
         "OK",
         null,
@@ -129,6 +130,7 @@ export function transactionTest(
         1,
         1,
         3.0,
+        1,
     ];
 }