sampler-fork/data/item.go

package data

import (
	"bufio"
	"errors"
	ui "github.com/gizak/termui/v3"
	"github.com/sqshq/sampler/config"
	"io"
	"os"
	"os/exec"
	"strings"
	"time"
)

type Item struct {
	Label            string
	SampleScript     string
	InitScript       *string
	TransformScript  *string
	Color            *ui.Color
	RateMs           int
	InteractiveShell *InteractiveShell
}

type InteractiveShell struct {
	StdoutCh chan string
	StderrCh chan string
	Stdin    io.WriteCloser
	Cmd      *exec.Cmd
}

func NewItems(cfgs []config.Item, rateMs int) []*Item {

	items := make([]*Item, 0)

	for _, i := range cfgs {
		item := &Item{
			Label:           *i.Label,
			SampleScript:    *i.SampleScript,
			InitScript:      i.InitScript,
			TransformScript: i.TransformScript,
			Color:           i.Color,
			RateMs:          rateMs,
		}
		items = append(items, item)
	}

	return items
}

func (i *Item) nextValue(variables []string) (string, error) {

	if i.InitScript != nil && i.InteractiveShell == nil {
		err := i.initInteractiveShell(variables)
		if err != nil {
			return "", err
		}
	}

	if i.InitScript != nil {
		return i.executeInteractiveShellCmd(variables)
	} else {
		return i.executeCmd(variables, i.SampleScript)
	}
}

func (i *Item) executeCmd(variables []string, script string) (string, error) {

	cmd := exec.Command("sh", "-c", script)
	enrichEnvVariables(cmd, variables)

	output, err := cmd.Output()

	if err != nil {
		return "", err
	}

	return string(output), nil
}

func (i *Item) initInteractiveShell(variables []string) error {

	cmd := exec.Command("sh", "-c", *i.InitScript)
	enrichEnvVariables(cmd, variables)

	stdout, err := cmd.StdoutPipe()
	if err != nil {
		return err
	}

	stderr, err := cmd.StderrPipe()
	if err != nil {
		return err
	}

	stdin, err := cmd.StdinPipe()
	if err != nil {
		return err
	}

	stdoutScanner := bufio.NewScanner(stdout)
	stderrScanner := bufio.NewScanner(stderr)

	stdoutCh := make(chan string)
	stderrCh := make(chan string)

	go func() {
		for stdoutScanner.Scan() {
			stdoutCh <- stdoutScanner.Text()
			stderrCh <- stderrScanner.Text()
		}
	}()

	i.InteractiveShell = &InteractiveShell{
		StdoutCh: stdoutCh,
		StderrCh: stderrCh,
		Stdin:    stdin,
		Cmd:      cmd,
	}

	err = cmd.Start()
	if err != nil {
		return err
	}

	return nil
}

func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) {

	_, err := io.WriteString(i.InteractiveShell.Stdin, i.SampleScript+"\n")
	if err != nil {
		return "", err
	}

	timeout := make(chan bool, 1)

	go func() {
		time.Sleep(time.Duration(i.RateMs / 2))
		timeout <- true
	}()

	var resultText strings.Builder
	var errorText strings.Builder

	for {
		select {
		case stdout := <-i.InteractiveShell.StdoutCh:
			if len(stdout) > 0 {
				resultText.WriteString(stdout)
				resultText.WriteString("\n")
			}
		case stderr := <-i.InteractiveShell.StderrCh:
			if len(stderr) > 0 {
				errorText.WriteString(stderr)
				errorText.WriteString("\n")
			}
		case <-timeout:
			if errorText.Len() > 0 {
				return "", errors.New(errorText.String())
			} else {
				return i.transformInteractiveShellCmd(resultText.String())
			}
		}
	}
}

func (i *Item) transformInteractiveShellCmd(sample string) (string, error) {

	if i.TransformScript != nil && len(sample) > 0 {
		return i.executeCmd([]string{"sample=" + sample}, *i.TransformScript)
	}

	return sample, nil
}

func enrichEnvVariables(cmd *exec.Cmd, variables []string) {
	cmd.Env = os.Environ()
	for _, variable := range variables {
		cmd.Env = append(cmd.Env, variable)
	}
}