unkin/node-lookup
3
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 20 Wiki Activity

Query piped nodes when counting with -R/-F #22

Merged
benvin merged 2 commits from benvin/count-stdin-query into main 2026-10-13 21:04:18 +11:00
Conversation 2 Commits 2 Files Changed 2
+185 -22
2 changed files with 185 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.go
+35 -22
View File
@@ -357,15 +357,38 @@ func run(cfg config, nodeName, factName, match string, showRole, partial, invers
return nil
}
if (nodeOnly || valueOnly || count || ansible) && !showRole && factName == "" {
return fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A")
hasQuery := showRole || factName != ""
// Bare -C counts already-looked-up piped lines; only it needs stdin to validate.
bareCount := count && !hasQuery && !nodeOnly && !valueOnly && !ansible && !jsonMode && nodeName == ""
countErr := fmt.Errorf("-R or -F must be used with -1, -2, -C, or -A")
if (nodeOnly || valueOnly || (count && !bareCount) || ansible) && !hasQuery {
return countErr
}
if (match != "" || partial || inverse) && !showRole && factName == "" {
if (match != "" || partial || inverse) && !hasQuery {
return fmt.Errorf("-R or -F must be used with -m, -p, or -i")
}
var reader *bufio.Reader
piped := false
if nodeName == "" {
reader, piped = stdinReader(os.Stdin)
}
if bareCount {
if !piped {
return countErr
}
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 +400,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 +427,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.
main_test.go
+150
View File
@@ -8,7 +8,9 @@ import (
"os"
"path/filepath"
"strings"
"sync/atomic"
"testing"
"time"
"gopkg.in/yaml.v3"
)
@@ -872,3 +874,151 @@ 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)
}
}
func TestRun_CountPipedWithFactQueriesEachNode(t *testing.T) {
withStdin(t, pipedLookup)
var queried atomic.Int32
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if !strings.Contains(r.URL.Query().Get("query"), "osfamily") {
t.Errorf("expected osfamily query, got %q", r.URL.Query().Get("query"))
}
queried.Add(1)
_ = json.NewEncoder(w).Encode([]fact{{Certname: "h", Name: "osfamily", Value: rawJSON("RedHat")}})
}))
t.Cleanup(srv.Close)
cfg := config{PuppetDBURL: srv.URL, RoleFact: "enc_role"}
out := strings.TrimSpace(captureStdout(t, func() {
if err := run(cfg, "", "osfamily", "", false, false, false, false, false, true, false, false, false); err != nil {
t.Fatalf("run returned error: %v", err)
}
}))
if out != "2: RedHat" {
t.Fatalf("expected fact 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_CountPipedWithOtherOutputFlagsRequiresRoleOrFact(t *testing.T) {
for name, mutate := range map[string]func(*runArgs){
"json": func(a *runArgs) { a.jsonMode = true },
"nodeOnly": func(a *runArgs) { a.nodeOnly = true },
"valueOnly": func(a *runArgs) { a.valueOnly = true },
} {
t.Run(name, func(t *testing.T) {
withStdin(t, pipedLookup)
a := runArgs{cfg: config{PuppetDBURL: "http://unused"}, count: true}
mutate(&a)
var err error
out := captureStdout(t, func() {
err = run(a.cfg, "", "", "", false, false, false, a.nodeOnly, a.valueOnly, true, false, a.jsonMode, false)
})
if err == nil || !strings.Contains(err.Error(), "-R or -F") {
t.Fatalf("expected -R/-F requirement error, got: %v (output %q)", err, out)
}
})
}
}
func TestRun_InvalidFlagsErrorWithoutReadingOpenPipe(t *testing.T) {
for name, a := range map[string]runArgs{
"match": {match: "foo"},
"countJSON": {count: true, jsonMode: true},
} {
t.Run(name, func(t *testing.T) {
r, w, err := os.Pipe()
if err != nil {
t.Fatal(err)
}
old := os.Stdin
os.Stdin = r
t.Cleanup(func() { os.Stdin = old; _ = w.Close(); _ = r.Close() })
done := make(chan error, 1)
go func() {
done <- run(config{PuppetDBURL: "http://unused"}, "", "", a.match, false, false, false, false, false, a.count, false, a.jsonMode, false)
}()
select {
case err := <-done:
if err == nil || !strings.Contains(err.Error(), "-R or -F") {
t.Fatalf("expected -R/-F requirement error, got: %v", err)
}
case <-time.After(2 * time.Second):
t.Fatal("run blocked on an open, empty stdin pipe before validating flags")
}
})
}
}