database.go 4.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227
  1. package database
  2. /*
  3. ArOZ Online Database Access Module
  4. author: tobychui
  5. This is an improved Object oriented base solution to the original
  6. aroz online database script.
  7. */
  8. import (
  9. "encoding/json"
  10. "errors"
  11. "log"
  12. "sync"
  13. "github.com/boltdb/bolt"
  14. )
  15. type Database struct {
  16. Db *bolt.DB
  17. Tables sync.Map
  18. ReadOnly bool
  19. }
  20. func NewDatabase(dbfile string, readOnlyMode bool) (*Database, error) {
  21. db, err := bolt.Open(dbfile, 0600, nil)
  22. log.Println("Key-value Database Service Started: " + dbfile)
  23. return &Database{
  24. Db: db,
  25. Tables: sync.Map{},
  26. ReadOnly: readOnlyMode,
  27. }, err
  28. }
  29. /*
  30. Create / Drop a table
  31. Usage:
  32. err := sysdb.NewTable("MyTable")
  33. err := sysdb.DropTable("MyTable")
  34. */
  35. func (d *Database) UpdateReadWriteMode(readOnly bool) {
  36. d.ReadOnly = readOnly
  37. }
  38. //Dump the whole db into a log file
  39. func (d *Database) Dump(filename string) ([]string, error) {
  40. results := []string{}
  41. d.Tables.Range(func(tableName, v interface{}) bool {
  42. entries, err := d.ListTable(tableName.(string))
  43. if err != nil {
  44. log.Println("Reading table " + tableName.(string) + " failed: " + err.Error())
  45. return false
  46. }
  47. for _, keypairs := range entries {
  48. results = append(results, string(keypairs[0])+":"+string(keypairs[1])+"\n")
  49. }
  50. return true
  51. })
  52. return results, nil
  53. }
  54. //Create a new table
  55. func (d *Database) NewTable(tableName string) error {
  56. if d.ReadOnly == true {
  57. return errors.New("Operation rejected in ReadOnly mode")
  58. }
  59. err := d.Db.Update(func(tx *bolt.Tx) error {
  60. _, err := tx.CreateBucketIfNotExists([]byte(tableName))
  61. if err != nil {
  62. return err
  63. }
  64. return nil
  65. })
  66. d.Tables.Store(tableName, "")
  67. return err
  68. }
  69. //Check is table exists
  70. func (d *Database) TableExists(tableName string) bool {
  71. if _, ok := d.Tables.Load(tableName); ok {
  72. return true
  73. }
  74. return false
  75. }
  76. //Drop the given table
  77. func (d *Database) DropTable(tableName string) error {
  78. if d.ReadOnly == true {
  79. return errors.New("Operation rejected in ReadOnly mode")
  80. }
  81. err := d.Db.Update(func(tx *bolt.Tx) error {
  82. err := tx.DeleteBucket([]byte(tableName))
  83. if err != nil {
  84. return err
  85. }
  86. return nil
  87. })
  88. return err
  89. }
  90. /*
  91. Write to database with given tablename and key. Example Usage:
  92. type demo struct{
  93. content string
  94. }
  95. thisDemo := demo{
  96. content: "Hello World",
  97. }
  98. err := sysdb.Write("MyTable", "username/message",thisDemo);
  99. */
  100. func (d *Database) Write(tableName string, key string, value interface{}) error {
  101. if d.ReadOnly == true {
  102. return errors.New("Operation rejected in ReadOnly mode")
  103. }
  104. jsonString, err := json.Marshal(value)
  105. if err != nil {
  106. return err
  107. }
  108. err = d.Db.Update(func(tx *bolt.Tx) error {
  109. _, err := tx.CreateBucketIfNotExists([]byte(tableName))
  110. b := tx.Bucket([]byte(tableName))
  111. err = b.Put([]byte(key), jsonString)
  112. return err
  113. })
  114. return err
  115. }
  116. /*
  117. Read from database and assign the content to a given datatype. Example Usage:
  118. type demo struct{
  119. content string
  120. }
  121. thisDemo := new(demo)
  122. err := sysdb.Read("MyTable", "username/message",&thisDemo);
  123. */
  124. func (d *Database) Read(tableName string, key string, assignee interface{}) error {
  125. err := d.Db.View(func(tx *bolt.Tx) error {
  126. b := tx.Bucket([]byte(tableName))
  127. v := b.Get([]byte(key))
  128. json.Unmarshal(v, &assignee)
  129. return nil
  130. })
  131. return err
  132. }
  133. func (d *Database) KeyExists(tableName string, key string) bool {
  134. resultIsNil := false
  135. err := d.Db.View(func(tx *bolt.Tx) error {
  136. b := tx.Bucket([]byte(tableName))
  137. v := b.Get([]byte(key))
  138. if v == nil {
  139. resultIsNil = true
  140. }
  141. return nil
  142. })
  143. if err != nil {
  144. return false
  145. } else {
  146. if resultIsNil {
  147. return false
  148. } else {
  149. return true
  150. }
  151. }
  152. }
  153. /*
  154. Delete a value from the database table given tablename and key
  155. err := sysdb.Delete("MyTable", "username/message");
  156. */
  157. func (d *Database) Delete(tableName string, key string) error {
  158. if d.ReadOnly == true {
  159. return errors.New("Operation rejected in ReadOnly mode")
  160. }
  161. err := d.Db.Update(func(tx *bolt.Tx) error {
  162. tx.Bucket([]byte(tableName)).Delete([]byte(key))
  163. return nil
  164. })
  165. return err
  166. }
  167. /*
  168. //List table example usage
  169. //Assume the value is stored as a struct named "groupstruct"
  170. entries := sysdb.ListTable("test")
  171. for _, keypairs := range entries{
  172. log.Println(string(keypairs[0]))
  173. group := new(groupstruct)
  174. json.Unmarshal(keypairs[1], &group)
  175. log.Println(group);
  176. }
  177. */
  178. func (d *Database) ListTable(tableName string) ([][][]byte, error) {
  179. var results [][][]byte
  180. err := d.Db.View(func(tx *bolt.Tx) error {
  181. b := tx.Bucket([]byte(tableName))
  182. c := b.Cursor()
  183. for k, v := c.First(); k != nil; k, v = c.Next() {
  184. results = append(results, [][]byte{k, v})
  185. }
  186. return nil
  187. })
  188. return results, err
  189. }
  190. func (d *Database) Close() {
  191. d.Db.Close()
  192. return
  193. }