Add expired status to assignments fetching params

Author: edshadi

pkg/api/assignmentfetcher/transport.go

@@ -50,6 +50,11 @@ func decodeAssignmentsFetcherRequest(_ context.Context, r *http.Request) (interf
 		as.TaskID = taskID[0]
 	}
 
+	status, ok := params["status"]
+	if ok && len(taskID) > 0 {
+		as.Status = assignment.Status(status[0])
+	}
+
 	return as, nil
 }
 

pkg/assignment/assignment.go

@@ -14,6 +14,7 @@ const (
 	Pending  Status = "pending"
 	Accepted Status = "accepted"
 	Rejected Status = "rejected"
+	Expired  Status = "expired"
 )
 
 type Params struct {

pkg/datastore/assignment.go

@@ -18,8 +18,10 @@ type Storage interface {
 	GetWhitelist(jobID uint64, workerID uint64) (*whitelist.Whitelist, error)
 	WorkerAlreadyAssigned(jobID uint64, workerID uint64) (bool, error)
 	DeleteAssignment(workerID uint64, jobID uint64) (bool, error)
+	DeleteAssignments(ids []uint64) error
 	UpdateAssignment(workerID uint64, jobID uint64, status string) (bool, error)
 	CreateSettings(assignment.Settings) (*assignment.Settings, error)
+	SelectExpiredAssignments() (assignment.Assignments, error)
 }
 
 type AssignmentStore struct {
@@ -33,10 +35,16 @@ func NewAssignmentStore(db *sqlx.DB) *AssignmentStore {
 }
 
 func (as *AssignmentStore) GetAssignments(p assignment.Params) (assignment.Assignments, error) {
+	// If we're looking for expired assignments, then return them
+	if p.Status == assignment.Expired {
+		return as.SelectExpiredAssignments()
+	}
+
 	assignments := assignment.Assignments{}
 	query := "SELECT * FROM assignments"
 	paramsQuery := []string{}
 	args := []interface{}{}
+
 	if p.WorkerID != "" {
 		args = append(args, p.WorkerID)
 		paramsQuery = append(paramsQuery, "worker_id=?")
@@ -200,3 +208,34 @@ func (as *AssignmentStore) CreateSettings(s assignment.Settings) (*assignment.Se
 
 	return set, nil
 }
+
+func (as *AssignmentStore) SelectExpiredAssignments() (assignment.Assignments, error) {
+	assignments := assignment.Assignments{}
+	err := as.DB.Select(
+		&assignments,
+		`SELECT * FROM assignments WHERE expires_at <= NOW()`,
+	)
+	if err != nil {
+		return nil, err
+	}
+	return assignments, nil
+}
+
+func (as *AssignmentStore) DeleteAssignments(ids []uint64) error {
+	query, args, _ := sqlx.In(`DELETE FROM assignments WHERE id IN (?)`, ids)
+	query = as.DB.Rebind(query)
+
+	res, err := as.DB.Exec(query, args...)
+
+	if err != nil {
+		return err
+	}
+
+	numRows, _ := res.RowsAffected()
+
+	if numRows != int64(len(ids)) {
+		return RecordsMismatch{}
+	}
+
+	return nil
+}

pkg/datastore/errors.go

@@ -28,3 +28,9 @@ type NoRowErr struct{}
 func (err NoRowErr) Error() string {
 	return "No Records"
 }
+
+type RecordsMismatch struct{}
+
+func (err RecordsMismatch) Error() string {
+	return "Some records where not processed properly"
+}

pkg/service/service.go

@@ -18,6 +18,7 @@ type AssignmentService interface {
 	DeleteAssignment(workerID uint64, jobID uint64) (bool, error)
 	UpdateAssignment(workerID uint64, jobID uint64, status string) (bool, error)
 	CreateSettings(assignment.Settings) (*assignment.Settings, error)
+	GetStore() datastore.Storage
 }
 
 type service struct {
@@ -32,6 +33,10 @@ func New(s datastore.Storage, a authorization.Authorizer) *service {
 	}
 }
 
+func (s *service) GetStore() datastore.Storage {
+	return s.store
+}
+
 func (s *service) Healthy() bool {
 	return true
 }