Node: Add command DBSize (valkey-io#1932)

shohamazon · Jul 18, 2024 · d63ec24 · d63ec24
1 parent ebb04c1
commit d63ec24
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -99,6 +99,7 @@
 * Python: Added XINFO STREAM command ([#1816](https://github.com/valkey-io/valkey-glide/pull/1816))
 * Python: Added transaction supports for DUMP, RESTORE, FUNCTION DUMP and FUNCTION RESTORE ([#1814](https://github.com/valkey-io/valkey-glide/pull/1814))
 * Node: Added FlushAll command ([#1958](https://github.com/valkey-io/valkey-glide/pull/1958))
+* Node: Added DBSize command ([#1932](https://github.com/valkey-io/valkey-glide/pull/1932))
 
 #### Breaking Changes
 * Node: Update XREAD to return a Map of Map ([#1494](https://github.com/valkey-io/valkey-glide/pull/1494))

diff --git a/node/src/Commands.ts b/node/src/Commands.ts
@@ -1696,3 +1696,10 @@ export function createLPos(
 
     return createCommand(RequestType.LPos, args);
 }
+
+/**
+ * @internal
+ */
+export function createDBSize(): command_request.Command {
+    return createCommand(RequestType.DBSize, []);
+}
diff --git a/node/src/GlideClient.ts b/node/src/GlideClient.ts
@@ -15,6 +15,7 @@ import {
     createConfigRewrite,
     createConfigSet,
     createCustomCommand,
+    createDBSize,
     createEcho,
     createFlushAll,
     createInfo,
@@ -355,4 +356,21 @@ export class GlideClient extends BaseClient {
             return this.createWritePromise(createFlushAll());
         }
     }
+
+    /**
+     * Returns the number of keys in the currently selected database.
+     *
+     * See https://valkey.io/commands/dbsize/ for more details.
+     *
+     * @returns The number of keys in the currently selected database.
+     *
+     * @example
+     * ```typescript
+     * const numKeys = await client.dbsize();
+     * console.log("Number of keys in the current database: ", numKeys);
+     * ```
+     */
+    public dbsize(): Promise<number> {
+        return this.createWritePromise(createDBSize());
+    }
 }
diff --git a/node/src/GlideClusterClient.ts b/node/src/GlideClusterClient.ts
@@ -15,6 +15,7 @@ import {
     createConfigRewrite,
     createConfigSet,
     createCustomCommand,
+    createDBSize,
     createEcho,
     createFlushAll,
     createInfo,
@@ -620,4 +621,24 @@ export class GlideClusterClient extends BaseClient {
             toProtobufRoute(route),
         );
     }
+
+    /**
+     * Returns the number of keys in the database.
+     *
+     * See https://valkey.io/commands/dbsize/ for more details.
+
+     * @param route - The command will be routed to all primaries, unless `route` is provided, in which
+     *     case the client will route the command to the nodes defined by `route`.
+     * @returns The number of keys in the database.
+     *     In the case of routing the query to multiple nodes, returns the aggregated number of keys across the different nodes.
+     *
+     * @example
+     * ```typescript
+     * const numKeys = await client.dbsize("allPrimaries");
+     * console.log("Number of keys across all primary nodes: ", numKeys);
+     * ```
+     */
+    public dbsize(route?: Routes): Promise<ClusterResponse<number>> {
+        return this.createWritePromise(createDBSize(), toProtobufRoute(route));
+    }
 }
diff --git a/node/src/Transaction.ts b/node/src/Transaction.ts
@@ -29,6 +29,7 @@ import {
     createConfigRewrite,
     createConfigSet,
     createCustomCommand,
+    createDBSize,
     createDecr,
     createDecrBy,
     createDel,
@@ -1717,6 +1718,17 @@ export class BaseTransaction<T extends BaseTransaction<T>> {
     public lpos(key: string, element: string, options?: LPosOptions): T {
         return this.addAndReturn(createLPos(key, element, options));
     }
+
+    /**
+     * Returns the number of keys in the currently selected database.
+     *
+     * See https://valkey.io/commands/dbsize/ for more details.
+     *
+     * Command Response - The number of keys in the currently selected database.
+     */
+    public dbsize(): T {
+        return this.addAndReturn(createDBSize());
+    }
 }
 
 /**

diff --git a/node/tests/SharedTests.ts b/node/tests/SharedTests.ts
@@ -3972,6 +3972,54 @@ export function runBaseTests<Context>(config: {
         },
         config.timeout,
     );
+
+    it.each([ProtocolVersion.RESP2, ProtocolVersion.RESP3])(
+        `dbsize test_%p`,
+        async (protocol) => {
+            await runTest(async (client: BaseClient) => {
+                // flush all data
+                expect(await client.flushall()).toBe("OK");
+
+                // check that DBSize is 0
+                expect(await client.dbsize()).toBe(0);
+
+                // set 10 random key-value pairs
+                for (let i = 0; i < 10; i++) {
+                    const key = `{key}:${uuidv4()}`;
+                    const value = "0".repeat(Math.random() * 7);
+
+                    expect(await client.set(key, value)).toBe("OK");
+                }
+
+                // check DBSIZE after setting
+                expect(await client.dbsize()).toBe(10);
+
+                // additional test for the standalone client
+                if (client instanceof GlideClient) {
+                    expect(await client.flushall()).toBe("OK");
+                    const key = uuidv4();
+                    expect(await client.set(key, "value")).toBe("OK");
+                    expect(await client.dbsize()).toBe(1);
+                    // switching to another db to check size
+                    expect(await client.select(1)).toBe("OK");
+                    expect(await client.dbsize()).toBe(0);
+                }
+
+                // additional test for the cluster client
+                if (client instanceof GlideClusterClient) {
+                    expect(await client.flushall()).toBe("OK");
+                    const key = uuidv4();
+                    expect(await client.set(key, "value")).toBe("OK");
+                    const primaryRoute: SingleNodeRoute = {
+                        type: "primarySlotKey",
+                        key: key,
+                    };
+                    expect(await client.dbsize(primaryRoute)).toBe(1);
+                }
+            }, protocol);
+        },
+        config.timeout,
+    );
 }
 
 export function runCommonTests<Context>(config: {

diff --git a/node/tests/TestUtilities.ts b/node/tests/TestUtilities.ts
@@ -317,6 +317,8 @@ export async function transactionTest(
     const args: ReturnType[] = [];
     baseTransaction.flushall();
     args.push("OK");
+    baseTransaction.dbsize();
+    args.push(0);
     baseTransaction.set(key1, "bar");
     args.push("OK");
     baseTransaction.objectEncoding(key1);