Merge pull request #126 from vcabbage/refactor-agent-cache

Refactor agent/cache package

agent/cache/cache.go

@@ -23,50 +23,52 @@ import (
 	"database/sql"
 	"fmt"
 	"os"
+	"path/filepath"
 
-	log "github.com/Sirupsen/logrus"
 	_ "github.com/mattn/go-sqlite3" // This look strange but is necessary - the sqlite package is used indirectly by database/sql
+	"github.com/pkg/errors"
 
 	"github.com/toddproject/todd/config"
 )
 
-func NewAgentCache(cfg config.Config) *AgentCache {
-	var ac AgentCache
-	ac.dbLoc = fmt.Sprintf("%s/agent_cache.db", cfg.LocalResources.OptDir)
-	return &ac
-}
-
+// AgentCache provides methods for interacting with the on disk cache.
 type AgentCache struct {
 	// Need similar abstractions to what you did in the tasks package
-	dbLoc string
+	db *sql.DB
 }
 
-// TODO(mierdin): Handling errors in this package?
+// New returns an initialized instance of AgentCache.
+func New(cfg config.Config) (*AgentCache, error) {
 
-// Init will set up the sqlite database to serve as this agent cache.
-func (ac AgentCache) Init() {
+	dbLoc := filepath.Join(cfg.LocalResources.OptDir, "agent_cache.db")
 
 	// Clean up any old cache data
-	os.Remove(ac.dbLoc)
+	err := os.Remove(dbLoc)
+	if err != nil && !os.IsNotExist(err) {
+		return nil, errors.Wrap(err, "removing existing DB file")
+	}
 
 	// Open connection
-	db, err := sql.Open("sqlite3", ac.dbLoc)
+	db, err := sql.Open("sqlite3", dbLoc)
 	if err != nil {
-		log.Fatal(err)
+		return nil, errors.Wrap(err, "opening DB file")
 	}
-	defer db.Close()
 
 	// Initialize database
-	sqlStmt := `
-    create table testruns (id integer not null primary key, uuid text, testlet text, args text, targets text, results text);
-    delete from testruns;
-    create table keyvalue (id integer not null primary key, key text, value text);
-    delete from keyvalue;
-    `
+	const sqlStmt = `
+		CREATE TABLE testruns (id INTEGER NOT NULL PRIMARY KEY, uuid TEXT, testlet TEXT, args TEXT, targets TEXT, results TEXT);
+		CREATE TABLE keyvalue (id INTEGER NOT NULL PRIMARY KEY, key TEXT, value TEXT);
+	`
 
 	_, err = db.Exec(sqlStmt)
 	if err != nil {
-		log.Errorf("%q: %s\n", err, sqlStmt)
-		return
+		return nil, fmt.Errorf("%q: %s", err, sqlStmt)
 	}
+
+	return &AgentCache{db: db}, nil
+}
+
+// Close closes the underlying database connection.
+func (ac *AgentCache) Close() error {
+	return ac.db.Close()
 }

agent/cache/keyvalue.go

@@ -9,116 +9,77 @@
 package cache
 
 import (
-	"database/sql"
-	"fmt"
-
 	log "github.com/Sirupsen/logrus"
-	_ "github.com/mattn/go-sqlite3" // This look strange but is necessary - the sqlite package is used indirectly by database/sql
+	"github.com/pkg/errors"
 )
 
 // GetKeyValue will retrieve a value from the agent cache using a key string
-func (ac AgentCache) GetKeyValue(key string) string {
-	// Open connection
-	db, err := sql.Open("sqlite3", ac.dbLoc)
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer db.Close()
-
-	log.Debugf("Retrieving value of key - %s", key)
+func (ac *AgentCache) GetKeyValue(key string) (string, error) {
+	log.Debugf("Retrieving value of key - %s\n", key)
 
 	// First, see if the key exists.
-	rows, err := db.Query(fmt.Sprintf("select value from keyvalue where key = \"%s\" ", key))
+	rows, err := ac.db.Query("SELECT value FROM keyvalue WHERE key = ?", key)
 	if err != nil {
-		log.Fatal(err)
+		return "", errors.Wrap(err, "querying DB")
 	}
-	value := ""
 	defer rows.Close()
+
+	var value string
 	for rows.Next() {
-		rows.Scan(&value)
+		err = rows.Scan(&value)
+		if err != nil {
+			return "", errors.Wrap(err, "scanning values retrieved from DB")
+		}
 	}
-	return value
+	return value, nil
 }
 
 // SetKeyValue sets a KeyValue pair within the agent cache
-func (ac AgentCache) SetKeyValue(key, value string) error {
-
-	// Open connection
-	db, err := sql.Open("sqlite3", ac.dbLoc)
-	if err != nil {
-		log.Fatal(err)
-	}
-	defer db.Close()
-
-	log.Debugf("Writing keyvalue pair to agent cache - %s:%s", key, value)
+func (ac *AgentCache) SetKeyValue(key, value string) error {
+	log.Debugf("Writing keyvalue pair to agent cache - %s:%s\n", key, value)
 
 	// First, see if the key exists.
-	rows, err := db.Query(fmt.Sprintf("select key, value FROM keyvalue WHERE KEY = \"%s\";", key))
+	rows, err := ac.db.Query("SELECT count(1) FROM keyvalue WHERE KEY = ?", key)
 	if err != nil {
-		log.Fatal(err)
+		return errors.Wrap(err, "querying count from DB")
 	}
-	rowcount := 0
 	defer rows.Close()
-	for rows.Next() {
-		rowcount++
-	}
-
-	if rowcount != 1 {
 
-		// If there is MORE than one row, we should delete the extras first
-		// TODO(mierdin): Is this really necessary?
-		if rowcount > 1 {
-			log.Warn("Extra keyvalue pair detected. Deleting and inserting new record.")
-			tx, err := db.Begin()
-			if err != nil {
-				log.Fatal(err)
-			}
-			stmt, err := tx.Prepare(fmt.Sprintf("DELETE FROM keyvalue WHERE KEY = \"%s\";", key))
-			if err != nil {
-				log.Fatal(err)
-			}
-			defer stmt.Close()
-			_, err = stmt.Exec()
-			if err != nil {
-				log.Fatal(err)
-			}
-		}
-
-		// Begin Insert
-		tx, err := db.Begin()
-		if err != nil {
-			log.Fatal(err)
-		}
-		stmt, err := tx.Prepare("insert into keyvalue(key, value) values(?, ?)")
+	var rowCount int
+	for rows.Next() {
+		err = rows.Scan(&rowCount)
 		if err != nil {
-			log.Fatal(err)
+			return errors.Wrap(err, "scanning rowCount from DB")
 		}
-		defer stmt.Close()
-		_, err = stmt.Exec(key, value)
+	}
+
+	if rowCount != 1 {
+		tx, err := ac.db.Begin()
 		if err != nil {
-			log.Fatal(err)
+			return errors.Wrap(err, "starting DB transaction")
 		}
-		tx.Commit()
 
-	} else {
+		if rowCount > 1 {
+			// If there is MORE than one row, we should delete the extras first
+			// TODO(mierdin): Is this really necessary?
+			log.Warn("Extra keyvalue pair detected. Deleting and inserting new record.")
 
-		// Begin Update
-		tx, err := db.Begin()
-		if err != nil {
-			log.Fatal(err)
+			_, err = tx.Exec("DELETE FROM keyvalue WHERE KEY = ?", key)
+			if err != nil {
+				tx.Rollback()
+				return errors.Wrap(err, "deleteing keyvalues")
+			}
 		}
 
-		stmt, err := tx.Prepare(fmt.Sprintf("update keyvalue set value = \"%s\" where key = \"%s\" ", value, key))
-		if err != nil {
-			log.Fatal(err)
-		}
-		defer stmt.Close()
-		_, err = stmt.Exec()
+		_, err = tx.Exec("INSERT INTO keyvalue(key, value) values(?, ?)", key, value)
 		if err != nil {
-			log.Fatal(err)
+			tx.Rollback()
+			return errors.Wrap(err, "inserting keyvalue into DB")
 		}
-		tx.Commit()
 
+		return errors.Wrap(tx.Commit(), "commmitting transaction")
 	}
-	return nil
+
+	_, err = ac.db.Exec("UPDATE keyvalue SET value = ? WHERE key = ?", value, key)
+	return errors.Wrap(err, "updating keyvalue")
 }