feat: add stdin support and retry logic for all search commandsHEADv1.1.1v0.1.14master

1 files changed, 98 insertions, 0 deletions

diff --git a/cmd/stdin.go b/cmd/stdin.go
new file mode 100644
index 0000000..e9240db
--- /dev/null
+++ b/cmd/stdin.go
@@ -0,0 +1,98 @@
+package cmd
+
+import (
+	"bufio"
+	"fmt"
+	"os"
+	"strings"
+
+	"git.db.org.ai/dborg/internal/utils"
+	"github.com/spf13/cobra"
+)
+
+func stdinPiped() bool {
+	fi, err := os.Stdin.Stat()
+	if err != nil {
+		return false
+	}
+	return (fi.Mode() & os.ModeCharDevice) == 0
+}
+
+func readStdinLines() ([]string, error) {
+	var lines []string
+	scanner := bufio.NewScanner(os.Stdin)
+	scanner.Buffer(make([]byte, 64*1024), 1024*1024)
+	for scanner.Scan() {
+		line := strings.TrimSpace(scanner.Text())
+		if line == "" || strings.HasPrefix(line, "#") {
+			continue
+		}
+		lines = append(lines, line)
+	}
+	if err := scanner.Err(); err != nil {
+		return nil, fmt.Errorf("failed to read stdin: %w", err)
+	}
+	return lines, nil
+}
+
+func argsOrStdin(n int) cobra.PositionalArgs {
+	return func(cmd *cobra.Command, args []string) error {
+		if len(args) == n {
+			return nil
+		}
+		if len(args) == 0 && stdinPiped() {
+			return nil
+		}
+		return fmt.Errorf("accepts %d arg(s), received %d", n, len(args))
+	}
+}
+
+func resolveQueries(args []string) ([]string, error) {
+	if len(args) > 0 {
+		return args, nil
+	}
+	if !stdinPiped() {
+		return nil, fmt.Errorf("no query provided")
+	}
+	lines, err := readStdinLines()
+	if err != nil {
+		return nil, err
+	}
+	if len(lines) == 0 {
+		return nil, fmt.Errorf("no queries read from stdin")
+	}
+	return lines, nil
+}
+
+func forEachQuery(args []string, fn func(query string) error) error {
+	queries, err := resolveQueries(args)
+	if err != nil {
+		return err
+	}
+
+	multi := len(queries) > 1
+	var firstErr error
+	for i, q := range queries {
+		if multi {
+			printQuerySeparator(i, q)
+		}
+		if err := fn(q); err != nil {
+			fmt.Fprintf(os.Stderr, "Error for %q: %v\n", q, err)
+			if firstErr == nil {
+				firstErr = err
+			}
+		}
+	}
+	return firstErr
+}
+
+func printQuerySeparator(i int, q string) {
+	if utils.IsTerminal() {
+		if i > 0 {
+			fmt.Println()
+		}
+		fmt.Printf("\033[1;36m━━━ %s ━━━\033[0m\n", q)
+	} else {
+		fmt.Printf("--- %s ---\n", q)
+	}
+}