// JobByIDRouteHandler provides the handler for jobs data for the given ID
func JobByIDRouteHandler(ren *render.Render, conf *config.Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
vars := mux.Vars(r)
jid := vars["id"]
jobID, err := strconv.Atoi(jid)
if err != nil {
log.Println(err)
ren.JSON(w, 400, map[string]error{"error": err})
return
}
db, err := database.NewDatabase(conf)
if err != nil {
log.Println(err)
ren.JSON(w, http.StatusOK, map[string]error{"error": err})
return
}
defer db.Conn.Close()
if t := db.GetJobByID(jobID); len(t) > 0 {
ren.JSON(w, http.StatusOK, map[string]interface{}{"task": t})
} else {
ren.JSON(w, http.StatusOK, map[string]interface{}{"task": ErrNoJobsFound.Error()})
}
}
}
func watchForNewJobs() error {
q, err := nsq.NewConsumer("new_job", "add", nsqConfig)
if err != nil {
return err
}
var j *database.Job
q.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
json.Unmarshal(message.Body, &j)
db, err := database.NewDatabase(Conf)
if err != nil {
log.Println(err)
}
defer db.Conn.Close()
db.AddJob(*j)
return nil
}))
err = q.ConnectToNSQD(fmt.Sprintf("%s:4150", queueHostFlag))
if err != nil {
return err
}
return nil
}
// Run starts the dispatcher
func (d *Dispatcher) Run() error {
db, err := database.NewDatabase(d.Conf)
if err != nil {
log.Fatalln(err)
}
defer db.Conn.Close()
log.Println("Starting scheduler...")
d.cron.Start()
defer d.cron.Stop()
d.AddExistingTasks()
for {
select {
case data := <-d.SenderChan:
if err := data.Send(db); err != nil {
log.Println(err)
}
// listen for new tasks and add them to the scheduler
case task := <-d.TaskProcChan:
if err := d.cron.AddFunc(task.Schedule, d.taskFuncFactory(&task)); err != nil {
log.Println(err)
}
// remove tasks from the scheduler
case task := <-d.RemoveTaskChan:
entries := d.cron.Entries()
//a = append(a[:i], a[i+1:]...)
log.Println(task)
for i, t := range entries {
log.Printf("%d: %+v\n", i, t)
}
}
}
}
// TaskDeleteByIDRouteHandler deletes the task data for the given ID
func TaskDeleteByIDRouteHandler(ren *render.Render, conf *config.Config, dispatch *dispatcher.Dispatcher) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
vars := mux.Vars(r)
tid := vars["id"]
taskID, err := strconv.Atoi(tid)
if err != nil {
log.Println(err)
}
db, err := database.NewDatabase(conf)
if err != nil {
log.Println(err)
}
defer db.Conn.Close()
task := db.GetTaskByID(taskID)
if len(task) > 0 {
dispatch.RemoveTaskChan <- task[0]
db.DeleteTask(taskID)
} else {
ren.JSON(w, http.StatusOK, map[string]interface{}{"error": ErrNoEntryFound.Error()})
return
}
ren.JSON(w, http.StatusOK, map[string]interface{}{"task": taskID})
}
}
// CommandsRouteHandler provides the handler for commands data
func CommandsRouteHandler(ren *render.Render, conf *config.Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
db, err := database.NewDatabase(conf)
if err != nil {
log.Println(err)
}
defer db.Conn.Close()
ren.JSON(w, http.StatusOK, map[string]interface{}{"commands": db.GetCommands()})
}
}
// AddExistingTasks adds active tasks from the database to the scheduler
func (d *Dispatcher) AddExistingTasks() {
db, err := database.NewDatabase(d.Conf)
if err != nil {
log.Fatalln(err)
}
defer db.Conn.Close()
tasks := db.GetTasks()
for _, task := range tasks {
d.cron.AddFunc(task.Schedule, d.taskFuncFactory(&task))
}
}
// Run runs the job manager
func (j *JobManager) Run() {
db, err := database.NewDatabase(j.Conf)
if err != nil {
log.Fatalln(err)
}
for {
select {
case data := <-j.JobChan:
db.AddJob(data)
case data := <-j.TaskChan:
db.AddTask(data)
case <-j.ExitChan:
return
}
}
}
// CommandDeleteByIDRouteHandler deletes the command data for the given ID
func CommandDeleteByIDRouteHandler(ren *render.Render, conf *config.Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
vars := mux.Vars(r)
tid := vars["id"]
commandID, err := strconv.Atoi(tid)
if err != nil {
log.Println(err)
}
db, err := database.NewDatabase(conf)
if err != nil {
log.Println(err)
}
defer db.Conn.Close()
db.DeleteCommand(commandID)
ren.JSON(w, http.StatusOK, map[string]interface{}{"command": commandID})
}
}
// CommandByIDRouteHandler provides the handler for commands data for the given ID
func CommandByIDRouteHandler(ren *render.Render, conf *config.Config) http.HandlerFunc {
return func(w http.ResponseWriter, r *http.Request) {
vars := mux.Vars(r)
tid := vars["id"]
commandID, err := strconv.Atoi(tid)
if err != nil {
log.Println(err)
}
db, err := database.NewDatabase(conf)
if err != nil {
log.Println(err)
}
defer db.Conn.Close()
if t := db.GetCommandByID(commandID); len(t) > 0 {
ren.JSON(w, http.StatusOK, map[string]interface{}{"command": t})
} else {
ren.JSON(w, http.StatusOK, map[string]interface{}{"command": ErrNoCommandsFound.Error()})
}
}
}
func watchForNewCommands() error {
q, err := nsq.NewConsumer("new_command", "add", nsqConfig)
if err != nil {
return err
}
var c *database.Command
q.AddHandler(nsq.HandlerFunc(func(message *nsq.Message) error {
json.Unmarshal(message.Body, &c)
db, err := database.NewDatabase(Conf)
if err != nil {
log.Println(err)
}
db.AddCommand(*c)
return nil
}))
err = q.ConnectToNSQD(fmt.Sprintf("%s:4150", queueHostFlag))
if err != nil {
return err
}
return nil
}
// generateTaskFunc generates a function suitable for use in the scheduler
func (d *Dispatcher) taskFuncFactory(task *database.Task) func() {
fn := func() {
var r database.Result
r.TaskID = task.ID
out, err := exec.Command(task.CMD, task.Args).Output()
if err != nil {
log.Println(err)
r.Error = err.Error()
}
db, err := database.NewDatabase(d.Conf)
if err != nil {
log.Fatalln(err)
}
defer db.Conn.Close()
r.Result = out
r.Send(db)
}
task.Func = fn
return fn
}
func main() {
flag.Parse()
signal.Notify(signalsChan, os.Interrupt)
go func() {
for sig := range signalsChan {
log.Printf("Exiting... %v\n", sig)
signalsChan = nil
os.Exit(1)
}
}()
if queueHostFlag == "" || dbUserFlag == "" || dbPassFlag == "" ||
dbHostFlag == "" || dbPortFlag == 0 || dbNameFlag == "" || resultWorkers < 3 {
flag.Usage()
os.Exit(1)
}
// assign
Conf = &config.Config{
Database: config.DBConf{
DBUser: dbUserFlag,
DBPass: dbPassFlag,
DBHost: dbHostFlag,
DBPort: dbPortFlag,
DBName: dbNameFlag,
},
QueueHost: queueHostFlag,
ResultWorkers: resultWorkers,
}
if dbSetupFlag {
fmt.Println(Conf.Database.DBUser, Conf.Database.DBPass, Conf.Database.DBHost, Conf.Database.DBPort, Conf.Database.DBName)
db, err := database.NewDatabase(Conf)
if err != nil {
log.Fatalln(err)
}
defer db.Conn.Close()
db.Setup()
os.Exit(0)
}
dispatch := dispatcher.NewDispatcher(Conf)
go dispatch.Run()
go watchForNewJobs()
go watchForNewTasks()
go watchForNewResults()
go watchForNewCommands()
// setup the renderer for returning our JSON
ren := render.New(render.Options{})
store := cookiestore.New([]byte(uuid.NewUUID().String()))
// initialize the web framework
n := negroni.New(
negroni.NewRecovery(),
negroni.NewLogger(),
negroni.NewStatic(http.Dir("public")),
)
n.Use(sessions.Sessions("session", store))
statsMiddleware := stats.New()
// create a router to handle the requests coming in to our endpoints
router := mux.NewRouter()
// Frontend Entry Point
router.HandleFunc(controllers.FrontEnd, controllers.FrontendHandler()).Methods("GET")
// Jobs Route
router.HandleFunc(controllers.JobsPath, controllers.JobsRouteHandler(ren, Conf)).Methods("GET")
// Job By ID Route
router.HandleFunc(controllers.JobByID, controllers.JobByIDRouteHandler(ren, Conf)).Methods("GET")
// Job Delete By ID Route
router.HandleFunc(controllers.JobByID, controllers.JobDeleteByIDRouteHandler(ren, Conf)).Methods("DELETE")
// New Jobs Route
router.HandleFunc(controllers.JobsPath, controllers.NewJobRouteHandler(ren, dispatch)).Methods("POST")
// Tasks Route
router.HandleFunc(controllers.TasksPath, controllers.TasksRouteHandler(ren, Conf)).Methods("GET")
// Task By ID Route
router.HandleFunc(controllers.TaskByID, controllers.TaskByIDRouteHandler(ren, Conf)).Methods("GET")
// Task Delete By ID Route
router.HandleFunc(controllers.TaskByID, controllers.TaskDeleteByIDRouteHandler(ren, Conf, dispatch)).Methods("DELETE")
// New Tasks Route
router.HandleFunc(controllers.TasksPath, controllers.NewTaskRouteHandler(ren, dispatch)).Methods("POST")
// Commands Route
router.HandleFunc(controllers.CommandsPath, controllers.CommandsRouteHandler(ren, Conf)).Methods("GET")
// Cmmand By ID Route
router.HandleFunc(controllers.CommandByID, controllers.CommandByIDRouteHandler(ren, Conf)).Methods("GET")
// Command Delete By ID Route
router.HandleFunc(controllers.CommandByID, controllers.CommandDeleteByIDRouteHandler(ren, Conf)).Methods("DELETE")
// New Commands Route
router.HandleFunc(controllers.CommandsPath, controllers.NewCommandRouteHandler(ren, dispatch)).Methods("POST")
// API Statistics Route
router.HandleFunc(controllers.APIStats, controllers.AdminAionAPIServerStats(statsMiddleware)).Methods("GET")
n.Use(statsMiddleware)
n.UseHandler(router)
n.Run(portFlag)
}