From ca931cd9eb0ce27700389a647516858cac7c2806 Mon Sep 17 00:00:00 2001
From: Brad Rydzewski <brad.rydzewski@gmail.com>
Date: Wed, 29 Apr 2015 15:26:22 -0700
Subject: [PATCH] ability to persist agent data

---
 common/agent.go         |  9 +++++++++
 datastore/bolt/agent.go |  1 +
 datastore/bolt/bolt.go  |  2 ++
 datastore/bolt/build.go | 29 +++++++++++++++++++++++++++--
 datastore/datastore.go  | 12 ++++++++++++
 5 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 common/agent.go
 create mode 100644 datastore/bolt/agent.go

diff --git a/common/agent.go b/common/agent.go
new file mode 100644
index 00000000..4e4a8f71
--- /dev/null
+++ b/common/agent.go
@@ -0,0 +1,9 @@
+package common
+
+// Agent represents a worker that has connected
+// to the system in order to perform work
+type Agent struct {
+	Name      string
+	Addr      string
+	IsHealthy bool
+}
diff --git a/datastore/bolt/agent.go b/datastore/bolt/agent.go
new file mode 100644
index 00000000..a3ce8fd0
--- /dev/null
+++ b/datastore/bolt/agent.go
@@ -0,0 +1 @@
+package bolt
diff --git a/datastore/bolt/bolt.go b/datastore/bolt/bolt.go
index cd0186e3..dfcbc5e5 100644
--- a/datastore/bolt/bolt.go
+++ b/datastore/bolt/bolt.go
@@ -23,6 +23,7 @@ var (
 	bucketRepoParams  = []byte("repo_params")
 	bucketRepoUsers   = []byte("repo_users")
 	bucketBuild       = []byte("build")
+	bucketBuildAgent  = []byte("build_agent")
 	bucketBuildStatus = []byte("build_status")
 	bucketBuildLogs   = []byte("build_logs")
 	bucketBuildSeq    = []byte("build_seq")
@@ -49,6 +50,7 @@ func New(path string) (*DB, error) {
 		tx.CreateBucketIfNotExists(bucketRepoParams)
 		tx.CreateBucketIfNotExists(bucketRepoUsers)
 		tx.CreateBucketIfNotExists(bucketBuild)
+		tx.CreateBucketIfNotExists(bucketBuildAgent)
 		tx.CreateBucketIfNotExists(bucketBuildStatus)
 		tx.CreateBucketIfNotExists(bucketBuildLogs)
 		tx.CreateBucketIfNotExists(bucketBuildSeq)
diff --git a/datastore/bolt/build.go b/datastore/bolt/build.go
index 9c0f90d8..90e34b25 100644
--- a/datastore/bolt/build.go
+++ b/datastore/bolt/build.go
@@ -3,10 +3,11 @@ package bolt
 import (
 	//"bytes"
 	"encoding/binary"
-	"github.com/boltdb/bolt"
-	"github.com/drone/drone/common"
 	"strconv"
 	"time"
+
+	"github.com/boltdb/bolt"
+	"github.com/drone/drone/common"
 )
 
 // Build gets the specified build number for the
@@ -213,3 +214,27 @@ func (db *DB) SetBuildTask(repo string, build int, task *common.Task) error {
 	})
 }
 
+// SetBuildAgent insert or updates the agent that is
+// running a build.
+func (db *DB) SetBuildAgent(repo string, build int, agent *common.Agent) error {
+	key := []byte(repo + "/" + strconv.Itoa(build))
+	return db.Update(func(t *bolt.Tx) error {
+		return update(t, bucketBuildAgent, key, agent)
+	})
+}
+
+func (db *DB) DelBuildAgent(repo string, build int, agent *common.Agent) error {
+	key := []byte(repo + "/" + strconv.Itoa(build))
+	return db.Update(func(t *bolt.Tx) error {
+		return delete(t, bucketBuildAgent, key)
+	})
+}
+
+func (db *DB) BuildAgent(repo string, build int) (*common.Agent, error) {
+	key := []byte(repo + "/" + strconv.Itoa(build))
+	agent := &common.Agent{}
+	err := db.View(func(t *bolt.Tx) error {
+		return get(t, bucketBuildAgent, key, agent)
+	})
+	return agent, err
+}
diff --git a/datastore/datastore.go b/datastore/datastore.go
index cae1993e..31df8614 100644
--- a/datastore/datastore.go
+++ b/datastore/datastore.go
@@ -101,6 +101,10 @@ type Datastore interface {
 	// named repository.
 	BuildLast(string) (*common.Build, error)
 
+	// BuildAgent returns the agent that is being
+	// used to execute the build.
+	// BuildAgent(string, int) (*common.Agent, error)
+
 	// SetBuild inserts or updates a build for the named
 	// repository. The build number is incremented and
 	// assigned to the provided build.
@@ -120,6 +124,14 @@ type Datastore interface {
 	// an error is returned.
 	SetBuildTask(string, int, *common.Task) error
 
+	// SetBuildAgent insert or updates the agent that is
+	// running a build.
+	// SetBuildAgent(string, int, *common.Agent) error
+
+	// DelBuildAgent purges the referce to the agent
+	// that ran a build.
+	// DelBuildAgent(string, int, *common.Agent) error
+
 	// LogReader gets the task logs at index N for
 	// the named repository and build number.
 	LogReader(string, int, int) (io.Reader, error)