cmd.go

package main

import (
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"sync"
	"time"
)

// ============================================================
//  Log
// ============================================================

// Log writes timestamped messages and serves as io.Writer for command output.
type Log struct {
	mu sync.Mutex
	w  io.Writer
}

// NewLog creates a Log writing to w.
func NewLog(w io.Writer) *Log { return &Log{w: w} }

// Printf writes a timestamped log line.
func (l *Log) Printf(format string, args ...any) {
	l.mu.Lock()
	defer l.mu.Unlock()
	fmt.Fprintf(l.w, "%s  %s\n", time.Now().Format("15:04:05"), fmt.Sprintf(format, args...))
}

// Write passes raw bytes through (for command output).
func (l *Log) Write(p []byte) (int, error) {
	l.mu.Lock()
	defer l.mu.Unlock()
	return l.w.Write(p)
}

// ============================================================
//  Env
// ============================================================

// Env is an environment variable list with a builder API.
type Env []string

// BaseEnv returns a copy of the current process environment.
func BaseEnv() Env { return Env(os.Environ()) }

// Set adds or replaces an environment variable.
func (e Env) Set(key, val string) Env {
	prefix := key + "="
	for i, entry := range e {
		if strings.HasPrefix(entry, prefix) {
			e[i] = prefix + val
			return e
		}
	}
	return append(e, prefix+val)
}

// SetInt is a convenience for integer values.
func (e Env) SetInt(key string, val int) Env {
	return e.Set(key, strconv.Itoa(val))
}

// Merge adds all entries from a map.
func (e Env) Merge(m map[string]string) Env {
	for k, v := range m {
		e = e.Set(k, v)
	}
	return e
}

// ============================================================
//  Command helpers
// ============================================================

// runCmd executes a command, logging its output.
func runCmd(log *Log, env Env, name string, args ...string) error {
	log.Printf("    $ %s %s", name, strings.Join(args, " "))
	cmd := exec.Command(name, args...)
	if env != nil {
		cmd.Env = []string(env)
	}
	cmd.Stdout = log
	cmd.Stderr = log
	return cmd.Run()
}

// stropSQL runs a SQL statement against the given database.
// Uses native clients (psql/mysql) because stroppy's inline SQL
// doesn't work reliably across all drivers yet.
// TODO: switch back to stroppy when it supports inline SQL everywhere.
func stropSQL(log *Log, db Database, dbname, sql string) error {
	log.Printf("    sql [%s]: %s", dbname, sql)
	var cmd *exec.Cmd
	host, port, _ := strings.Cut(db.Host, ":")
	switch db.Driver {
	case Postgres:
		cmd = exec.Command("psql", db.URL(dbname), "-c", sql)
	case MySQL:
		args := []string{"-h", host, "-P", port, "-u", db.User}
		if db.Password != "" {
			args = append(args, "-p"+db.Password)
		}
		if dbname != "" {
			args = append(args, "-D", dbname)
		}
		args = append(args, "-e", sql)
		cmd = exec.Command("mysql", args...)
	default:
		return fmt.Errorf("stropSQL: unsupported driver %s", db.Driver)
	}
	cmd.Stdout = log
	cmd.Stderr = log
	return cmd.Run()
}

// stropLoad runs stroppy data loading (all steps except workload).
func stropLoad(log *Log, db Database, url string, scale int) error {
	log.Printf("    loading data (scale=%d)", scale)
	args := []string{"run", db.Preset}
	if db.DriverArg != "" {
		args = append(args, db.DriverArg)
	}
	args = append(args, "--no-steps", "workload", "--", "-q")

	env := BaseEnv().
		Set("DRIVER_URL", url).
		SetInt("SCALE_FACTOR", scale).
		SetInt("POOL_SIZE", LoadPoolSize)

	cmd := exec.Command("stroppy", args...)
	cmd.Env = []string(env)
	cmd.Stdout = log
	cmd.Stderr = log
	return cmd.Run()
}