package storage import ( "github.com/sirupsen/logrus" ) type Worker struct { Id int64 `json:"id"` Created int64 `json:"created"` Alias string `json:"alias,omitempty"` Secret []byte `json:"secret"` Paused bool `json:"paused"` } type WorkerStats struct { Alias string `json:"alias"` ClosedTaskCount int64 `json:"closed_task_count"` Paused bool `json:"paused"` Id int64 `json:"id"` } type WorkerAccess struct { Submit bool `json:"submit"` Assign bool `json:"assign"` Request bool `json:"request"` Worker Worker `json:"worker"` Project int64 `json:"project"` } func (database *Database) SaveWorker(worker *Worker) { db := database.getDB() row := db.QueryRow(`INSERT INTO worker (created, secret, alias) VALUES ($1,$2,$3) RETURNING id`, worker.Created, worker.Secret, worker.Alias) err := row.Scan(&worker.Id) handleErr(err) logrus.WithFields(logrus.Fields{ "newId": worker.Id, }).Trace("Database.saveWorker INSERT worker") } func (database *Database) GetWorker(id int64) *Worker { if database.workerCache[id] != nil { return database.workerCache[id] } db := database.getDB() worker := &Worker{} row := db.QueryRow("SELECT id, created, secret, alias FROM worker WHERE id=$1", id) err := row.Scan(&worker.Id, &worker.Created, &worker.Secret, &worker.Alias) if err != nil { logrus.WithFields(logrus.Fields{ "id": id, }).Warn("Database.getWorker SELECT worker NOT FOUND") return nil } logrus.WithFields(logrus.Fields{ "worker": worker, }).Trace("Database.getWorker SELECT worker") database.workerCache[id] = worker return worker } func (database *Database) UpdateWorker(worker *Worker) bool { db := database.getDB() res, err := db.Exec(`UPDATE worker SET alias=$1, paused=$2 WHERE id=$3`, worker.Alias, worker.Paused, worker.Id) handleErr(err) rowsAffected, _ := res.RowsAffected() logrus.WithFields(logrus.Fields{ "rowsAffected": rowsAffected, "worker": worker, }).Trace("Database.UpdateWorker UPDATE worker") database.workerCache[worker.Id] = worker return rowsAffected == 1 } func (database *Database) SaveAccessRequest(wa *WorkerAccess) bool { db := database.getDB() res, err := db.Exec(`INSERT INTO worker_access(worker, project, role_assign, role_submit, request) VALUES ($1,$2,$3,$4,TRUE)`, wa.Worker.Id, wa.Project, wa.Assign, wa.Submit) if err != nil { return false } rowsAffected, _ := res.RowsAffected() logrus.WithFields(logrus.Fields{ "rowsAffected": rowsAffected, }).Trace("Database.SaveAccessRequest INSERT") return rowsAffected == 1 } func (database *Database) AcceptAccessRequest(worker int64, projectId int64) bool { db := database.getDB() database.invalidateAccessCache(worker) res, err := db.Exec(`UPDATE worker_access SET request=FALSE WHERE worker=$1 AND project=$2`, worker, projectId) handleErr(err) rowsAffected, _ := res.RowsAffected() logrus.WithFields(logrus.Fields{ "rowsAffected": rowsAffected, }).Trace("Database.AcceptAccessRequest") return rowsAffected == 1 } func (database *Database) RejectAccessRequest(workerId int64, projectId int64) bool { db := database.getDB() database.invalidateAccessCache(workerId) res, err := db.Exec(`DELETE FROM worker_access WHERE worker=$1 AND project=$2`, workerId, projectId) handleErr(err) rowsAffected, _ := res.RowsAffected() logrus.WithFields(logrus.Fields{ "rowsAffected": rowsAffected, "workerId": workerId, "projectId": projectId, }).Trace("Database.RejectAccessRequest DELETE") return rowsAffected == 1 } func (database *Database) GetAllAccesses(projectId int64) *[]WorkerAccess { db := database.getDB() rows, err := db.Query(`SELECT id, alias, created, paused, role_assign, role_submit, request FROM worker_access INNER JOIN worker w on worker_access.worker = w.id WHERE project=$1 ORDER BY request, alias`, projectId) handleErr(err) requests := make([]WorkerAccess, 0) for rows.Next() { wa := WorkerAccess{ Project: projectId, } _ = rows.Scan(&wa.Worker.Id, &wa.Worker.Alias, &wa.Worker.Created, &wa.Worker.Paused, &wa.Assign, &wa.Submit, &wa.Request) requests = append(requests, wa) } return &requests } func (database *Database) GetAllWorkerStats() *[]WorkerStats { db := database.getDB() rows, err := db.Query(`SELECT alias, closed_task_count, paused, worker.id FROM worker WHERE closed_task_count>0 LIMIT 50`) handleErr(err) if err != nil { return nil } stats := make([]WorkerStats, 0) for rows.Next() { s := WorkerStats{} _ = rows.Scan(&s.Alias, &s.ClosedTaskCount, &s.Paused, &s.Id) stats = append(stats, s) } return &stats }