added pty interactive shell setting, user now is able to choose between basic interactive shell and PTY interactive shell

component/runchart/runchart.go

@@ -133,13 +133,13 @@ func (c *RunChart) Draw(buffer *ui.Buffer) {
 	c.grid = c.newChartGrid()
 
 	drawArea := image.Rect(
-		c.Inner.Min.X+c.grid.minTimeWidth+1, c.Inner.Min.Y,
+		c.Inner.Min.X+c.grid.minTimeWidth+2, c.Inner.Min.Y,
 		c.Inner.Max.X, c.Inner.Max.Y-xAxisLabelsHeight-1,
 	)
 
+	c.renderAxes(buffer)
 	c.renderLines(buffer, drawArea)
 	c.renderLegend(buffer, drawArea)
-	c.renderAxes(buffer)
 	component.RenderAlert(c.alert, c.Rectangle, buffer)
 	c.mutex.Unlock()
 }

config/component.go

@@ -104,6 +104,7 @@ type LegendConfig struct {
 type Item struct {
 	Label           *string   `yaml:"label,omitempty"`
 	Color           *ui.Color `yaml:"color,omitempty"`
+	Pty             *bool     `yaml:"pty,omitempty"`
 	InitScript      *string   `yaml:"init,omitempty"`
 	SampleScript    *string   `yaml:"sample"`
 	TransformScript *string   `yaml:"transform,omitempty"`

config/default.go

@@ -12,7 +12,7 @@ const (
 
 func (c *Config) setDefaults() {
 	c.setDefaultValues()
-	c.setDefaultColors()
+	c.setDefaultItemSettings()
 	c.setDefaultArrangement()
 }
 
@@ -173,17 +173,21 @@ func setDefaultTriggersValues(triggers []TriggerConfig) {
 	}
 }
 
-func (c *Config) setDefaultColors() {
+func (c *Config) setDefaultItemSettings() {
 
 	palette := console.GetPalette(*c.Theme)
 	colorsCount := len(palette.ContentColors)
+	defaultPty := false
 
 	for _, ch := range c.RunCharts {
 		for j, item := range ch.Items {
 			if item.Color == nil {
 				item.Color = &palette.ContentColors[j%colorsCount]
-				ch.Items[j] = item
 			}
+			if item.Pty == nil {
+				item.Pty = &defaultPty
+			}
+			ch.Items[j] = item
 		}
 	}
 
@@ -191,20 +195,49 @@ func (c *Config) setDefaultColors() {
 		for j, item := range b.Items {
 			if item.Color == nil {
 				item.Color = &palette.ContentColors[j%colorsCount]
-				b.Items[j] = item
 			}
+			if item.Pty == nil {
+				item.Pty = &defaultPty
+			}
+			b.Items[j] = item
 		}
 	}
 
 	for i, s := range c.SparkLines {
 		s.Gradient = &palette.GradientColors[i%(len(palette.GradientColors))]
+		if s.Item.Pty == nil {
+			s.Item.Pty = &defaultPty
+		}
 		c.SparkLines[i] = s
 	}
 
 	for i, g := range c.Gauges {
+		if g.Min.Pty == nil {
+			g.Min.Pty = &defaultPty
+		}
+		if g.Max.Pty == nil {
+			g.Max.Pty = &defaultPty
+		}
+		if g.Cur.Pty == nil {
+			g.Cur.Pty = &defaultPty
+		}
 		if g.Color == nil {
 			g.Color = &palette.ContentColors[i%colorsCount]
+		}
 		c.Gauges[i] = g
 	}
+
+	for i, a := range c.AsciiBoxes {
+		if a.Item.Pty == nil {
+			a.Item.Pty = &defaultPty
+		}
+		c.AsciiBoxes[i] = a
+	}
+
+	for i, t := range c.TextBoxes {
+		if t.Item.Pty == nil {
+			t.Item.Pty = &defaultPty
+		}
+		c.TextBoxes[i] = t
 	}
 }

data/int_pty.go

@@ -0,0 +1,129 @@
+package data
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"github.com/kr/pty"
+	"github.com/lunixbochs/vtclean"
+	"io"
+	"os/exec"
+	"strings"
+	"time"
+)
+
+const (
+	startupTimeout  = 100 * time.Millisecond
+	minAwaitTimeout = 100 * time.Millisecond
+	maxAwaitTimeout = 1 * time.Second
+)
+
+/**
+ * Experimental
+ */
+type PtyInteractiveShell struct {
+	item      *Item
+	variables []string
+	cmd       *exec.Cmd
+	File      io.WriteCloser
+	ch        chan string
+	errCount  int
+}
+
+func (s *PtyInteractiveShell) init() error {
+
+	cmd := exec.Command("sh", "-c", *s.item.initScript)
+	enrichEnvVariables(cmd, s.variables)
+
+	file, err := pty.Start(cmd)
+	if err != nil {
+		return err
+	}
+
+	scanner := bufio.NewScanner(file)
+	channel := make(chan string)
+
+	go func() {
+		for scanner.Scan() {
+			channel <- scanner.Text()
+		}
+	}()
+
+	s.cmd = cmd
+	s.File = file
+	s.ch = channel
+
+	_, err = file.Read(make([]byte, 4096))
+	if err != nil {
+		return err
+	}
+
+	time.Sleep(startupTimeout)
+
+	return nil
+}
+
+func (s *PtyInteractiveShell) execute() (string, error) {
+
+	_, err := io.WriteString(s.File, fmt.Sprintf(" %s\n", s.item.sampleScript))
+	if err != nil {
+		s.errCount++
+		if s.errCount > errorThreshold {
+			s.item.ptyShell = nil // restart session
+		}
+		return "", errors.New(fmt.Sprintf("Failed to execute command: %s", err))
+	}
+
+	softTimeout := make(chan bool, 1)
+	hardTimeout := make(chan bool, 1)
+
+	go func() {
+		time.Sleep(s.getAwaitTimeout() / 2)
+		softTimeout <- true
+		time.Sleep(s.getAwaitTimeout() * 100)
+		hardTimeout <- true
+	}()
+
+	var builder strings.Builder
+	softTimeoutElapsed := false
+
+await:
+	for {
+		select {
+		case out := <-s.ch:
+			cout := vtclean.Clean(out, false)
+			if len(cout) > 0 && !strings.Contains(cout, s.item.sampleScript) {
+				builder.WriteString(cout)
+				builder.WriteString("\n")
+				if softTimeoutElapsed {
+					break await
+				}
+			}
+		case <-softTimeout:
+			if builder.Len() > 0 {
+				break await
+			} else {
+				softTimeoutElapsed = true
+			}
+		case <-hardTimeout:
+			break await
+		}
+	}
+
+	sample := strings.TrimSpace(builder.String())
+
+	return s.item.transform(sample)
+}
+
+func (s *PtyInteractiveShell) getAwaitTimeout() time.Duration {
+
+	timeout := time.Duration(s.item.rateMs) * time.Millisecond
+
+	if timeout > maxAwaitTimeout {
+		return maxAwaitTimeout
+	} else if timeout < minAwaitTimeout {
+		return minAwaitTimeout
+	}
+
+	return timeout
+}

data/int_shell.go

@@ -0,0 +1,110 @@
+package data
+
+import (
+	"bufio"
+	"errors"
+	"fmt"
+	"io"
+	"os/exec"
+	"strings"
+	"time"
+)
+
+type BasicInteractiveShell struct {
+	item      *Item
+	variables []string
+	stdoutCh  chan string
+	stderrCh  chan string
+	stdin     io.WriteCloser
+	cmd       *exec.Cmd
+	errCount  int
+}
+
+func (s *BasicInteractiveShell) init() error {
+
+	cmd := exec.Command("sh", "-c", *s.item.initScript)
+	enrichEnvVariables(cmd, s.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()
+		}
+	}()
+
+	s.stdoutCh = stdoutCh
+	s.stderrCh = stderrCh
+	s.stdin = stdin
+	s.cmd = cmd
+
+	err = cmd.Start()
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (s *BasicInteractiveShell) execute() (string, error) {
+
+	_, err := io.WriteString(s.stdin, s.item.sampleScript+"\n")
+	if err != nil {
+		s.errCount++
+		if s.errCount > errorThreshold {
+			s.item.ptyShell = nil // restart session
+		}
+		return "", errors.New(fmt.Sprintf("Failed to execute command: %s", err))
+	}
+
+	timeout := make(chan bool, 1)
+
+	go func() {
+		time.Sleep(time.Duration(s.item.rateMs / 2))
+		timeout <- true
+	}()
+
+	var resultText strings.Builder
+	var errorText strings.Builder
+
+	for {
+		select {
+		case stdout := <-s.stdoutCh:
+			if len(stdout) > 0 {
+				resultText.WriteString(stdout)
+				resultText.WriteString("\n")
+			}
+		case stderr := <-s.stderrCh:
+			if len(stderr) > 0 {
+				errorText.WriteString(stderr)
+				errorText.WriteString("\n")
+			}
+		case <-timeout:
+			if errorText.Len() > 0 {
+				return "", errors.New(errorText.String())
+			} else {
+				return resultText.String(), nil
+			}
+		}
+	}
+}

data/item.go

@@ -1,26 +1,13 @@
 package data
 
 import (
-	"bufio"
-	"errors"
-	"fmt"
 	ui "github.com/gizak/termui/v3"
-	"github.com/kr/pty"
-	"github.com/lunixbochs/vtclean"
 	"github.com/sqshq/sampler/config"
-	"io"
 	"os"
 	"os/exec"
-	"strings"
-	"time"
 )
 
-const (
-	interactiveShellStartupTimeout  = 100 * time.Millisecond
-	interactiveShellMinAwaitTimeout = 100 * time.Millisecond
-	interactiveShellMaxAwaitTimeout = 1 * time.Second
-	interactiveShellErrorThreshold  = 10
-)
+const errorThreshold = 10
 
 type Item struct {
 	label           string
@@ -29,14 +16,9 @@ type Item struct {
 	transformScript *string
 	color           *ui.Color
 	rateMs          int
-	errorsCount      int
-	interactiveShell *InteractiveShell
-}
-
-type InteractiveShell struct {
-	Channel chan string
-	File    io.WriteCloser
-	Cmd     *exec.Cmd
+	pty             bool
+	basicShell      *BasicInteractiveShell
+	ptyShell        *PtyInteractiveShell
 }
 
 func NewItems(cfgs []config.Item, rateMs int) []*Item {
@@ -51,30 +33,32 @@ func NewItems(cfgs []config.Item, rateMs int) []*Item {
 			transformScript: i.TransformScript,
 			color:           i.Color,
 			rateMs:          rateMs,
+			pty:             *i.Pty,
 		}
 		items = append(items, item)
 	}
-
 	return items
 }
 
 func (i *Item) nextValue(variables []string) (string, error) {
 
-	if i.initScript != nil && i.interactiveShell == nil {
+	if i.initScript != nil && i.basicShell == nil && i.ptyShell == nil {
 		err := i.initInteractiveShell(variables)
 		if err != nil {
-			return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err))
+			return "", err
 		}
 	}
 
-	if i.initScript != nil {
-		return i.executeInteractiveShellCmd(variables)
+	if i.basicShell != nil {
+		return i.basicShell.execute()
+	} else if i.ptyShell != nil {
+		return i.ptyShell.execute()
 	} else {
-		return i.executeCmd(variables, i.sampleScript)
+		return i.execute(variables, i.sampleScript)
 	}
 }
 
-func (i *Item) executeCmd(variables []string, script string) (string, error) {
+func (i *Item) execute(variables []string, script string) (string, error) {
 
 	cmd := exec.Command("sh", "-c", script)
 	enrichEnvVariables(cmd, variables)
@@ -85,132 +69,31 @@ func (i *Item) executeCmd(variables []string, script string) (string, error) {
 		return "", err
 	}
 
-	result := vtclean.Clean(string(output), false)
-
-	return result, nil
+	return string(output), nil
 }
 
-func (i *Item) initInteractiveShell(variables []string) error {
-
-	cmd := exec.Command("sh", "-c", *i.initScript)
-	enrichEnvVariables(cmd, variables)
-
-	file, err := pty.Start(cmd)
-	if err != nil {
-		return err
-	}
-
-	scanner := bufio.NewScanner(file)
-	channel := make(chan string)
-
-	go func() {
-		for scanner.Scan() {
-			channel <- scanner.Text()
-		}
-	}()
-
-	i.interactiveShell = &InteractiveShell{
-		Channel: channel,
-		File:    file,
-		Cmd:     cmd,
-	}
-
-	_, err = file.Read(make([]byte, 4096))
-	if err != nil {
-		return err
-	}
-
-	time.Sleep(interactiveShellStartupTimeout)
-
-	return nil
-}
-
-func (i *Item) executeInteractiveShellCmd(variables []string) (string, error) {
-
-	_, err := io.WriteString(i.interactiveShell.File, fmt.Sprintf(" %s\n", i.sampleScript))
-	if err != nil {
-		i.errorsCount++
-		if i.errorsCount > interactiveShellErrorThreshold {
-			i.interactiveShell = nil // restart session
-			i.errorsCount = 0
-		}
-		return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err))
-	}
-
-	softTimeout := make(chan bool, 1)
-	hardTimeout := make(chan bool, 1)
-
-	go func() {
-		time.Sleep(i.getAwaitTimeout() / 4)
-		softTimeout <- true
-		time.Sleep(i.getAwaitTimeout() * 100)
-		hardTimeout <- true
-	}()
-
-	var builder strings.Builder
-	softTimeoutElapsed := false
-
-await:
-	for {
-		select {
-		case output := <-i.interactiveShell.Channel:
-			o := vtclean.Clean(output, false)
-			if len(o) > 0 && !strings.Contains(o, i.sampleScript) {
-				builder.WriteString(o)
-				builder.WriteString("\n")
-				if softTimeoutElapsed {
-					break await
-				}
-			}
-		case <-softTimeout:
-			if builder.Len() > 0 {
-				break await
+func (i *Item) initInteractiveShell(v []string) error {
+	if i.pty {
+		i.ptyShell = &PtyInteractiveShell{item: i, variables: v}
+		return i.ptyShell.init()
 	} else {
-				softTimeoutElapsed = true
+		i.basicShell = &BasicInteractiveShell{item: i, variables: v}
+		return i.basicShell.init()
 	}
-		case <-hardTimeout:
-			break await
-		}
-	}
-
-	sample := strings.TrimSpace(builder.String())
-
-	return i.transformInteractiveShellCmd(sample)
 }
 
-func (i *Item) transformInteractiveShellCmd(sample string) (string, error) {
+func (i *Item) transform(sample string) (string, error) {
 
 	if i.transformScript != nil && len(sample) > 0 {
-		return i.executeCmd([]string{"sample=" + sample}, *i.transformScript)
+		return i.execute([]string{"sample=" + sample}, *i.transformScript)
 	}
 
 	return sample, nil
 }
 
-func (i *Item) getAwaitTimeout() time.Duration {
-
-	timeout := time.Duration(i.rateMs) * time.Millisecond
-
-	if timeout > interactiveShellMaxAwaitTimeout {
-		return interactiveShellMaxAwaitTimeout
-	} else if timeout < interactiveShellMinAwaitTimeout {
-		return interactiveShellMinAwaitTimeout
-	}
-
-	return timeout
-}
-
 func enrichEnvVariables(cmd *exec.Cmd, variables []string) {
 	cmd.Env = os.Environ()
 	for _, variable := range variables {
 		cmd.Env = append(cmd.Env, variable)
 	}
 }
-
-func cleanupOutput(output string) string {
-	s := strings.TrimSpace(output)
-	if idx := strings.Index(s, "\r"); idx != -1 {
-		return s[idx+1:]
-	}
-	return s
-}

data/item_test.go

@@ -1,25 +0,0 @@
-package data
-
-import "testing"
-
-func Test_cleanupOutput(t *testing.T) {
-	type args struct {
-		output string
-	}
-	tests := []struct {
-		name string
-		args args
-		want string
-	}{
-		{"should trim everything before carriage return", args{">>\rtext"}, "text"},
-		{"should trim carriage return at the end", args{"text\r"}, "text"},
-		{"should remove tabs and spaces", args{"\t\t\ntext "}, "text"},
-	}
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			if got := cleanupOutput(tt.args.output); got != tt.want {
-				t.Errorf("cleanupOutput() = %v, want %v", got, tt.want)
-			}
-		})
-	}
-}

example-interactive-shell.yml

@@ -8,6 +8,7 @@ variables:
 textboxes:
 - title: Neo4j
   position: [[0, 0], [10, 40]]
+  pty: true
   init: $neo4jconnection
   sample: RETURN rand();
   transform: echo "$sample" | tail -n 1
@@ -17,6 +18,7 @@ textboxes:
   sample: select random();
 - title: MySQL
   position: [[19, 0], [10, 40]]
+  pty: true
   init: $mysqlconnection
   sample: select rand();
 - title: MongoDB
@@ -26,5 +28,6 @@ textboxes:
   sample: sleep(3000);Date.now();
 - title: SSH
   position: [[39, 0], [41, 40]]
+  pty: true
   init: $sshconnection
   sample: top