add shrink/save/drop command

Sora233 · Sora233 · commit 9e13b220ae73 · 2021-01-11T14:28:29.000+08:00
diff --git a/README.md b/README.md
@@ -37,10 +37,13 @@ go get -u -v github.com/Sora233/buntdb-cli
     * show
     * keys
     * use
+    * shrink
+    * save
 
 You can provide -h flag for command to print help message.
 ![get](https://user-images.githubusercontent.com/11474360/104104364-81e09e00-52e2-11eb-8863-391420bf6064.jpg)
 
 ### TODO
 
-- [ ] create index
+- [ ] create index (Index is memory-only, You need to create index everytime you restart, so I am considering whether to
+  impl it)
diff --git a/cli/completer.go b/cli/completer.go
@@ -14,15 +14,15 @@ func BuntdbCompleter(d prompt.Document) []prompt.Suggest {
 	args := strings.Split(d.TextBeforeCursor(), " ")
 	if len(args) == 1 {
 		// input command
-		return cmdCompleter(args[0])
+		return cmdCompleter(d, args[0])
 	} else {
-		return optionCompleter(args[0], args[1:])
+		return optionCompleter(d, args[0], args[1:])
 	}
 }
 
-func cmdCompleter(cmd string) []prompt.Suggest {
+func cmdCompleter(d prompt.Document, cmd string) []prompt.Suggest {
 	if Debug {
-		fmt.Printf("cmdCompleter %v\n", cmd)
+		fmt.Printf("|cmdCompleter %v|\n", cmd)
 	}
 	cmds := []prompt.Suggest{
 		{Text: "get", Description: "get command"},
@@ -33,12 +33,15 @@ func cmdCompleter(cmd string) []prompt.Suggest {
 		{Text: "keys", Description: "iterate keys"},
 		{Text: "use", Description: "change db"},
 		{Text: "exit", Description: "exit buntdb shell client"},
+		{Text: "drop", Description: "drop the index"},
 	}
 	tx, _ := db.GetCurrentTransaction()
 	if tx == nil {
 		cmds = append(cmds,
 			prompt.Suggest{Text: "rbegin", Description: "open a readonly transaction"},
 			prompt.Suggest{Text: "rwbegin", Description: "open a read/write transaction"},
+			prompt.Suggest{Text: "shrink", Description: "shrink command"},
+			prompt.Suggest{Text: "save", Description: "save db to file"},
 		)
 	} else {
 		cmds = append(cmds,
@@ -49,28 +52,69 @@ func cmdCompleter(cmd string) []prompt.Suggest {
 	return prompt.FilterHasPrefix(cmds, cmd, true)
 }
 
-func optionCompleter(cmd string, args []string) []prompt.Suggest {
+func optionCompleter(d prompt.Document, cmd string, args []string) []prompt.Suggest {
 	if Debug {
-		fmt.Printf("optionCompleter %v %v\n", cmd, args)
+		fmt.Printf("|optionCompleter %v [%v]|\n", cmd, strings.Join(args, ":"))
 	}
+	var result = make([]prompt.Suggest, 0)
 	switch cmd {
 	case "get":
 	case "set":
 	case "del":
 	case "ttl":
 	case "show":
-		return []prompt.Suggest{
+		result = []prompt.Suggest{
 			{Text: "index"},
 			{Text: "db"},
 		}
+		if len(args) == 0 {
+			break
+		}
+		arg := args[0]
+		if Debug {
+			fmt.Printf("|arg %v|", arg)
+		}
+		switch arg {
+		case "index":
+			result = []prompt.Suggest{}
+		case "db":
+			result = []prompt.Suggest{}
+		default:
+			result = prompt.FilterHasPrefix(result, arg, true)
+		}
 	case "keys":
 	case "use":
 	case "rbegin":
 	case "rwbegin":
 	case "rollback":
 	case "commit":
+	case "shrink":
+	case "save":
+	case "drop":
+		result = []prompt.Suggest{
+			{Text: "index"},
+		}
+		if len(args) == 0 {
+			break
+		}
+		switch args[0] {
+		case "index":
+			result = []prompt.Suggest{}
+			tx, _, closeTx := db.GetCurrentOrNewTransaction()
+			defer closeTx()
+			indexes, err := tx.Indexes()
+			if err == nil {
+				for _, index := range indexes {
+					result = append(result, prompt.Suggest{Text: index})
+				}
+				if len(args) >= 2 {
+					result = prompt.FilterHasPrefix(result, args[1], true)
+				}
+			}
+		default:
+			result = prompt.FilterHasPrefix(result, args[0], true)
+		}
 	default:
-		return []prompt.Suggest{}
 	}
-	return []prompt.Suggest{}
+	return result
 }
diff --git a/cli/completer_test.go b/cli/completer_test.go
@@ -26,10 +26,14 @@ func TestBuntdbCompleter(t *testing.T) {
 	d = buf.Document()
 	assert.NotNil(t, d)
 	sug = BuntdbCompleter(*d)
-	assert.Len(t, sug, 2)
-	text := []string{sug[0].Text, sug[1].Text}
+	var text []string
+	for _, s := range sug {
+		text = append(text, s.Text)
+	}
 	assert.Contains(t, text, "set")
 	assert.Contains(t, text, "show")
+	assert.Contains(t, text, "shrink")
+	assert.Contains(t, text, "save")
 
 	buf.DeleteBeforeCursor(999)
 	buf.InsertText("get ", false, true)
@@ -58,6 +62,7 @@ func TestBuntdbCompleter(t *testing.T) {
 	buf.DeleteBeforeCursor(999)
 
 	db.InitBuntDB(":memory:")
+	defer db.Close()
 	db.Begin(true)
 	buf.InsertText("r", false, true)
 	d = buf.Document()
@@ -66,6 +71,24 @@ func TestBuntdbCompleter(t *testing.T) {
 	assert.Len(t, sug, 1)
 	assert.Equal(t, "rollback", sug[0].Text)
 	db.Rollback()
-	db.Close()
 
+	buf.DeleteBeforeCursor(999)
+	buf.InsertText("dr", false, true)
+	d = buf.Document()
+	assert.NotNil(t, d)
+	sug = BuntdbCompleter(*d)
+	assert.Len(t, sug, 1)
+	assert.Equal(t, "drop", sug[0].Text)
+
+	buf.InsertText("op ", false, true)
+	d = buf.Document()
+	assert.NotNil(t, d)
+	sug = BuntdbCompleter(*d)
+	assert.Len(t, sug, 1)
+	assert.Equal(t, "index", sug[0].Text)
+	buf.InsertText("index", false, true)
+	d = buf.Document()
+	assert.NotNil(t, d)
+	sug = BuntdbCompleter(*d)
+	assert.Len(t, sug, 0)
 }
diff --git a/cli/executor.go b/cli/executor.go
@@ -7,6 +7,25 @@ import (
 	"strings"
 )
 
+type transactionRequireType int64
+
+const (
+	noNeed transactionRequireType = iota
+	any
+	nonNil
+)
+
+func commandTransactionRequireType(command string) transactionRequireType {
+	switch command {
+	case "rwbegin", "rbegin", "rollback", "commit", "shrink", "save":
+		return noNeed
+	case "use":
+		return any
+	default:
+		return nonNil
+	}
+}
+
 func BuntdbExecutor(s string) {
 	s = strings.TrimSpace(s)
 	if s == "" || s == "exit" {
@@ -32,52 +51,40 @@ func BuntdbExecutor(s string) {
 		return
 	}
 	cmd := ctx.Selected().Name
-	if cmd == "rwbegin" || cmd == "rbegin" || cmd == "rollback" || cmd == "commit" {
+	switch commandTransactionRequireType(cmd) {
+	case noNeed:
 		err = ctx.Run()
 		if err != nil {
 			fmt.Printf("ERR: %v\n", err)
 		}
 		return
-	}
-	tx, rw := db.GetCurrentTransaction()
-	if ctx.Selected().Name == "use" {
-		err = ctx.Run(tx)
-		if err != nil {
-			fmt.Printf("ERR: %v\n", err)
-		}
-		return
-	}
-	if tx != nil {
-		if Debug {
-			fmt.Printf("got current %v transaction\n", db.RWDescribe(rw))
-		}
-		err = ctx.Run(tx)
-		if err != nil {
-			fmt.Printf("ERR: %v\n", err)
+	case any:
+		tx, _ := db.GetCurrentTransaction()
+		if cmd == "use" {
+			err = ctx.Run(tx)
+			if err != nil {
+				fmt.Printf("ERR: %v\n", err)
+			}
 			return
 		}
-	} else {
+	case nonNil:
+		tx, rw, closeTx := db.GetCurrentOrNewTransaction()
 		if Debug {
-			fmt.Printf("no transaction, create a rw transaction\n")
-		}
-		tx, err := db.Begin(true)
-		if err != nil {
-			fmt.Printf("ERR: %v\n", err)
-			return
+			fmt.Printf("GetCurrentOrNewTransaction(%v)\n", db.RWDescribe(rw))
 		}
 		defer func() {
-			if Debug {
-				fmt.Printf("transaction commit\n")
-			}
-			err := db.Commit()
+			err = closeTx()
 			if err != nil {
-				fmt.Printf("ERR: commit error %v\n", err)
+				fmt.Printf("ERR: %v\n", err)
 			}
 		}()
 		err = ctx.Run(tx)
 		if err != nil {
 			fmt.Printf("ERR: %v\n", err)
 			return
 		}
+	default:
+		fmt.Printf("ERR: unknown transaction require\n")
+		return
 	}
 }
diff --git a/cli/executor_test.go b/cli/executor_test.go
@@ -97,6 +97,10 @@ func TestBuntdbExecutor(t *testing.T) {
 	BuntdbExecutor("rbegin")
 	BuntdbExecutor("del x")
 	BuntdbExecutor("del y")
+	BuntdbExecutor("shrink")
+	BuntdbExecutor("save testcli_save")
+	_, err = os.Lstat("testcli_save")
+	assert.Nil(t, err)
 	BuntdbExecutor("commit")
 	BuntdbExecutor("rollback")
 	bd.View(func(tx *buntdb.Tx) error {
@@ -108,15 +112,24 @@ func TestBuntdbExecutor(t *testing.T) {
 		assert.Equal(t, "x", val)
 		return nil
 	})
+	BuntdbExecutor("shrink")
+	BuntdbExecutor("set a xy")
+	BuntdbExecutor("save testcli_save")
+	_, err = os.Lstat("testcli_save")
+	assert.Nil(t, err)
+	BuntdbExecutor("save testcli_save")
+	BuntdbExecutor("save --force testcli_save")
 
 	BuntdbExecutor("use testcli-2")
 	BuntdbExecutor("use -c testcli-2")
 	assert.Equal(t, db.GetDbPath(), "testcli-2")
 	BuntdbExecutor("use -c testcli")
 	assert.Equal(t, db.GetDbPath(), "testcli")
+	BuntdbExecutor("use :memory:")
 	BuntdbExecutor("exit")
 	BuntdbExecutor("")
 
+	os.Remove("testcli_save")
 	os.Remove("testcli")
 	os.Remove("testcli-2")
 }
diff --git a/cli/grammar.go b/cli/grammar.go
@@ -123,13 +123,11 @@ func (u *UseGrammar) Run(ctx *kong.Context, tx *buntdb.Tx) error {
 		if u.Create {
 			return db.InitBuntDB(u.Path)
 		} else {
-			fmt.Fprintf(ctx.Stdout, "%v does not exist, set --create to create it.\n", u.Path)
-			return nil
+			return fmt.Errorf("%v does not exist, set --create to create it", u.Path)
 		}
 	}
 	if f.IsDir() {
-		fmt.Fprintf(ctx.Stdout, "%v is a dir.\n", u.Path)
-		return nil
+		return fmt.Errorf("%v is a dir", u.Path)
 	}
 	return db.InitBuntDB(u.Path)
 }
@@ -146,10 +144,9 @@ func (t *TTLGrammar) Run(ctx *kong.Context, tx *buntdb.Tx) error {
 			return nil
 		}
 		return err
-	} else {
-		fmt.Fprintln(ctx.Stdout, int64(ttl.Seconds()))
-		return nil
 	}
+	fmt.Fprintln(ctx.Stdout, int64(ttl.Seconds()))
+	return nil
 }
 
 type RWBeginGrammar struct{}
@@ -178,6 +175,46 @@ func (r *RollbackGrammar) Run(ctx *kong.Context) error {
 	return db.Rollback()
 }
 
+type ShrinkGrammar struct{}
+
+func (s *ShrinkGrammar) Run(ctx *kong.Context) error {
+	return db.Shrink()
+}
+
+type SaveGrammar struct {
+	Path  string `arg:"" help:"the path to save"`
+	Force bool   `optional:"" help:"overwrite if the path exists"`
+}
+
+func (s *SaveGrammar) Run(ctx *kong.Context) error {
+	f, err := os.Lstat(s.Path)
+	if err == nil {
+		if f.IsDir() {
+			return fmt.Errorf("%v is a dir", s.Path)
+		}
+		if !s.Force {
+			return fmt.Errorf("%v exist, use --force to overwrite it", s.Path)
+		}
+	}
+	file, err := os.Create(s.Path)
+	if err != nil {
+		return err
+	}
+	return db.Save(file)
+}
+
+type DropGrammar struct {
+	Index DropIndexGrammar `cmd:"" help:"drop the index with the given name"`
+}
+
+type DropIndexGrammar struct {
+	Name string `arg:"" help:"the index name to drop"`
+}
+
+func (s *DropIndexGrammar) Run(ctx *kong.Context, tx *buntdb.Tx) error {
+	return tx.DropIndex(s.Name)
+}
+
 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."`
@@ -190,6 +227,9 @@ type Grammar struct {
 	RBegin   RBeginGrammar   `cmd:"" name:"rbegin" help:"begin a readonly transaction"`
 	Commit   CommitGrammar   `cmd:"" help:"commit a transaction"`
 	Rollback RollbackGrammar `cmd:"" help:"rollback a transaction"`
+	Shrink   ShrinkGrammar   `cmd:"" help:"run database shrink command"`
+	Save     SaveGrammar     `cmd:"" help:"save the db to file"`
+	Drop     DropGrammar     `cmd:"" help:"drop command"`
 	Exit     bool            `kong:"-"`
 }
 
diff --git a/db/buntdb.go b/db/buntdb.go
diff --git a/db/buntdb_test.go b/db/buntdb_test.go
