dbwrapper.go 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. package main
  2. import (
  3. "encoding/json"
  4. "errors"
  5. "os"
  6. "os/user"
  7. "strings"
  8. "time"
  9. "github.com/boltdb/bolt"
  10. )
  11. // newBoltDB ...
  12. func (dbw DBWrapper) connect() (db *bolt.DB) {
  13. usr, err := user.Current()
  14. handleError(err)
  15. dbpath := usr.HomeDir + "/.rain"
  16. err = os.MkdirAll(dbpath, 0755)
  17. handleError(err)
  18. fullpath := dbpath + "/bolt.db"
  19. db, err = bolt.Open(fullpath, 0600, &bolt.Options{Timeout: 1 * time.Second})
  20. handleError(err)
  21. err = db.Update(func(tx *bolt.Tx) error {
  22. _, err := tx.CreateBucketIfNotExists([]byte("servers"))
  23. handleError(err)
  24. return err
  25. })
  26. handleError(err)
  27. return db
  28. }
  29. // DBWrapper ...
  30. type DBWrapper struct {
  31. }
  32. // DeleteServer ...
  33. func (dbw DBWrapper) DeleteServer(alias string) (err error) {
  34. bdb := dbw.connect()
  35. defer bdb.Close()
  36. err = bdb.Update(func(tx *bolt.Tx) error {
  37. b := tx.Bucket([]byte("servers"))
  38. //TODO: check if key exists first
  39. err := b.Delete([]byte(alias))
  40. return err
  41. })
  42. handleError(err)
  43. return
  44. }
  45. // AddServer ...
  46. func (dbw DBWrapper) AddServer(s Server) (err error) {
  47. bdb := dbw.connect()
  48. defer bdb.Close()
  49. err = bdb.Update(func(tx *bolt.Tx) error {
  50. b := tx.Bucket([]byte("servers"))
  51. handleError(err)
  52. encoded, err := json.Marshal(s)
  53. handleError(err)
  54. err = b.Put([]byte(s.Alias), encoded)
  55. return err
  56. })
  57. handleError(err)
  58. return
  59. }
  60. // ServerSearch ...
  61. func (dbw DBWrapper) ServerSearch(search string) (servers []Server, err error) {
  62. bdb := dbw.connect()
  63. defer bdb.Close()
  64. err = bdb.View(func(tx *bolt.Tx) error {
  65. b := tx.Bucket([]byte("servers"))
  66. cursor := b.Cursor()
  67. for alias, host := cursor.First(); alias != nil; alias, host = cursor.Next() {
  68. var s Server
  69. err := json.Unmarshal(host, &s)
  70. handleError(err)
  71. searchL := strings.ToLower(search)
  72. aliasL := strings.ToLower(string(alias))
  73. hostL := strings.ToLower(s.Hostname)
  74. notesL := strings.ToLower(s.Notes)
  75. if strings.Contains(aliasL, searchL) || strings.Contains(hostL, searchL) || strings.Contains(notesL, searchL) {
  76. servers = append(servers, s)
  77. }
  78. }
  79. return nil
  80. })
  81. return
  82. }
  83. // AllServers ... TODO merge with search?
  84. func (dbw DBWrapper) AllServers() (servers []Server, err error) {
  85. bdb := dbw.connect()
  86. defer bdb.Close()
  87. err = bdb.View(func(tx *bolt.Tx) error {
  88. b := tx.Bucket([]byte("servers"))
  89. c := b.Cursor()
  90. for k, v := c.First(); k != nil; k, v = c.Next() {
  91. var s Server
  92. err := json.Unmarshal(v, &s)
  93. handleError(err)
  94. servers = append(servers, s)
  95. }
  96. return nil
  97. })
  98. handleError(err)
  99. return
  100. }
  101. // GetServer ...
  102. func (dbw DBWrapper) GetServer(alias string) (s Server, err error) {
  103. bdb := dbw.connect()
  104. defer bdb.Close()
  105. err = bdb.View(func(tx *bolt.Tx) error {
  106. b := tx.Bucket([]byte("servers"))
  107. encoded := b.Get([]byte(alias))
  108. if encoded == nil {
  109. return errors.New("alias " + alias + " not found.")
  110. }
  111. err := json.Unmarshal(encoded, &s)
  112. handleError(err)
  113. return nil
  114. })
  115. return
  116. }
  117. // UpdateServer ...
  118. func (dbw DBWrapper) UpdateServer(s Server) (err error) {
  119. return dbw.AddServer(s)
  120. }