From 323fac9dae980c4568cc0c8a212fcab3b049973b Mon Sep 17 00:00:00 2001
From: sora233 <sora@sora233.me>
Date: Wed, 2 Feb 2022 00:30:16 +0800
Subject: [PATCH] impl dbsize command

---
 cli/executor.go      | 10 ++++------
 cli/executor_test.go |  6 ++++++
 cli/grammar.go       | 37 ++++++++++++++++++++++++++++++++++---
 3 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/cli/executor.go b/cli/executor.go
index a99c039..24dab6f 100644
--- a/cli/executor.go
+++ b/cli/executor.go
@@ -62,13 +62,11 @@ func BuntdbExecutor(s string) {
 		return
 	case any:
 		tx, _ := db.GetCurrentTransaction()
-		if cmd == "use" {
-			err = ctx.Run(tx)
-			if err != nil {
-				fmt.Printf("ERR: %v\n", err)
-			}
-			return
+		err = ctx.Run(tx)
+		if err != nil {
+			fmt.Printf("ERR: %v\n", err)
 		}
+		return
 	case nonNil:
 		tx, rw, closeTx := db.GetCurrentOrNewTransaction()
 		if Debug {
diff --git a/cli/executor_test.go b/cli/executor_test.go
index 233d375..406c2d4 100644
--- a/cli/executor_test.go
+++ b/cli/executor_test.go
@@ -53,6 +53,7 @@ func TestBuntdbExecutor(t *testing.T) {
 		assert.Equal(t, "b", val)
 		return nil
 	})
+	BuntdbExecutor("dbsize")
 	BuntdbExecutor("set a c")
 	bd.View(func(tx *buntdb.Tx) error {
 		val, err := tx.Get("a")
@@ -95,6 +96,7 @@ func TestBuntdbExecutor(t *testing.T) {
 	BuntdbExecutor("rwbegin")
 	BuntdbExecutor("set x y")
 	BuntdbExecutor("set y x")
+	BuntdbExecutor("dbsize")
 	BuntdbExecutor("commit")
 
 	bd.View(func(tx *buntdb.Tx) error {
@@ -111,6 +113,7 @@ func TestBuntdbExecutor(t *testing.T) {
 	BuntdbExecutor("del x")
 	BuntdbExecutor("del y")
 	BuntdbExecutor("set x xxx")
+	BuntdbExecutor("dbsize")
 	BuntdbExecutor("rollback")
 	bd.View(func(tx *buntdb.Tx) error {
 		val, err := tx.Get("x")
@@ -126,6 +129,7 @@ func TestBuntdbExecutor(t *testing.T) {
 	BuntdbExecutor("del x")
 	BuntdbExecutor("del y")
 	BuntdbExecutor("shrink")
+	BuntdbExecutor("dbsize")
 	BuntdbExecutor("save testcli_save")
 	_, err = os.Lstat("testcli_save")
 	assert.Nil(t, err)
@@ -140,6 +144,7 @@ func TestBuntdbExecutor(t *testing.T) {
 		assert.Equal(t, "x", val)
 		return nil
 	})
+	BuntdbExecutor("dbsize")
 	BuntdbExecutor("shrink")
 	BuntdbExecutor("set a xy")
 	BuntdbExecutor("save testcli_save")
@@ -156,6 +161,7 @@ func TestBuntdbExecutor(t *testing.T) {
 	BuntdbExecutor("use :memory:")
 	BuntdbExecutor("exit")
 	BuntdbExecutor("")
+	BuntdbExecutor("dbsize")
 
 	os.Remove("testcli_save")
 	os.Remove("testcli")
diff --git a/cli/grammar.go b/cli/grammar.go
index 93d5fc2..48663b4 100644
--- a/cli/grammar.go
+++ b/cli/grammar.go
@@ -10,7 +10,7 @@ import (
 	"time"
 )
 
-var null = "<nil>"
+const null = "<nil>"
 
 type GetGrammar struct {
 	Key          string `arg:"" help:"the key to get"`
@@ -262,13 +262,43 @@ type HelpGrammar struct {
 }
 
 func (h *HelpGrammar) Run(ctx *kong.Context) (err error) {
-	commands := "get\nset\ndel\nttl\nrbegin (begin a readonly transaction)\nrwbegin (begin a read/write transaction)\ncommit\nrollback\nshow\nkeys\nsearch\nuse\nshrink"
+	commands := "get\n" +
+		"set\n" +
+		"del\n" +
+		"ttl\n" +
+		"rbegin (begin a readonly transaction)\n" +
+		"rwbegin (begin a read/write transaction)\n" +
+		"commit\n" +
+		"rollback\n" +
+		"show\n" +
+		"keys\n" +
+		"search\n" +
+		"use\n" +
+		"shrink\n" +
+		"size"
 	_, err = fmt.Fprintln(ctx.Stdout,
 		"Commands available:\n----------\n"+commands+"\n----------\nFor more help, try running a command with the -h switch.",
 	)
 	return
 }
 
+type DBSizeGrammar struct {
+}
+
+func (h *DBSizeGrammar) Run(ctx *kong.Context, tx *buntdb.Tx) (err error) {
+	var count int
+	err = tx.Ascend("", func(key, value string) bool {
+		count++
+		return true
+	})
+	if err != nil {
+		fmt.Fprintf(ctx.Stdout, "ERR: Ascend error %v\n", err)
+		return
+	}
+	fmt.Fprintf(ctx.Stdout, "%v\n", count)
+	return nil
+}
+
 type Grammar struct {
 	Get      GetGrammar      `cmd:"" help:"get a value from key, return the value if key exists, or <nil> if non-exists."`
 	Set      SetGrammar      `cmd:"" help:"set a key-value [ttl], return the old value, or <nil> if old value doesn't exist."`
@@ -285,7 +315,8 @@ type Grammar struct {
 	Save     SaveGrammar     `cmd:"" help:"save the db to file"`
 	Drop     DropGrammar     `cmd:"" help:"drop command"`
 	Search   SearchGrammar   `cmd:"" help:"Search for a string contained in any values"`
-	Help     HelpGrammar     `cmd:""`
+	Help     HelpGrammar     `cmd:"" help:"show available command"`
+	DBSize   DBSizeGrammar   `cmd:"" name:"dbsize" help:"show db size"`
 	Exit     bool            `kong:"-"`
 }