From 8b9ac394a058daf6022308c3f310e25aad53d917 Mon Sep 17 00:00:00 2001 From: unkin-agent Date: Tue, 13 Oct 2026 20:38:55 +1100 Subject: [PATCH] Query piped nodes for -R/-F with -C; count input only for bare -C --- main.go | 47 +++++++++++++++++++--------------- main_test.go | 72 ++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 99 insertions(+), 20 deletions(-) diff --git a/main.go b/main.go index 2b22e86..af30b01 100644 --- a/main.go +++ b/main.go @@ -357,15 +357,32 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers return nil } - if (nodeOnly || valueOnly || count || ansible) && !showRole && factName == "" { + var reader *bufio.Reader + piped := false + if nodeName == "" { + reader, piped = stdinReader(os.Stdin) + } + // Bare -C on piped input counts the already-looked-up lines instead of querying. + passthrough := count && piped && !showRole && factName == "" + + if (nodeOnly || valueOnly || (count && !passthrough) || ansible) && !showRole && factName == "" { return fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A") } if (match != "" || partial || inverse) && !showRole && factName == "" { return fmt.Errorf("-R or -F must be used with -m, -p, or -i") } + if passthrough { + var lines []string + scanner := bufio.NewScanner(reader) + for scanner.Scan() { + lines = append(lines, scanner.Text()) + } + fmt.Println(strings.Join(countResults(lines), "\n")) + return nil + } + var collected []fact - var stdinLines []string doQuery := func(node string) error { query := buildQuery(node, factName, match, cfg.RoleFact, showRole, partial, inverse) @@ -377,21 +394,15 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers return nil } - if reader, ok := stdinReader(os.Stdin); ok && nodeName == "" { + if piped { scanner := bufio.NewScanner(reader) - if count { - for scanner.Scan() { - stdinLines = append(stdinLines, scanner.Text()) + for scanner.Scan() { + fields := strings.Fields(scanner.Text()) + if len(fields) == 0 { + continue } - } else { - for scanner.Scan() { - fields := strings.Fields(scanner.Text()) - if len(fields) == 0 { - continue - } - if err := doQuery(qualifyNode(fields[0], cfg.Domain)); err != nil { - fmt.Fprintln(os.Stderr, "error:", err) - } + if err := doQuery(qualifyNode(fields[0], cfg.Domain)); err != nil { + fmt.Fprintln(os.Stderr, "error:", err) } } } else { @@ -410,11 +421,7 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers _ = enc.Encode(factsByHost(collected, factName, cfg.RoleFact, showRole)) case count: - values := stdinLines - if len(values) == 0 { - values = returnData - } - fmt.Println(strings.Join(countResults(values), "\n")) + fmt.Println(strings.Join(countResults(returnData), "\n")) case ansible: // Attach each host's queried fact(s) as inventory host vars, e.g. diff --git a/main_test.go b/main_test.go index c65f7b5..a595bcd 100644 --- a/main_test.go +++ b/main_test.go @@ -8,6 +8,7 @@ import ( "os" "path/filepath" "strings" + "sync/atomic" "testing" "gopkg.in/yaml.v3" @@ -872,3 +873,74 @@ func TestRun_AllFacts_PrintsSortedByName(t *testing.T) { t.Fatalf("facts not sorted by name: %v", lines) } } + +// ---- run: count with piped stdin -------------------------------------------- + +func withStdin(t *testing.T, content string) { + t.Helper() + path := filepath.Join(t.TempDir(), "stdin") + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + f, err := os.Open(path) + if err != nil { + t.Fatal(err) + } + old := os.Stdin + os.Stdin = f + t.Cleanup(func() { os.Stdin = old; _ = f.Close() }) +} + +const pipedLookup = "hosta.main.unkin.net puppetdbapi.service.consul\nhostb.main.unkin.net puppetdbapi.service.consul\n" + +func TestRun_CountPipedWithRoleQueriesEachNode(t *testing.T) { + withStdin(t, pipedLookup) + var queried atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + queried.Add(1) + _ = json.NewEncoder(w).Encode([]fact{{Certname: "h", Name: "enc_role", Value: rawJSON("roles::vault")}}) + })) + t.Cleanup(srv.Close) + cfg := config{PuppetDBURL: srv.URL, RoleFact: "enc_role"} + + out := strings.TrimSpace(captureStdout(t, func() { + if err := run(cfg, "", "", "", true, false, false, false, false, true, false, false, false); err != nil { + t.Fatalf("run returned error: %v", err) + } + })) + if out != "2: roles::vault" { + t.Fatalf("expected role counts of piped nodes, got: %q", out) + } + if n := queried.Load(); n != 2 { + t.Fatalf("expected one query per piped node, got %d", n) + } +} + +func TestRun_CountPipedWithoutQueryCountsInput(t *testing.T) { + withStdin(t, pipedLookup+"hostc.main.unkin.net other\n") + cfg := config{PuppetDBURL: "http://unused", RoleFact: "enc_role"} + + out := strings.TrimSpace(captureStdout(t, func() { + if err := run(cfg, "", "", "", false, false, false, false, false, true, false, false, false); err != nil { + t.Fatalf("run returned error: %v", err) + } + })) + if out != "1: other\n2: puppetdbapi.service.consul" { + t.Fatalf("expected counts of piped values, got: %q", out) + } +} + +func TestRun_CountWithoutStdinRequiresRoleOrFact(t *testing.T) { + devnull, err := os.Open(os.DevNull) + if err != nil { + t.Fatal(err) + } + old := os.Stdin + os.Stdin = devnull + t.Cleanup(func() { os.Stdin = old; _ = devnull.Close() }) + + err = run(config{PuppetDBURL: "http://unused"}, "", "", "", false, false, false, false, false, true, false, false, false) + if err == nil || !strings.Contains(err.Error(), "-R or -F") { + t.Fatalf("expected -R/-F requirement error, got: %v", err) + } +}