Browse Source

Initial Commit

Signed-off-by: jolheiser <john.olheiser@gmail.com>
main
jolheiser 1 year ago
commit
7c565c4b5c
No known key found for this signature in database GPG Key ID: 83E486E71AFEB820
  1. 2
      .gitignore
  2. 5
      go.mod
  3. 4
      go.sum
  4. 13
      options.go
  5. 18
      serial.go
  6. 26
      shock.go
  7. 35
      shock_test.go
  8. 96
      store.go
  9. 98
      store_test.go

2
.gitignore

@ -0,0 +1,2 @@
# GoLand
.idea/

5
go.mod

@ -0,0 +1,5 @@
module go.jolheiser.com/shock
go 1.14
require go.etcd.io/bbolt v1.3.5

4
go.sum

@ -0,0 +1,4 @@
go.etcd.io/bbolt v1.3.5 h1:XAzx9gjCb0Rxj7EoqcClPD1d5ZBxZJk0jbuoPHenBt0=
go.etcd.io/bbolt v1.3.5/go.mod h1:G5EMThwa9y8QZGBClrRx5EY+Yw9kAhnjy3bSjsnlVTQ=
golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5 h1:LfCXLvNmTYH9kEmVgqbnsWfruoXZIrh4YBgqVHtDvw0=
golang.org/x/sys v0.0.0-20200202164722-d101bd2416d5/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=

13
options.go

@ -0,0 +1,13 @@
package shock
import "go.etcd.io/bbolt"
type Options struct {
Bolt *bbolt.Options
Serializer Serializer
}
var DefaultOptions = &Options{
Bolt: bbolt.DefaultOptions,
Serializer: JSONSerializer{},
}

18
serial.go

@ -0,0 +1,18 @@
package shock
import "encoding/json"
type Serializer interface {
Marshal(v interface{}) ([]byte, error)
Unmarshal(data []byte, v interface{}) error
}
type JSONSerializer struct{}
func (j JSONSerializer) Marshal(v interface{}) ([]byte, error) {
return json.Marshal(v)
}
func (j JSONSerializer) Unmarshal(data []byte, v interface{}) error {
return json.Unmarshal(data, &v)
}

26
shock.go

@ -0,0 +1,26 @@
package shock
import (
"os"
"go.etcd.io/bbolt"
)
type DB struct {
Bolt *bbolt.DB
Serializer Serializer
}
func Open(path string, mode os.FileMode, options *Options) (*DB, error) {
if options == nil {
options = DefaultOptions
}
db, err := bbolt.Open(path, mode, options.Bolt)
if err != nil {
return nil, err
}
return &DB{
Bolt: db,
Serializer: options.Serializer,
}, nil
}

35
shock_test.go

@ -0,0 +1,35 @@
package shock
import (
"fmt"
"io/ioutil"
"os"
"path"
"testing"
)
var db *DB
func TestMain(m *testing.M) {
dir, err := ioutil.TempDir(os.TempDir(), "shock")
if err != nil {
panic(err)
}
dbPath := path.Join(dir, "shock.db")
db, err = Open(dbPath, os.ModePerm, DefaultOptions)
if err != nil {
panic(err)
}
exit := m.Run()
if err := db.Bolt.Close(); err != nil {
fmt.Printf("Could not close DB %s: %v\n", dbPath, err)
}
if err := os.RemoveAll(dir); err != nil {
fmt.Printf("Could not delete temp dir %s: %v\n", dir, err)
}
os.Exit(exit)
}

96
store.go

@ -0,0 +1,96 @@
package shock
import (
"fmt"
"go.etcd.io/bbolt"
)
type Sequencer interface {
AssignSeq(uint64)
}
func (d *DB) Put(bucket string, val Sequencer) error {
if err := d.initBucket(bucket); err != nil {
return err
}
return d.Bolt.Update(func(tx *bbolt.Tx) error {
b := tx.Bucket([]byte(bucket))
seq, err := b.NextSequence()
if err != nil {
return err
}
val.AssignSeq(seq)
serial, err := d.Serializer.Marshal(val)
if err != nil {
return err
}
return b.Put([]byte(fmt.Sprintf("%d", seq)), serial)
})
}
func (d *DB) Get(bucket string, seq uint64, val interface{}) error {
if err := d.initBucket(bucket); err != nil {
return err
}
if err := d.Bolt.View(func(tx *bbolt.Tx) error {
serial := tx.Bucket([]byte(bucket)).Get([]byte(fmt.Sprintf("%d", seq)))
return d.Serializer.Unmarshal(serial, val)
}); err != nil {
return err
}
return nil
}
func (d *DB) Count(bucket string) (int, error) {
count := 0
if err := d.initBucket(bucket); err != nil {
return count, err
}
if err := d.Bolt.View(func(tx *bbolt.Tx) error {
return tx.Bucket([]byte(bucket)).ForEach(func(_, _ []byte) error {
count++
return nil
})
}); err != nil {
return count, err
}
return count, nil
}
type eachFn func(seq, serial []byte) error
func (d *DB) ViewEach(bucket string, fn eachFn) error {
return d.forEach(bucket, false, fn)
}
func (d *DB) UpdateEach(bucket string, fn eachFn) error {
return d.forEach(bucket, true, fn)
}
func (d *DB) forEach(bucket string, writable bool, fn eachFn) error {
if err := d.initBucket(bucket); err != nil {
return err
}
tx, err := d.Bolt.Begin(writable)
if err != nil {
return err
}
if err := tx.Bucket([]byte(bucket)).ForEach(fn); err != nil {
return err
}
if writable {
return tx.Commit()
}
return tx.Rollback()
}
func (d *DB) initBucket(bucket string) error {
return d.Bolt.Update(func(tx *bbolt.Tx) error {
_, err := tx.CreateBucketIfNotExists([]byte(bucket))
return err
})
}

98
store_test.go

@ -0,0 +1,98 @@
package shock
import (
"testing"
)
type Test struct {
ID uint64 `json:"id"`
Name string `json:"name"`
Age int `json:"age"`
Admin bool `json:"admin"`
}
func (t *Test) AssignSeq(seq uint64) {
t.ID = seq
}
func (t Test) Equal(tt Test) bool {
return t.ID == tt.ID && t.Name == tt.Name &&
t.Age == tt.Age && t.Admin == tt.Admin
}
func TestStore(t *testing.T) {
tt := []*Test{
{
Name: "user1",
Age: 25,
Admin: true,
},
{
Name: "user2",
Age: 30,
Admin: false,
},
{
Name: "user3",
Age: 40,
Admin: false,
},
}
for _, tc := range tt {
t.Run(tc.Name, func(t *testing.T) {
if err := db.Put("test", tc); err != nil {
t.Log(err)
t.FailNow()
}
var tcc Test
err := db.Get("test", tc.ID, &tcc)
if err != nil {
t.Log(err)
t.FailNow()
}
if !tcc.Equal(*tc) {
t.Log("Serialized struct is not the same")
t.FailNow()
}
})
}
t.Run("count", func(t *testing.T) {
count, err := db.Count("test")
if err != nil {
t.Log(err)
t.FailNow()
}
if count != len(tt) {
t.Log("Count is off")
t.FailNow()
}
})
t.Run("list", func(t *testing.T) {
list := make([]Test, 0, len(tt))
if err := db.ViewEach("test", func(_, serial []byte) error {
var t Test
if err := db.Serializer.Unmarshal(serial, &t); err != nil {
return err
}
list = append(list, t)
return nil
}); err != nil {
t.Log(err)
t.FailNow()
}
for idx, l := range list {
if !tt[idx].Equal(l) {
t.Logf("List doesn't match:\n%v\n%v\n", tt[idx], l)
t.FailNow()
}
}
})
}
Loading…
Cancel
Save