From 4d5fa35642f04b647505517bcb1b4bbdae84f4a0 Mon Sep 17 00:00:00 2001 From: sqshq Date: Mon, 13 May 2019 23:24:29 -0400 Subject: [PATCH] switch interactive shell to PTY to support cases like python or neo4j --- component/statusbar.go | 7 ++- config-interactive-shell.yml | 29 ++++++++++++ console/console.go | 1 + data/item.go | 85 ++++++++++++++++-------------------- data/item_test.go | 25 +++++++++++ data/sampler.go | 2 +- go.mod | 1 + go.sum | 2 + main.go | 2 + 9 files changed, 104 insertions(+), 50 deletions(-) create mode 100644 config-interactive-shell.yml create mode 100644 data/item_test.go diff --git a/component/statusbar.go b/component/statusbar.go index 99f9ed5..d8eb989 100644 --- a/component/statusbar.go +++ b/component/statusbar.go @@ -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 { diff --git a/config-interactive-shell.yml b/config-interactive-shell.yml new file mode 100644 index 0000000..7b34d61 --- /dev/null +++ b/config-interactive-shell.yml @@ -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}' diff --git a/console/console.go b/console/console.go index 9d7ab29..fd25f9b 100644 --- a/console/console.go +++ b/console/console.go @@ -14,6 +14,7 @@ const ( RowsCount = 40 AppTitle = "sampler" AppVersion = "0.9.0" + AppLicenseWarning = "UNLICENSED. FOR PERSONAL USE ONLY. VISIT SAMPLER.DEV" ) const ( diff --git a/data/item.go b/data/item.go index c2cd5ca..d4a7304 100644 --- a/data/item.go +++ b/data/item.go @@ -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 +} diff --git a/data/item_test.go b/data/item_test.go new file mode 100644 index 0000000..88cce01 --- /dev/null +++ b/data/item_test.go @@ -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) + } + }) + } +} diff --git a/data/sampler.go b/data/sampler.go index c3e5b25..c19f904 100644 --- a/data/sampler.go +++ b/data/sampler.go @@ -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) } } } diff --git a/go.mod b/go.mod index bdb6652..b087f31 100644 --- a/go.mod +++ b/go.mod @@ -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 diff --git a/go.sum b/go.sum index 3e0d9cb..221e9f1 100644 --- a/go.sum +++ b/go.sum @@ -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= diff --git a/main.go b/main.go index aedbc44..6d925a1 100644 --- a/main.go +++ b/main.go @@ -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() {