Adds job whitelist and worker eligibility

Author: edshadi

migrations/sql/3_create_whitelists.down.sql

@@ -0,0 +1 @@
+DROP table IF EXISTS `whitelists`

migrations/sql/3_create_whitelists.up.sql

@@ -0,0 +1,7 @@
+CREATE TABLE `whitelists` (
+  `id` int(11) unsigned NOT NULL AUTO_INCREMENT,
+  `job_id` int(11) unsigned NOT NULL,
+  `worker_id` int(11) unsigned NOT NULL,
+  PRIMARY KEY (`id`),
+  UNIQUE KEY `job_worker` (`job_id`,`worker_id`)
+)

migrations/version

@@ -1 +1 @@
-2
+3

pkg/assignment/errors.go

@@ -29,3 +29,9 @@ type WorkerNotEnoughFunds struct{}
 func (err WorkerNotEnoughFunds) Error() string {
 	return "Worker doesn't have enough funds"
 }
+
+type WorkerNotWhitelisted struct{}
+
+func (err WorkerNotWhitelisted) Error() string {
+	return "Worker has not been whitelisted"
+}

pkg/datastore/assignment.go

@@ -6,6 +6,7 @@ import (
 	"strings"
 
 	"github.com/gemsorg/assignment/pkg/assignment"
+	"github.com/gemsorg/assignment/pkg/whitelist"
 	"github.com/jmoiron/sqlx"
 )
 
@@ -14,6 +15,7 @@ type Storage interface {
 	GetAssignment(id string) (*assignment.Assignment, error)
 	CreateAssignment(assignment.NewAssignment) (*assignment.Assignment, error)
 	GetSettings(jobID uint64) (*assignment.Settings, error)
+	GetWhitelist(jobID uint64, workerID uint64) (*whitelist.Whitelist, error)
 }
 
 type AssignmentStore struct {
@@ -105,3 +107,14 @@ func (as *AssignmentStore) GetSettings(jobID uint64) (*assignment.Settings, erro
 
 	return set, nil
 }
+
+func (as *AssignmentStore) GetWhitelist(jobID uint64, workerID uint64) (*whitelist.Whitelist, error) {
+	wl := &whitelist.Whitelist{}
+	err := as.DB.Get(wl, "SELECT * FROM whitelists WHERE job_id = ? AND worker_id = ?", jobID, workerID)
+
+	if err != nil {
+		return nil, err
+	}
+
+	return wl, nil
+}

pkg/service/service.go

@@ -45,6 +45,13 @@ func (s *service) GetAssignment(id string) (*assignment.Assignment, error) {
 }
 
 func (s *service) CreateAssignment(a assignment.NewAssignment, set *assignment.Settings) (*assignment.Assignment, error) {
+	// if job has a whitelist, check if worker is part of it
+	if set.Whitelist {
+		wl, err := s.store.GetWhitelist(a.JobID, a.WorkerID)
+		if wl == nil || err != nil {
+			return nil, assignment.WorkerNotWhitelisted{}
+		}
+	}
 	// Check if the assignment is allowed
 	allowed, err := a.IsAllowed(set)
 	if !allowed {

pkg/whitelist/whitelist.go

@@ -0,0 +1,7 @@
+package whitelist
+
+type Whitelist struct {
+	ID       uint64 `db:"id"`
+	JobID    uint64 `db:"job_id"`
+	WorkerID uint64 `db:"worker_id"`
+}