switch interactive shell to PTY to support cases like python or neo4j

This commit is contained in:
sqshq 2019-05-13 23:24:29 -04:00
parent 809d361ca7
commit 4d5fa35642
9 changed files with 104 additions and 50 deletions

View File

@ -34,7 +34,12 @@ func NewStatusLine(configFileName string, palette console.Palette) *StatusBar {
func (s *StatusBar) Draw(buffer *ui.Buffer) {
buffer.Fill(ui.NewCell(' ', ui.NewStyle(console.ColorClear, console.MenuColorBackground)), s.GetRect())
buffer.SetString(fmt.Sprintf(" %s %s @ %s", console.AppTitle, console.AppVersion, s.configFileName), ui.NewStyle(console.MenuColorText, console.MenuColorBackground), s.Min)
if false { // TODO check license
buffer.SetString(fmt.Sprintf(" %s", console.AppLicenseWarning), ui.NewStyle(console.MenuColorText, console.MenuColorBackground), s.Min)
} else {
buffer.SetString(fmt.Sprintf(" %s %s @ %s", console.AppTitle, console.AppVersion, s.configFileName), ui.NewStyle(console.MenuColorText, console.MenuColorBackground), s.Min)
}
indent := bindingsIndent
for _, binding := range s.keyBindings {

View File

@ -0,0 +1,29 @@
variables:
PGPASSWORD: fred
mongoconnection: mongo --quiet --host=localhost blog
mysqlconnection: mysql -u root -s --database mysql --skip-column-names
neo4jconnection: cypher-shell -u neo4j -p 121314 --format plain
postgresconnection: psql -h localhost -U postgres --no-align --tuples-only
sshconnection: ssh -i ~/sqshq.pem ec2-user@3.215.108.82
textboxes:
- title: Neo4j
position: [[0, 0], [13, 40]]
init: $neo4jconnection
sample: match (n) return count(n);
transform: echo "$sample" | tail -n 1
- title: Postgres
position: [[13, 0], [14, 40]]
init: $postgresconnection
sample: select random();
- title: MySQL
position: [[27, 0], [14, 40]]
init: $mysqlconnection
sample: select rand();
- title: MongoDB
position: [[41, 0], [13, 40]]
init: $mongoconnection
sample: db.getCollection('posts').find({status:'ACTIVE'}).itcount()
- title: SSH
position: [[54, 0], [13, 40]]
init: $sshconnection
sample: ps -A -o %cpu | awk '{s+=$1} END {print s}'

View File

@ -14,6 +14,7 @@ const (
RowsCount = 40
AppTitle = "sampler"
AppVersion = "0.9.0"
AppLicenseWarning = "UNLICENSED. FOR PERSONAL USE ONLY. VISIT SAMPLER.DEV"
)
const (

View File

@ -3,7 +3,9 @@ package data
import (
"bufio"
"errors"
"fmt"
ui "github.com/gizak/termui/v3"
"github.com/kr/pty"
"github.com/sqshq/sampler/config"
"io"
"os"
@ -12,6 +14,8 @@ import (
"time"
)
const interactiveShellStartupTimeout = time.Second
type Item struct {
Label string
SampleScript string
@ -23,10 +27,9 @@ type Item struct {
}
type InteractiveShell struct {
StdoutCh chan string
StderrCh chan string
Stdin io.WriteCloser
Cmd *exec.Cmd
Channel chan string
File io.WriteCloser
Cmd *exec.Cmd
}
func NewItems(cfgs []config.Item, rateMs int) []*Item {
@ -53,7 +56,7 @@ func (i *Item) nextValue(variables []string) (string, error) {
if i.InitScript != nil && i.InteractiveShell == nil {
err := i.initInteractiveShell(variables)
if err != nil {
return "", err
return "", errors.New(fmt.Sprintf("Failed to init interactive shell: %s", err))
}
}
@ -83,84 +86,62 @@ func (i *Item) initInteractiveShell(variables []string) error {
cmd := exec.Command("sh", "-c", *i.InitScript)
enrichEnvVariables(cmd, variables)
stdout, err := cmd.StdoutPipe()
file, err := pty.Start(cmd)
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)
scanner := bufio.NewScanner(file)
channel := make(chan string)
go func() {
for stdoutScanner.Scan() {
stdoutCh <- stdoutScanner.Text()
stderrCh <- stderrScanner.Text()
for scanner.Scan() {
channel <- scanner.Text()
}
}()
i.InteractiveShell = &InteractiveShell{
StdoutCh: stdoutCh,
StderrCh: stderrCh,
Stdin: stdin,
Cmd: cmd,
Channel: channel,
File: file,
Cmd: cmd,
}
err = cmd.Start()
_, 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.Stdin, i.SampleScript+"\n")
_, err := io.WriteString(i.InteractiveShell.File, fmt.Sprintf(" %s\n", i.SampleScript))
if err != nil {
return "", err
return "", errors.New(fmt.Sprintf("Failed to execute interactive shell cmd: %s", err))
}
timeout := make(chan bool, 1)
go func() {
time.Sleep(time.Duration(i.RateMs / 2))
time.Sleep(time.Duration(i.RateMs))
timeout <- true
}()
var resultText strings.Builder
var errorText strings.Builder
var outputText 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 output := <-i.InteractiveShell.Channel:
if !strings.Contains(output, i.SampleScript) && len(output) > 0 {
outputText.WriteString(output)
outputText.WriteString("\n")
}
case <-timeout:
if errorText.Len() > 0 {
return "", errors.New(errorText.String())
} else {
return i.transformInteractiveShellCmd(resultText.String())
}
sample := cleanupOutput(outputText.String())
return i.transformInteractiveShellCmd(sample)
}
}
}
@ -180,3 +161,11 @@ func enrichEnvVariables(cmd *exec.Cmd, variables []string) {
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
}

25
data/item_test.go Normal file
View File

@ -0,0 +1,25 @@
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)
}
})
}
}

View File

@ -39,7 +39,7 @@ func NewSampler(consumer *Consumer, items []*Item, triggers []*Trigger, options
select {
case sample := <-sampler.triggersChannel:
for _, t := range sampler.triggers {
t.Execute(sample)
go t.Execute(sample)
}
}
}

1
go.mod
View File

@ -5,6 +5,7 @@ require (
github.com/hajimehoshi/go-mp3 v0.1.1
github.com/hajimehoshi/oto v0.1.1
github.com/jessevdk/go-flags v1.4.0
github.com/kr/pty v1.1.4
github.com/mattn/go-runewidth v0.0.4
github.com/mbndr/figlet4go v0.0.0-20190224160619-d6cef5b186ea
github.com/mitchellh/go-wordwrap v1.0.0 // indirect

2
go.sum
View File

@ -12,6 +12,8 @@ github.com/hajimehoshi/oto v0.1.1 h1:EG+WxxeAfde1mI0adhLYvGbKgDCxm7bCTd6g+JIA6vI
github.com/hajimehoshi/oto v0.1.1/go.mod h1:hUiLWeBQnbDu4pZsAhOnGqMI1ZGibS6e2qhQdfpwz04=
github.com/jessevdk/go-flags v1.4.0 h1:4IU2WS7AumrZ/40jfhf4QVDMsQwqA7VEHozFRrGARJA=
github.com/jessevdk/go-flags v1.4.0/go.mod h1:4FA24M0QyGHXBuZZK/XkWh8h0e1EYbRYJSGM75WSRxI=
github.com/kr/pty v1.1.4 h1:5Myjjh3JY/NaAi4IsUbHADytDyl1VE1Y9PXDlL+P/VQ=
github.com/kr/pty v1.1.4/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
github.com/mattn/go-runewidth v0.0.2/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=
github.com/mattn/go-runewidth v0.0.4 h1:2BvfKmzob6Bmd4YsL0zygOqfdFnK7GR4QL06Do4/p7Y=
github.com/mattn/go-runewidth v0.0.4/go.mod h1:LwmH8dsx7+W8Uxz3IHJYH5QSwggIsqBzpuz5H//U1FU=

View File

@ -15,6 +15,7 @@ import (
"github.com/sqshq/sampler/console"
"github.com/sqshq/sampler/data"
"github.com/sqshq/sampler/event"
"time"
)
type Starter struct {
@ -30,6 +31,7 @@ func (s *Starter) start(drawable ui.Drawable, consumer *data.Consumer, component
items := data.NewItems(itemsConfig, *componentConfig.RateMs)
data.NewSampler(consumer, items, triggers, s.opt, s.cfg.Variables, *componentConfig.RateMs)
s.lout.AddComponent(cpt)
time.Sleep(100 * time.Millisecond) // desync coroutines
}
func main() {