datagraph/script_scheduler.go

package main

import (
	// External
	werr "git.gibonuddevalla.se/go/wrappederror"

	// Standard
	"bytes"
	"database/sql"
	"fmt"
	"os/exec"
	"strings"
	"time"
)

type ScriptScheduler struct {
	EventQueue chan string
}

type ScriptExecution struct {
	ID           int
	TimeStart    sql.NullTime `db:"time_start"`
	TimeEnd      sql.NullTime `db:"time_end"`
	Source       []byte
	Data         []byte
	SSH          string
	OutputStdout sql.NullString `db:"output_stdout"`
	OutputStderr sql.NullString `db:"output_stderr"`
	ExitCode     sql.NullInt16
}

func NewScriptScheduler() (sched ScriptScheduler) {
	sched.EventQueue = make(chan string, 64)
	return
}

func (self ScriptScheduler) Loop() { // {{{

	// Lets check for somehow missed executions every minute.
	// An event SHOULD be received for each new created, but let's be sure.
	tick := time.NewTicker(time.Second * 60)

	var event string
	for {
		select {
		case <-tick.C:
			self.HandleNextExecution()

		case event = <-self.EventQueue:
			if event == "SCRIPT_SCHEDULED" {
				self.HandleNextExecution()
			}
		}
	}
} // }}}
func (self ScriptScheduler) HandleNextExecution() { // {{{
	se, err := self.GetNextExecution()
	if err != nil {
		logger.Error("script_scheduler", "error", err)
		return
	}

	if se.ID == 0 {
		return
	}

	// Setting the time_start value on the database row makes sure it doesn't get handled again.
	se.TimeStart.Time = time.Now()
	se.TimeStart.Valid = true
	se.Update()

	logger.Info("script_scheduler", "op", "execute", "id", se.ID)
	fname, err := se.GetScriptTempFilename()
	if err != nil {
		err = werr.Wrap(err)
		logger.Error("script_execution", "op", "get_script_temp_filename", "id", se.ID, "error", err)
		return
	}

	err = se.UploadScript(fname)
	if err != nil {
		err = werr.Wrap(err)
		logger.Error("script_execution", "op", "upload_script", "id", se.ID, "error", err)
		return
	}

	err = se.RunScript(fname)
	if err != nil {
		err = werr.Wrap(err)
		logger.Error("script_execution", "op", "run_script", "id", se.ID, "error", err)
		return
	}

	se.SSHCommand([]byte{}, false, fmt.Sprintf("rm %s", fname))

	logger.Info("script_scheduler", "op", "handled", "script", fname)
} // }}}
func (self ScriptScheduler) GetNextExecution() (e ScriptExecution, err error) { // {{{
	row := db.QueryRowx(`
		SELECT
			e.id,
			time_start,
			time_end,
			data,
			ssh,
			output_stdout,
			output_stderr,
			exitcode,
			sl.source
		FROM execution e
		INNER JOIN script_log sl ON e.script_log_id = sl.id
		WHERE
			time_start IS NULL
		ORDER BY
			id ASC
	`)
	err = row.StructScan(&e)

	// Returned ScriptExecution is having an ID of 0 if none was returned
	if err == sql.ErrNoRows {
		err = nil
		return
	}

	if err != nil {
		err = werr.Wrap(err)
		return
	}

	return
} // }}}

func (se *ScriptExecution) Update() (err error) { // {{{
	_, err = db.Exec(`
		UPDATE public.execution
		SET
			time_start = $2,
			time_end = $3,
			output_stdout = $4,
			output_stderr = $5,
			exitcode = $6

		WHERE
			id=$1`,
		se.ID,
		se.TimeStart,
		se.TimeEnd,
		se.OutputStdout,
		se.OutputStderr,
		se.ExitCode,
	)
	if err != nil {
		err = werr.Wrap(err)
		logger.Error("script_execution", "op", "execute", "id", se.ID, "error", err)
		return
	}
	return
} // }}}
func (se *ScriptExecution) SSHCommand(stdin []byte, log bool, args ...string) (stdoutString string, err error) { // {{{
	params := []string{se.SSH}
	params = append(params, args...)
	cmd := exec.Command("ssh", params...)

	stdout := new(bytes.Buffer)
	stderr := new(bytes.Buffer)
	cmd.Stdin = bytes.NewReader(stdin)
	cmd.Stdout = stdout
	cmd.Stderr = stderr
	err = cmd.Run()

	// A cleanup command is run after the script. This shouldn't overwrite the output from the actual script.
	if log {
		se.OutputStdout.String = stdout.String()
		se.OutputStderr.String = stderr.String()
		se.ExitCode.Int16 = int16(cmd.ProcessState.ExitCode())
		se.OutputStdout.Valid = true
		se.OutputStderr.Valid = true
		se.ExitCode.Valid = true
	}

	se.TimeEnd.Time = time.Now()
	se.TimeEnd.Valid = true
	se.Update()

	if err != nil {
		err = werr.Wrap(err)
		return
	}

	return stdout.String(), nil
} // }}}
func (se *ScriptExecution) GetScriptTempFilename() (fname string, err error) { // {{{
	fname, err = se.SSHCommand([]byte{}, true, "mktemp -t datagraph.XXXXXX")
	if err != nil {
		err = werr.Wrap(err)
		return
	}
	fname = strings.TrimSpace(fname)
	return
} // }}}
func (se *ScriptExecution) UploadScript(fname string) (err error) { // {{{
	_, err = se.SSHCommand(se.Source, true, fmt.Sprintf("sh -c 'touch %s && chmod 700 %s && cat >%s'", fname, fname, fname))
	if err != nil {
		err = werr.Wrap(err)
	}
	return
} // }}}
func (se *ScriptExecution) RunScript(fname string) (err error) { // {{{
	_, err = se.SSHCommand(se.Data, true, fname)
	if err != nil {
		err = werr.Wrap(err)
	}
	return
} // }}}