423 lines
11 KiB
Go
423 lines
11 KiB
Go
package main
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
"sort"
|
|
|
|
"gopkg.in/yaml.v3"
|
|
)
|
|
|
|
// exit codes
|
|
const (
|
|
exitOK = 0
|
|
exitZeroMatch = 1
|
|
exitError = 2
|
|
)
|
|
|
|
// queryResult is the stable JSON shape for selector queries.
|
|
type queryResult struct {
|
|
Env string `json:"env"`
|
|
Kind string `json:"kind"`
|
|
Selector string `json:"selector"`
|
|
Matches []match `json:"matches"`
|
|
}
|
|
|
|
// envProjection is the stable JSON shape when no selector is given.
|
|
type envProjection struct {
|
|
Env string `json:"test_env"`
|
|
Profile string `json:"profile"`
|
|
LastUpdatedAt string `json:"last_updated_at"`
|
|
Source interface{} `json:"source,omitempty"`
|
|
Edge interface{} `json:"edge,omitempty"`
|
|
Build interface{} `json:"build,omitempty"`
|
|
}
|
|
|
|
// match is a single path/value pair in a selector query result.
|
|
type match struct {
|
|
Path string `json:"path"`
|
|
Value interface{} `json:"value"`
|
|
}
|
|
|
|
// flagState holds parsed CLI flags.
|
|
type flagState struct {
|
|
env string
|
|
model string
|
|
node string
|
|
provider string
|
|
help bool
|
|
}
|
|
|
|
func main() {
|
|
os.Exit(run(os.Args[1:], os.Stdout, os.Stderr, loadInventory))
|
|
}
|
|
|
|
// run executes the query flow and returns the process exit code.
|
|
func run(args []string, stdout io.Writer, stderr io.Writer, loader func(string) (map[string]interface{}, error)) int {
|
|
fs, err := parseFlags(args)
|
|
if err != nil {
|
|
fmt.Fprintln(stderr, err)
|
|
return exitError
|
|
}
|
|
|
|
if fs.help {
|
|
printUsage(stdout)
|
|
return exitOK
|
|
}
|
|
|
|
data, err := loader(fs.env)
|
|
if err != nil {
|
|
fmt.Fprintln(stderr, err)
|
|
return exitError
|
|
}
|
|
|
|
var selector string
|
|
var kind string
|
|
switch {
|
|
case fs.node != "":
|
|
selector = fs.node
|
|
kind = "node"
|
|
case fs.provider != "":
|
|
selector = fs.provider
|
|
kind = "provider"
|
|
case fs.model != "":
|
|
selector = fs.model
|
|
kind = "model"
|
|
}
|
|
|
|
if err := validateInventory(data, fs.env, kind); err != nil {
|
|
fmt.Fprintln(stderr, err)
|
|
return exitError
|
|
}
|
|
|
|
if selector == "" {
|
|
proj := buildEnvProjection(data)
|
|
if err := encodeJSON(stdout, proj); err != nil {
|
|
fmt.Fprintf(stderr, "json encode error: %v\n", err)
|
|
return exitError
|
|
}
|
|
return exitOK
|
|
}
|
|
|
|
matches := queryInventory(data, kind, selector)
|
|
if len(matches) == 0 {
|
|
fmt.Fprintf(stderr, "no match for %s=%q\n", kind, selector)
|
|
return exitZeroMatch
|
|
}
|
|
|
|
result := queryResult{
|
|
Env: fs.env,
|
|
Kind: kind,
|
|
Selector: selector,
|
|
Matches: matches,
|
|
}
|
|
if err := encodeJSON(stdout, result); err != nil {
|
|
fmt.Fprintf(stderr, "json encode error: %v\n", err)
|
|
return exitError
|
|
}
|
|
return exitOK
|
|
}
|
|
|
|
// parseFlags parses CLI flags. --env is required (dev|dev-corp). At most one of --model, --node, --provider is allowed.
|
|
func parseFlags(args []string) (flagState, error) {
|
|
var fs flagState
|
|
var selectors int
|
|
|
|
i := 0
|
|
for i < len(args) {
|
|
switch args[i] {
|
|
case "--env":
|
|
i++
|
|
if i >= len(args) {
|
|
return fs, fmt.Errorf("--env requires a value")
|
|
}
|
|
fs.env = args[i]
|
|
case "--model":
|
|
i++
|
|
if i >= len(args) {
|
|
return fs, fmt.Errorf("--model requires a value")
|
|
}
|
|
fs.model = args[i]
|
|
selectors++
|
|
case "--node":
|
|
i++
|
|
if i >= len(args) {
|
|
return fs, fmt.Errorf("--node requires a value")
|
|
}
|
|
fs.node = args[i]
|
|
selectors++
|
|
case "--provider":
|
|
i++
|
|
if i >= len(args) {
|
|
return fs, fmt.Errorf("--provider requires a value")
|
|
}
|
|
fs.provider = args[i]
|
|
selectors++
|
|
case "--help", "-h":
|
|
fs.help = true
|
|
default:
|
|
return fs, fmt.Errorf("unknown flag: %s", args[i])
|
|
}
|
|
i++
|
|
}
|
|
|
|
if !fs.help {
|
|
if fs.env == "" {
|
|
return fs, fmt.Errorf("--env is required (dev|dev-corp)")
|
|
}
|
|
if fs.env != "dev" && fs.env != "dev-corp" {
|
|
return fs, fmt.Errorf("invalid environment %q: --env must be dev or dev-corp", fs.env)
|
|
}
|
|
if selectors > 1 {
|
|
return fs, fmt.Errorf("at most one of --model, --node, --provider is allowed")
|
|
}
|
|
}
|
|
return fs, nil
|
|
}
|
|
|
|
func printUsage(w io.Writer) {
|
|
fmt.Fprintln(w, "usage: inventory-query --env <env> [--model <model>] | [--node <node>] | [--provider <provider>]")
|
|
fmt.Fprintln(w, " --env required. inventory environment: dev or dev-corp")
|
|
fmt.Fprintln(w, " --model match a model alias/id")
|
|
fmt.Fprintln(w, " --node match a node id or alias")
|
|
fmt.Fprintln(w, " --provider match a provider id")
|
|
}
|
|
|
|
// loadInventory reads the inventory.yaml for the given env from agent-test/<env>/.
|
|
func loadInventory(env string) (map[string]interface{}, error) {
|
|
path := fmt.Sprintf("agent-test/%s/inventory.yaml", env)
|
|
data, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("cannot read inventory for %s: %w", env, err)
|
|
}
|
|
var doc map[string]interface{}
|
|
if err := yaml.Unmarshal(data, &doc); err != nil {
|
|
return nil, fmt.Errorf("invalid YAML in %s: %w", path, err)
|
|
}
|
|
return doc, nil
|
|
}
|
|
|
|
// validateInventory validates the inventory structure.
|
|
func validateInventory(data map[string]interface{}, requestedEnv string, kind string) error {
|
|
testEnv, ok := data["test_env"].(string)
|
|
if !ok {
|
|
return fmt.Errorf("missing or invalid type for test_env")
|
|
}
|
|
if testEnv != requestedEnv {
|
|
return fmt.Errorf("env identity mismatch: requested %s, found %s", requestedEnv, testEnv)
|
|
}
|
|
|
|
if _, ok := data["profile"].(string); !ok {
|
|
return fmt.Errorf("missing or invalid type for profile")
|
|
}
|
|
if _, ok := data["last_updated_at"].(string); !ok {
|
|
return fmt.Errorf("missing or invalid type for last_updated_at")
|
|
}
|
|
if _, ok := data["source"].(map[string]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for source")
|
|
}
|
|
if _, ok := data["edge"].(map[string]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for edge")
|
|
}
|
|
if _, ok := data["build"].(map[string]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for build")
|
|
}
|
|
|
|
if kind == "model" {
|
|
m, ok := data["model"].(map[string]interface{})
|
|
if !ok {
|
|
return fmt.Errorf("missing or invalid type for model")
|
|
}
|
|
if _, ok := m["aliases"].(map[string]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for model.aliases")
|
|
}
|
|
}
|
|
|
|
if kind == "provider" {
|
|
if _, ok := data["model"].(map[string]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for model")
|
|
}
|
|
if _, ok := data["nodes"].([]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for nodes")
|
|
}
|
|
}
|
|
|
|
if kind == "node" {
|
|
if _, ok := data["nodes"].([]interface{}); !ok {
|
|
return fmt.Errorf("missing or invalid type for nodes")
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
// buildEnvProjection returns the bounded env projection without model/nodes.
|
|
func buildEnvProjection(data map[string]interface{}) envProjection {
|
|
proj := envProjection{}
|
|
if v, ok := data["test_env"].(string); ok {
|
|
proj.Env = v
|
|
}
|
|
if v, ok := data["profile"].(string); ok {
|
|
proj.Profile = v
|
|
}
|
|
if v, ok := data["last_updated_at"].(string); ok {
|
|
proj.LastUpdatedAt = v
|
|
}
|
|
if v, ok := data["source"]; ok {
|
|
proj.Source = v
|
|
}
|
|
if v, ok := data["edge"]; ok {
|
|
proj.Edge = v
|
|
}
|
|
if v, ok := data["build"]; ok {
|
|
proj.Build = v
|
|
}
|
|
return proj
|
|
}
|
|
|
|
// queryInventory searches the inventory for matches of the given kind and selector.
|
|
func queryInventory(data map[string]interface{}, kind, selector string) []match {
|
|
var matches []match
|
|
switch kind {
|
|
case "node":
|
|
matches = queryNodes(data, selector)
|
|
case "provider":
|
|
matches = queryProviders(data, selector)
|
|
case "model":
|
|
matches = queryModels(data, selector)
|
|
}
|
|
sortMatches(matches)
|
|
return matches
|
|
}
|
|
|
|
// queryNodes finds nodes by id or alias at top-level nodes[].
|
|
func queryNodes(data map[string]interface{}, selector string) []match {
|
|
var matches []match
|
|
nodes, ok := data["nodes"].([]interface{})
|
|
if !ok {
|
|
return matches
|
|
}
|
|
for i, n := range nodes {
|
|
obj, ok := n.(map[string]interface{})
|
|
if !ok {
|
|
continue
|
|
}
|
|
path := fmt.Sprintf("nodes[%d]", i)
|
|
if matchString(obj, "id", selector) || matchString(obj, "alias", selector) {
|
|
matches = append(matches, match{Path: path, Value: obj})
|
|
}
|
|
}
|
|
sortMatches(matches)
|
|
return matches
|
|
}
|
|
|
|
// queryProviders finds providers by id in any canonical provider shapes.
|
|
func queryProviders(data map[string]interface{}, selector string) []match {
|
|
collector := newProviderCollector(selector)
|
|
collector.collect(data, "")
|
|
sortMatches(collector.matches)
|
|
return collector.matches
|
|
}
|
|
|
|
func joinPath(base, elem string) string {
|
|
if base == "" {
|
|
return elem
|
|
}
|
|
return base + "." + elem
|
|
}
|
|
|
|
// queryModels finds models by alias key in model.aliases or by id/alias/model_id in model subtree entities.
|
|
func queryModels(data map[string]interface{}, selector string) []match {
|
|
var matches []match
|
|
seen := make(map[string]bool)
|
|
|
|
model, ok := data["model"].(map[string]interface{})
|
|
if !ok {
|
|
return matches
|
|
}
|
|
|
|
// Check model.aliases keys.
|
|
if aliases, ok := model["aliases"].(map[string]interface{}); ok {
|
|
for key, val := range aliases {
|
|
if key == selector && !seen["model.aliases."+key] {
|
|
matches = append(matches, match{Path: "model.aliases." + key, Value: val})
|
|
seen["model.aliases."+key] = true
|
|
}
|
|
}
|
|
}
|
|
|
|
// Check model.subtree entities recursively for id/alias/model_id.
|
|
entityMatches := collectEntityMatches(model, "model", selector, seen)
|
|
matches = append(matches, entityMatches...)
|
|
|
|
sortMatches(matches)
|
|
return matches
|
|
}
|
|
|
|
// collectEntityMatches recursively searches a subtree for entities matching id/alias/model_id.
|
|
func collectEntityMatches(node interface{}, basePath, selector string, seen map[string]bool) []match {
|
|
var matches []match
|
|
|
|
switch n := node.(type) {
|
|
case map[string]interface{}:
|
|
if id, ok := n["id"].(string); ok && id == selector {
|
|
key := basePath + ".id"
|
|
if !seen[key] {
|
|
matches = append(matches, match{Path: key, Value: id})
|
|
seen[key] = true
|
|
}
|
|
}
|
|
if a, ok := n["alias"].(string); ok && a == selector {
|
|
key := basePath + ".alias"
|
|
if !seen[key] {
|
|
matches = append(matches, match{Path: key, Value: a})
|
|
seen[key] = true
|
|
}
|
|
}
|
|
if mid, ok := n["model_id"].(string); ok && mid == selector {
|
|
key := basePath + ".model_id"
|
|
if !seen[key] {
|
|
matches = append(matches, match{Path: key, Value: mid})
|
|
seen[key] = true
|
|
}
|
|
}
|
|
|
|
for k, v := range n {
|
|
childPath := basePath + "." + k
|
|
matches = append(matches, collectEntityMatches(v, childPath, selector, seen)...)
|
|
}
|
|
|
|
case []interface{}:
|
|
for i, v := range n {
|
|
childPath := fmt.Sprintf("%s[%d]", basePath, i)
|
|
matches = append(matches, collectEntityMatches(v, childPath, selector, seen)...)
|
|
}
|
|
}
|
|
|
|
return matches
|
|
}
|
|
|
|
// matchString checks if the given key in obj has the exact string value.
|
|
func matchString(obj map[string]interface{}, key, value string) bool {
|
|
if v, ok := obj[key].(string); ok && v == value {
|
|
return true
|
|
}
|
|
return false
|
|
}
|
|
|
|
// sortMatches sorts matches by path ascending.
|
|
func sortMatches(matches []match) {
|
|
sort.Slice(matches, func(i, j int) bool {
|
|
return matches[i].Path < matches[j].Path
|
|
})
|
|
}
|
|
|
|
// encodeJSON encodes v as indented JSON to w.
|
|
func encodeJSON(w io.Writer, v interface{}) error {
|
|
enc := json.NewEncoder(w)
|
|
enc.SetIndent("", " ")
|
|
enc.SetEscapeHTML(false)
|
|
return enc.Encode(v)
|
|
}
|