From 39dc8d83312b48460237e51b19d2971949067b0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Martin=20Odstr=C4=8Dil=C3=ADk?=
 <martin.odstrcilik@lablabs.io>
Date: Tue, 8 Oct 2024 08:44:09 +0200
Subject: [PATCH] fix(metrics/script): patch pipe processing

---
 go.mod                           |  1 +
 go.sum                           |  4 +++
 internal/scrape/script/metric.go |  5 ++--
 internal/scrape/script/script.go | 45 +++++++++++---------------------
 4 files changed, 22 insertions(+), 33 deletions(-)

diff --git a/go.mod b/go.mod
index 2adff70..c94efab 100644
--- a/go.mod
+++ b/go.mod
@@ -6,6 +6,7 @@ require (
 	github.com/aws/aws-sdk-go-v2 v1.26.0
 	github.com/aws/aws-sdk-go-v2/config v1.27.4
 	github.com/aws/aws-sdk-go-v2/service/servicequotas v1.21.1
+	github.com/go-cmd/cmd v1.4.3
 	github.com/itchyny/gojq v0.12.14
 	github.com/jessevdk/go-flags v1.5.0
 	github.com/prometheus/client_golang v1.18.0
diff --git a/go.sum b/go.sum
index 9267523..3321a40 100644
--- a/go.sum
+++ b/go.sum
@@ -34,6 +34,10 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-cmd/cmd v1.4.3 h1:6y3G+3UqPerXvPcXvj+5QNPHT02BUw7p6PsqRxLNA7Y=
+github.com/go-cmd/cmd v1.4.3/go.mod h1:u3hxg/ry+D5kwh8WvUkHLAMe2zQCaXd00t35WfQaOFk=
+github.com/go-test/deep v1.1.0 h1:WOcxcdHcvdgThNXjw0t76K42FXTU7HpNQWHpA2HHNlg=
+github.com/go-test/deep v1.1.0/go.mod h1:5C2ZWiW0ErCdrYzpqxLbTX7MG14M9iiw8DgHncVwcsE=
 github.com/golang/protobuf v1.5.0/go.mod h1:FsONVRAS9T7sI+LIUmWTfcYkHO4aIWwzhcaSAoJOfIk=
 github.com/google/go-cmp v0.5.5/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
 github.com/google/go-cmp v0.5.9 h1:O2Tfq5qg4qc4AmwVlvv0oLiVAGB7enBSJ2x2DqQFi38=
diff --git a/internal/scrape/script/metric.go b/internal/scrape/script/metric.go
index 435b760..7594902 100644
--- a/internal/scrape/script/metric.go
+++ b/internal/scrape/script/metric.go
@@ -3,7 +3,6 @@ package script
 import (
 	"bufio"
 	"fmt"
-	"io"
 	"regexp"
 	"strconv"
 	"strings"
@@ -11,10 +10,10 @@ import (
 
 const metricRegex = `,`
 
-func ParseStdout(r io.Reader) ([]Data, error) {
+func ParseStdout(r []string) ([]Data, error) {
 	out := make([]Data, 0)
 	pr := NewParser()
-	scanner := bufio.NewScanner(r)
+	scanner := bufio.NewScanner(strings.NewReader(strings.Join(r, "\n")))
 	for scanner.Scan() {
 		mt, err := pr.ParseMetric(scanner.Text())
 		if err != nil {
diff --git a/internal/scrape/script/script.go b/internal/scrape/script/script.go
index ef98ace..57b272e 100644
--- a/internal/scrape/script/script.go
+++ b/internal/scrape/script/script.go
@@ -1,13 +1,10 @@
 package script
 
 import (
-	"bytes"
 	"context"
 	"fmt"
-	"io"
+	"github.com/go-cmd/cmd"
 	"os"
-	"os/exec"
-	"strings"
 )
 
 type Data struct {
@@ -27,37 +24,25 @@ func (d Data) LabelNames() []string {
 }
 
 func Run(ctx context.Context, cfg Config) ([]Data, error) {
-	cmd := exec.CommandContext(ctx, "bash", "-c", cfg.Script)
-	envs := make([]string, 0)
-	envs = append(envs, os.Environ()...)
-	envs = append(envs, cfg.FormatEnvs()...)
-	cmd.Env = envs
-	stderr, err := cmd.StderrPipe()
-	if err != nil {
-		return nil, err
-	}
-	stdout := &bytes.Buffer{}
-	cmd.Stdout = stdout
-	err = cmd.Run()
-	if err != nil {
-		errString, err := errorString(stderr)
+	c := cmd.NewCmd("bash", "-c", cfg.Script)
+	c.Env = append(c.Env, os.Environ()...)
+	c.Env = append(c.Env, cfg.FormatEnvs()...)
+
+	<-c.Start()
+	select {
+	case <-ctx.Done():
+		err := c.Stop()
 		if err != nil {
-			return nil, err
+			return nil, fmt.Errorf("script error: %w, std err: %s", err, c.Status().Stderr)
 		}
-		return nil, fmt.Errorf("script error: %w, std err: %s", err, errString)
 	}
-	data, err := ParseStdout(stdout)
+	err := c.Status().Error
 	if err != nil {
-		return nil, fmt.Errorf("unable to parse response from command: %v", cfg.Script)
+		return nil, fmt.Errorf("script error: %w, std err: %s", err, c.Status().Stderr)
 	}
-	return data, nil
-}
-
-func errorString(r io.Reader) (string, error) {
-	b := strings.Builder{}
-	_, err := io.Copy(&b, r)
+	data, err := ParseStdout(c.Status().Stdout)
 	if err != nil {
-		return "", err
+		return nil, fmt.Errorf("unable to parse response from command: %v", cfg.Script)
 	}
-	return b.String(), nil
+	return data, nil
 }