Create a database abstraction

README.md

@@ -79,6 +79,14 @@ create /etc/eremetic/eremetic.yml with:
     loglevel: DEBUG
     logformat: json
 
+## Database Backing
+Eremetic uses a BoltDB based database to store task information. The location of
+it can be set by adding
+
+    database: '/tmp/eremeticdb/eremetic.db'
+
+to the eremetic.yml
+
 ### Authentication
 To enable mesos framework authentication add the location of credential file to your configuration:
 

database/task.go → database/boltdb.go

@@ -2,19 +2,66 @@ package database
 
 import (
 	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"path/filepath"
 
 	"github.com/boltdb/bolt"
 	"github.com/klarna/eremetic/types"
 )
 
-// PutTask stores a requested task in the database
-func PutTask(task *types.EremeticTask) error {
-	err := ensureDB()
-	if err != nil {
-		return err
+type boltDriver struct {
+	database *bolt.DB
+}
+
+func boltDB(file string) (TaskDB, error) {
+	if file == "" {
+		return nil, errors.New("Missing BoltDB database loctation.")
+	}
+
+	if !filepath.IsAbs(file) {
+		dir, _ := os.Getwd()
+		file = fmt.Sprintf("%s/../%s", dir, file)
 	}
+	os.MkdirAll(filepath.Dir(file), 0755)
+
+	db, err := bolt.Open(file, 0600, nil)
+
+	wrapped := wrap(db)
+
+	return wrapped, err
+}
 
-	return boltdb.Update(func(tx *bolt.Tx) error {
+func wrap(db *bolt.DB) TaskDB {
+	return boltDriver{
+		database: db,
+	}
+}
+
+// Close is used to Close the database
+func (db boltDriver) Close() {
+	if db.database != nil {
+		db.database.Close()
+	}
+}
+
+// Clean is used to delete the tasks bucket
+func (db boltDriver) Clean() error {
+	return db.database.Update(func(tx *bolt.Tx) error {
+		if err := tx.DeleteBucket([]byte("tasks")); err != nil {
+			return err
+		}
+		if _, err := tx.CreateBucketIfNotExists([]byte("tasks")); err != nil {
+			return err
+		}
+		return nil
+	})
+}
+
+// PutTask stores a requested task in the database
+func (db boltDriver) PutTask(task *types.EremeticTask) error {
+	return db.database.Update(func(tx *bolt.Tx) error {
 		b, err := tx.CreateBucketIfNotExists([]byte("tasks"))
 		if err != nil {
 			return err
@@ -31,8 +78,8 @@ func PutTask(task *types.EremeticTask) error {
 
 // ReadTask fetches a task from the database and applies a mask to the
 // MaskedEnvironment field
-func ReadTask(id string) (types.EremeticTask, error) {
-	task, err := ReadUnmaskedTask(id)
+func (db boltDriver) ReadTask(id string) (types.EremeticTask, error) {
+	task, err := db.ReadUnmaskedTask(id)
 
 	for k := range task.MaskedEnvironment {
 		task.MaskedEnvironment[k] = "*******"
@@ -46,15 +93,10 @@ func ReadTask(id string) (types.EremeticTask, error) {
 // This function should be considered internal to Eremetic, and is used where
 // we need to fetch a task and then re-save it to the database. It should not
 // be returned to the API.
-func ReadUnmaskedTask(id string) (types.EremeticTask, error) {
+func (db boltDriver) ReadUnmaskedTask(id string) (types.EremeticTask, error) {
 	var task types.EremeticTask
 
-	err := ensureDB()
-	if err != nil {
-		return task, err
-	}
-
-	err = boltdb.View(func(tx *bolt.Tx) error {
+	err := db.database.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte("tasks"))
 		if b == nil {
 			return bolt.ErrBucketNotFound
@@ -69,15 +111,10 @@ func ReadUnmaskedTask(id string) (types.EremeticTask, error) {
 
 // ListNonTerminalTasks returns a list of tasks that are not yet finished in one
 // way or another.
-func ListNonTerminalTasks() ([]*types.EremeticTask, error) {
+func (db boltDriver) ListNonTerminalTasks() ([]*types.EremeticTask, error) {
 	var tasks []*types.EremeticTask
 
-	err := ensureDB()
-	if err != nil {
-		return tasks, err
-	}
-
-	err = boltdb.View(func(tx *bolt.Tx) error {
+	err := db.database.View(func(tx *bolt.Tx) error {
 		b := tx.Bucket([]byte("tasks"))
 		if b == nil {
 			return bolt.ErrBucketNotFound
@@ -96,9 +133,3 @@ func ListNonTerminalTasks() ([]*types.EremeticTask, error) {
 
 	return tasks, err
 }
-
-func applyMask(task *types.EremeticTask) {
-	for k := range task.MaskedEnvironment {
-		task.MaskedEnvironment[k] = "*******"
-	}
-}

database/database_test.go → database/boltdb_test.go

@@ -13,12 +13,23 @@ import (
 	. "github.com/smartystreets/goconvey/convey"
 )
 
-var testDB string
+var (
+	testDB string
+	db     boltDriver
+)
 
 func setup() error {
 	dir, _ := ioutil.TempDir("", "eremetic")
 	testDB = fmt.Sprintf("%s/test.db", dir)
-	return NewDB(testDB)
+	adb, err := NewDB("boltdb", testDB)
+
+	if err != nil {
+		return err
+	}
+
+	db = adb.(boltDriver)
+
+	return nil
 }
 
 func teardown() {
@@ -35,53 +46,43 @@ func TestDatabase(t *testing.T) {
 
 	Convey("NewDB", t, func() {
 		Convey("With an absolute path", func() {
-			err := setup()
+			setup()
 			defer teardown()
-			defer Close()
-
-			So(boltdb.Path(), ShouldNotBeEmpty)
-			So(filepath.IsAbs(boltdb.Path()), ShouldBeTrue)
-			So(err, ShouldBeNil)
-		})
-
-		Convey("With a relative path", func() {
-			NewDB("db/test.db")
-			defer Close()
+			defer db.Close()
 
-			dir, _ := os.Getwd()
-			So(filepath.IsAbs(boltdb.Path()), ShouldBeTrue)
-			So(boltdb.Path(), ShouldEqual, fmt.Sprintf("%s/../db/test.db", dir))
+			So(db.database.Path(), ShouldNotBeEmpty)
+			So(filepath.IsAbs(db.database.Path()), ShouldBeTrue)
 		})
 	})
 
 	Convey("Close", t, func() {
 		setup()
 		defer teardown()
-		Close()
+		db.Close()
 
-		So(boltdb.Path(), ShouldBeEmpty)
+		So(db.database.Path(), ShouldBeEmpty)
 	})
 
 	Convey("Clean", t, func() {
 		setup()
 		defer teardown()
-		defer Close()
+		defer db.Close()
 
-		PutTask(&types.EremeticTask{ID: "1234"})
-		task, _ := ReadTask("1234")
+		db.PutTask(&types.EremeticTask{ID: "1234"})
+		task, _ := db.ReadTask("1234")
 		So(task, ShouldNotEqual, types.EremeticTask{})
 		So(task.ID, ShouldNotBeEmpty)
 
-		Clean()
+		db.Clean()
 
-		task, _ = ReadTask("1234")
+		task, _ = db.ReadTask("1234")
 		So(task, ShouldBeZeroValue)
 	})
 
 	Convey("Put and Read Task", t, func() {
 		setup()
 		defer teardown()
-		defer Close()
+		defer db.Close()
 
 		var maskedEnv = make(map[string]string)
 		maskedEnv["foo"] = "bar"
@@ -100,21 +101,21 @@ func TestDatabase(t *testing.T) {
 			MaskedEnvironment: maskedEnv,
 		}
 
-		PutTask(&task1)
-		PutTask(&task2)
+		db.PutTask(&task1)
+		db.PutTask(&task2)
 
-		t1, err := ReadTask(task1.ID)
+		t1, err := db.ReadTask(task1.ID)
 		So(t1, ShouldResemble, task1)
 		So(err, ShouldBeNil)
-		t2, err := ReadTask(task2.ID)
+		t2, err := db.ReadTask(task2.ID)
 		So(err, ShouldBeNil)
 		So(t2.MaskedEnvironment["foo"], ShouldEqual, "*******")
 	})
 
 	Convey("Read unmasked task", t, func() {
 		setup()
 		defer teardown()
-		defer Close()
+		defer db.Close()
 
 		var maskedEnv = make(map[string]string)
 		maskedEnv["foo"] = "bar"
@@ -131,9 +132,9 @@ func TestDatabase(t *testing.T) {
 			Image:             "busybox",
 			MaskedEnvironment: maskedEnv,
 		}
-		PutTask(&task)
+		db.PutTask(&task)
 
-		t, err := ReadUnmaskedTask(task.ID)
+		t, err := db.ReadUnmaskedTask(task.ID)
 		So(t, ShouldResemble, task)
 		So(err, ShouldBeNil)
 		So(t.MaskedEnvironment, ShouldContainKey, "foo")
@@ -144,12 +145,12 @@ func TestDatabase(t *testing.T) {
 	Convey("List non-terminal tasks", t, func() {
 		setup()
 		defer teardown()
-		defer Close()
+		defer db.Close()
 
-		Clean()
+		db.Clean()
 
 		// A terminated task
-		PutTask(&types.EremeticTask{
+		db.PutTask(&types.EremeticTask{
 			ID: "1234",
 			Status: []types.Status{
 				types.Status{
@@ -168,7 +169,7 @@ func TestDatabase(t *testing.T) {
 		})
 
 		// A running task
-		PutTask(&types.EremeticTask{
+		db.PutTask(&types.EremeticTask{
 			ID: "2345",
 			Status: []types.Status{
 				types.Status{
@@ -182,7 +183,7 @@ func TestDatabase(t *testing.T) {
 			},
 		})
 
-		tasks, err := ListNonTerminalTasks()
+		tasks, err := db.ListNonTerminalTasks()
 		So(err, ShouldBeNil)
 		So(tasks, ShouldHaveLength, 1)
 		task := tasks[0]

database/database.go

@@ -1,54 +1,22 @@
 package database
 
-import (
-	"fmt"
-	"os"
-	"path/filepath"
-
-	"github.com/boltdb/bolt"
-	"github.com/spf13/viper"
-)
-
-var boltdb *bolt.DB
-
-// NewDB is used to load the database handler into memory.
-// It will create a new database file if it doesn't already exist.
-func NewDB(file string) error {
-	if !filepath.IsAbs(file) {
-		dir, _ := os.Getwd()
-		file = fmt.Sprintf("%s/../%s", dir, file)
-	}
-	os.MkdirAll(filepath.Dir(file), 0755)
-
-	db, err := bolt.Open(file, 0600, nil)
-	boltdb = db
-	return err
-}
-
-// Clean is used to delete the tasks bucket
-func Clean() error {
-	return boltdb.Update(func(tx *bolt.Tx) error {
-		if err := tx.DeleteBucket([]byte("tasks")); err != nil {
-			return err
-		}
-		if _, err := tx.CreateBucketIfNotExists([]byte("tasks")); err != nil {
-			return err
-		}
-		return nil
-	})
+import "github.com/klarna/eremetic/types"
+
+type TaskDB interface {
+	Clean() error
+	Close()
+	PutTask(task *types.EremeticTask) error
+	ReadTask(id string) (types.EremeticTask, error)
+	ReadUnmaskedTask(id string) (types.EremeticTask, error)
+	ListNonTerminalTasks() ([]*types.EremeticTask, error)
 }
 
-// Close is used to Close the database
-func Close() {
-	if boltdb != nil {
-		boltdb.Close()
-	}
+func NewDB(driver string, location string) (TaskDB, error) {
+	return boltDB(location)
 }
 
-func ensureDB() error {
-	if boltdb == nil {
-		err := NewDB(viper.GetString("database"))
-		return err
+func applyMask(task *types.EremeticTask) {
+	for k := range task.MaskedEnvironment {
+		task.MaskedEnvironment[k] = "*******"
 	}
-	return nil
 }