diff --git a/README.md b/README.md index 975f78c..13d4295 100644 --- 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 \ No newline at end of file +- [ ] create index (Index is memory-only, You need to create index everytime you restart, so I am considering whether to + impl it) \ No newline at end of file diff --git a/cli/completer.go b/cli/completer.go index e27c740..1b17342 100644 --- 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 index c0b46fd..1ce1f14 100644 --- 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 index 2234de0..21db9d3 100644 --- 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,46 +51,31 @@ 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) @@ -79,5 +83,8 @@ func BuntdbExecutor(s string) { 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 index 4cedc0d..db468a3 100644 --- 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 index 029c69f..fe41ead 100644 --- 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 if non-exists."` Set SetGrammar `cmd:"" help:"set a key-value [ttl], return the old value, or 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 index 4551eba..f66533b 100644 --- a/db/buntdb.go +++ b/db/buntdb.go @@ -3,6 +3,7 @@ package db import ( "errors" "github.com/tidwall/buntdb" + "io" "io/ioutil" "os" "path" @@ -123,10 +124,35 @@ func Rollback() error { } } +func Shrink() error { + if tx != nil { + return ErrTransactionExist + } + return db.Shrink() +} + +func Save(writer io.Writer) error { + if tx != nil { + return ErrTransactionExist + } + return db.Save(writer) +} + +// GetCurrentTransaction return current transaction, return nil if no current transaction func GetCurrentTransaction() (*buntdb.Tx, bool) { return tx, writable } +// GetCurrentOrNewTransaction return current transaction, or create new one if no transaction, +// make sure the func is called. +func GetCurrentOrNewTransaction() (*buntdb.Tx, bool, func() error) { + if tx != nil { + return tx, writable, func() error { return nil } + } + Begin(true) + return tx, writable, func() error { return Commit() } +} + func RWDescribe(writable bool) string { if writable { return "rw" diff --git a/db/buntdb_test.go b/db/buntdb_test.go index cd96846..4a196be 100644 --- a/db/buntdb_test.go +++ b/db/buntdb_test.go @@ -3,6 +3,7 @@ package db import ( "github.com/stretchr/testify/assert" "github.com/tidwall/buntdb" + "io/ioutil" "os" "path/filepath" "strings" @@ -45,6 +46,7 @@ func TestSwitchBuntDB(t *testing.T) { assert.Nil(t, err) assert.Nil(t, InitBuntDB(testDb2)) db, err = GetClient() + assert.Nil(t, err) err = db.View(func(tx *buntdb.Tx) error { _, err := tx.Get("a") assert.Equal(t, buntdb.ErrNotFound, err) @@ -92,6 +94,7 @@ func TestBegin(t *testing.T) { }) tx, err = Begin(false) + assert.Nil(t, err) _, _, err = tx.Set("a", "a", nil) assert.Equal(t, buntdb.ErrTxNotWritable, err) assert.NotNil(t, Commit()) @@ -124,9 +127,40 @@ func TestGetCurrentTransaction(t *testing.T) { tx, rw = GetCurrentTransaction() assert.False(t, rw) assert.NotNil(t, tx) + Shrink() Rollback() tx, rw = GetCurrentTransaction() assert.Nil(t, tx) + Shrink() + os.Remove(testDb) + os.Remove(testDb2) +} + +func TestSave(t *testing.T) { + assert.Nil(t, InitBuntDB(testDb)) + db, err := GetClient() + assert.Nil(t, err) + err = db.Update(func(tx *buntdb.Tx) error { + _, _, err := tx.Set("a", "testsave", nil) + return err + }) + assert.Nil(t, err) + f, err := ioutil.TempFile("", "test_save") + assert.Nil(t, err) + Save(f) + Close() + f.Close() + + db, err = buntdb.Open(f.Name()) + assert.Nil(t, err) + db.View(func(tx *buntdb.Tx) error { + val, err := tx.Get("a") + assert.Nil(t, err) + assert.Equal(t, "testsave", val) + return nil + }) + db.Close() + os.Remove(f.Name()) os.Remove(testDb) os.Remove(testDb2) }