Implement DBSIZE command (#159)

Implement DBSIZE Command - @NicoleStrel

README.md

@@ -206,6 +206,7 @@ Benchmark script options:
 <a name="commands-generic"></a>
 ## GENERIC
 * [COPY](https://sugardb.io/docs/commands/generic/copy)
+* [DBSIZE](https://sugardb.io/docs/commands/generic/dbsize)
 * [DECR](https://sugardb.io/docs/commands/generic/decr)
 * [DECRBY](https://sugardb.io/docs/commands/generic/decrby)
 * [DEL](https://sugardb.io/docs/commands/generic/del)

docs/docs/commands/generic/dbsize.mdx

@@ -0,0 +1,47 @@
+import Tabs from '@theme/Tabs';
+import TabItem from '@theme/TabItem';
+
+# DBSIZE
+
+### Syntax
+```
+DBSIZE
+```
+
+### Module
+<span className="acl-category">generic</span>
+
+### Categories 
+<span className="acl-category">fast</span>
+<span className="acl-category">read</span>
+<span className="acl-category">keyspace</span>
+
+### Description
+Return the number of keys in the currently-selected database.
+
+### Examples
+
+<Tabs
+  defaultValue="go"
+  values={[
+    { label: 'Go (Embedded)', value: 'go', },
+    { label: 'CLI', value: 'cli', },
+  ]}
+>
+  <TabItem value="go">
+  Get the number of keys in the database:
+  ```go
+  db, err := sugardb.NewSugarDB()
+  if err != nil {
+    log.Fatal(err)
+  }
+  key, err := db.DBSize()
+  ```
+  </TabItem>
+  <TabItem value="cli">
+  Get the number of keys in the database:
+  ```
+  > DBSIZE
+  ```
+  </TabItem>
+</Tabs> 

internal/modules/generic/commands.go

@@ -714,6 +714,11 @@ func handleRandomKey(params internal.HandlerFuncParams) ([]byte, error) {
 	return []byte(fmt.Sprintf("+%v\r\n", key)), nil
 }
 
+func handleDBSize(params internal.HandlerFuncParams) ([]byte, error) {
+	count := params.DBSize(params.Context)
+	return []byte(fmt.Sprintf(":%d\r\n", count)), nil
+}
+
 func handleGetdel(params internal.HandlerFuncParams) ([]byte, error) {
 	keys, err := getDelKeyFunc(params.Command)
 	if err != nil {
@@ -1285,6 +1290,16 @@ Delete all the keys in the currently selected database. This command is always s
 			KeyExtractionFunc: randomKeyFunc,
 			HandlerFunc:       handleRandomKey,
 		},
+		{
+			Command:           "dbsize",
+			Module:            constants.GenericModule,
+			Categories:        []string{constants.KeyspaceCategory, constants.ReadCategory, constants.FastCategory},
+			Description:       "(DBSIZE) Return the number of keys in the currently selected database.",
+			Sync:              false,
+			Type:              "BUILT_IN",
+			KeyExtractionFunc: dbSizeKeyFunc,
+			HandlerFunc:       handleDBSize,
+		},
 		{
 			Command:           "getdel",
 			Module:            constants.GenericModule,

internal/modules/generic/commands_test.go

@@ -2909,6 +2909,45 @@ func Test_Generic(t *testing.T) {
 		}
 	})
 
+	t.Run("Test_HandleDBSIZE", func(t *testing.T) {
+		conn, err := internal.GetConnection("localhost", port)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		defer func() {
+			_ = conn.Close()
+		}()
+		client := resp.NewConn(conn)
+
+		// Populate the store with a few keys
+		expectedSize := 5
+		for i := 0; i < expectedSize; i++ {
+			_, _, err := mockServer.Set(
+				fmt.Sprintf("DBSizeKey%d", i),
+				fmt.Sprintf("Value%d", i),
+				sugardb.SETOptions{},
+			)
+			if err != nil {
+				t.Error(err)
+				return
+			}
+		}
+
+		if err = client.WriteArray([]resp.Value{resp.StringValue("DBSIZE")}); err != nil {
+			t.Error(err)
+		}
+
+		res, _, err := client.ReadValue()
+		if err != nil {
+			t.Error(err)
+		}
+
+		if res.Integer() != int(expectedSize) {
+			t.Errorf("expected dbsize %d, got %d", expectedSize, res.Integer())
+		}
+	})
+
 	t.Run("Test_HandleGETDEL", func(t *testing.T) {
 		t.Parallel()
 		conn, err := internal.GetConnection("localhost", port)

internal/modules/generic/key_funcs.go

@@ -211,6 +211,17 @@ func randomKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
 	}, nil
 }
 
+func dbSizeKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
+	if len(cmd) != 1 {
+		return internal.KeyExtractionFuncResult{}, errors.New(constants.WrongArgsResponse)
+	}
+	return internal.KeyExtractionFuncResult{
+		Channels:  make([]string, 0),
+		ReadKeys:  make([]string, 0),
+		WriteKeys: make([]string, 0),
+	}, nil
+}
+
 func getDelKeyFunc(cmd []string) (internal.KeyExtractionFuncResult, error) {
 	if len(cmd) != 2 {
 		return internal.KeyExtractionFuncResult{}, errors.New(constants.WrongArgsResponse)

internal/types.go

@@ -196,6 +196,8 @@ type HandlerFuncParams struct {
 	Flush func(database int)
 	// RandomKey returns a random key
 	RandomKey func(ctx context.Context) string
+	// DBSize returns the number of keys in the currently selected database.
+	DBSize func(ctx context.Context) int
 	// (TOUCH key [key ...]) Alters the last access time or access count of the key(s)
 	// depending on whether LFU or LRU strategy was used.
 	// A key is ignored if it does not exist.

sugardb/api_generic.go

@@ -632,6 +632,16 @@ func (server *SugarDB) RandomKey() (string, error) {
 	return internal.ParseStringResponse(b)
 }
 
+// DBSize returns the number of keys in the currently-selected database.
+// Returns: An integer number of keys
+func (server *SugarDB) DBSize() (int, error) {
+	b, err := server.handleCommand(server.context, internal.EncodeCommand([]string{"DBSIZE"}), nil, false, true)
+	if err != nil {
+		return 0, err
+	}
+	return internal.ParseIntegerResponse(b)
+}
+
 // GetDel retrieves the value at the provided key and deletes that key.
 //
 // Parameters:

sugardb/api_generic_test.go

@@ -1410,6 +1410,37 @@ func TestSugarDB_RANDOMKEY(t *testing.T) {
 
 }
 
+func TestSugarDB_DBSize(t *testing.T) {
+	server := createSugarDB()
+	got, err := server.DBSize()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if got != 0 {
+		t.Errorf("DBSIZE error, expected 0, got %d", got)
+	}
+
+	// test with keys
+	testkeys := []string{"1", "2", "3"}
+	for _, k := range testkeys {
+		err := presetValue(server, context.Background(), k, "")
+		if err != nil {
+			t.Error(err)
+			return
+		}
+	}
+
+	got, err = server.DBSize()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	if got != len(testkeys) {
+		t.Errorf("DBSIZE error, expected %d, got %d", len(testkeys), got)
+	}
+}
+
 func TestSugarDB_GETDEL(t *testing.T) {
 	server := createSugarDB()
 

sugardb/keyspace.go

@@ -771,6 +771,14 @@ func (server *SugarDB) randomKey(ctx context.Context) string {
 	return randkey
 }
 
+func (server *SugarDB) dbSize(ctx context.Context) int {
+	server.storeLock.RLock()
+	defer server.storeLock.RUnlock()
+
+	database := ctx.Value("Database").(int)
+	return len(server.store[database])
+}
+
 func (server *SugarDB) getObjectFreq(ctx context.Context, key string) (int, error) {
 	database := ctx.Value("Database").(int)
 

sugardb/modules.go

@@ -18,12 +18,13 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"github.com/echovault/sugardb/internal"
-	"github.com/echovault/sugardb/internal/clock"
-	"github.com/echovault/sugardb/internal/constants"
 	"io"
 	"net"
 	"strings"
+
+	"github.com/echovault/sugardb/internal"
+	"github.com/echovault/sugardb/internal/clock"
+	"github.com/echovault/sugardb/internal/constants"
 )
 
 func (server *SugarDB) getCommand(cmd string) (internal.Command, error) {
@@ -61,6 +62,7 @@ func (server *SugarDB) getHandlerFuncParams(ctx context.Context, cmd []string, c
 		GetClock:              server.getClock,
 		Flush:                 server.Flush,
 		RandomKey:             server.randomKey,
+		DBSize:                server.dbSize,
 		TouchKey:              server.updateKeysInCache,
 		GetObjectFrequency:    server.getObjectFreq,
 		GetObjectIdleTime:     server.getObjectIdleTime,