feat: search multiple NetBox instances with isolated host identities

2026-10-27 23:14:42 +01:00
1 parent a2a6dd7591
commit 7f53b90266
@@ -15,6 +15,7 @@ import (
 )

 type Client struct {
+	sources      map[string]*Client
 	baseURL      string
 	token        string
 	tokenVersion int
@@ -38,6 +39,9 @@ func NewClient(baseURL, token string, tokenVersion int) *Client {
 // Search queries up to 50 devices and VMs in parallel and merges the results.
 // Use SearchOptions to restrict by kind or tag.
 func (c *Client) Search(ctx context.Context, query string, opts SearchOptions) ([]HostEntry, error) {
+	if c.sources != nil {
+		return c.searchSources(ctx, query, opts, false)
+	}
 	if opts.Kind != "" && opts.Kind != "device" && opts.Kind != "vm" {
 		return nil, fmt.Errorf("invalid host kind %q", opts.Kind)
 	}
@@ -95,6 +99,9 @@ func (c *Client) Search(ctx context.Context, query string, opts SearchOptions) (
 // SearchAll paginates through all matching devices and VMs, fetching every page.
 // Intended for cache refresh; use Search for interactive queries.
 func (c *Client) SearchAll(ctx context.Context, query string, opts SearchOptions) ([]HostEntry, error) {
+	if c.sources != nil {
+		return c.searchSources(ctx, query, opts, true)
+	}
 	if opts.Kind != "" && opts.Kind != "device" && opts.Kind != "vm" {
 		return nil, fmt.Errorf("invalid host kind %q", opts.Kind)
 	}
@@ -146,6 +153,13 @@ func (c *Client) SearchAll(ctx context.Context, query string, opts SearchOptions
 // GetIPs returns all IP addresses assigned to a host, used by resolver strategies
 // that need more than just the primary IP.
 func (c *Client) GetIPs(ctx context.Context, entry HostEntry) ([]string, error) {
+	if c.sources != nil {
+		source, host, err := c.sourceForHost(entry)
+		if err != nil {
+			return nil, err
+		}
+		return source.GetIPs(ctx, host)
+	}
 	var apiURL string
 	switch entry.Kind {
 	case "device":
@@ -170,6 +184,9 @@ func (c *Client) GetIPs(ctx context.Context, entry HostEntry) ([]string, error)

 // GetIPsWithFilter calls /api/ipam/ip-addresses/ with arbitrary filter query parameters.
 func (c *Client) GetIPsWithFilter(ctx context.Context, filterParams string) ([]string, error) {
+	if c.sources != nil {
+		return nil, fmt.Errorf("IP filters require a host instance; use GetIPsWithFilterForHost")
+	}
 	apiURL := fmt.Sprintf("%s/api/ipam/ip-addresses/?%s&limit=100", c.baseURL, filterParams)
 	var resp netboxIPListResponse
 	if err := c.get(ctx, apiURL, &resp); err != nil {
@@ -360,12 +377,13 @@ func (c *Client) FindHost(ctx context.Context, name, kind string) (HostEntry, er
 		return HostEntry{}, fmt.Errorf("invalid kind %q", kind)
 	}
 	entries, err := c.SearchAll(ctx, name, SearchOptions{Kind: kind})
+	_, hostname := SplitTarget(name)
 	if err != nil {
 		return HostEntry{}, err
 	}
 	var matches []HostEntry
 	for _, e := range entries {
-		if strings.EqualFold(e.Name, name) {
+		if strings.EqualFold(e.Name, hostname) {
 			matches = append(matches, e)
 		}
 	}
@@ -373,6 +391,9 @@ func (c *Client) FindHost(ctx context.Context, name, kind string) (HostEntry, er
 		return HostEntry{}, fmt.Errorf("%w: %q", ErrHostNotFound, name)
 	}
 	if len(matches) != 1 {
+		if c.sources != nil {
+			return HostEntry{}, fmt.Errorf("host %q: found %d exact matches; use instance/hostname and --kind to select a unique host", name, len(matches))
+		}
 		return HostEntry{}, fmt.Errorf("host %q: found %d exact matches; use --kind or a unique hostname", name, len(matches))
 	}
 	return matches[0], nil
@@ -395,6 +416,13 @@ func (c *Client) hostURL(host HostEntry) (string, error) {
 }

 func (c *Client) GetComments(ctx context.Context, host HostEntry) (string, string, error) {
+	if c.sources != nil {
+		source, host, err := c.sourceForHost(host)
+		if err != nil {
+			return "", "", err
+		}
+		return source.GetComments(ctx, host)
+	}
 	u, err := c.hostURL(host)
 	if err != nil {
 		return "", "", err
@@ -407,6 +435,13 @@ func (c *Client) GetComments(ctx context.Context, host HostEntry) (string, strin
 }

 func (c *Client) SetComments(ctx context.Context, host HostEntry, comments, etag string) error {
+	if c.sources != nil {
+		source, host, err := c.sourceForHost(host)
+		if err != nil {
+			return err
+		}
+		return source.SetComments(ctx, host, comments, etag)
+	}
 	u, err := c.hostURL(host)
 	if err != nil {
 		return err
@@ -1,5 +1,7 @@
 package netbox

+import "strings"
+
 // SearchOptions filters a Search or SearchAll query.
 // Zero value means no filtering (return all kinds, no tag filter).
 type SearchOptions struct {
@@ -9,6 +11,7 @@ type SearchOptions struct {

 // HostEntry is a unified model for both devices and virtual machines from NetBox.
 type HostEntry struct {
+	Source     string // named NetBox instance; empty for the legacy single-instance config
 	ID         int
 	Name       string
 	Kind       string // "device" | "vm"
@@ -58,3 +61,19 @@ type netboxIPListResponse struct {
 		} `json:"assigned_object"`
 	} `json:"results"`
 }
+
+// Target returns an unambiguous selector for cache, completion, and SSH.
+func (h HostEntry) Target() string {
+	if h.Source == "" {
+		return h.Name
+	}
+	return h.Source + "/" + h.Name
+}
+
+// SplitTarget separates an optional instance prefix from a hostname.
+func SplitTarget(target string) (source, name string) {
+	if source, name, ok := strings.Cut(target, "/"); ok {
+		return strings.ToLower(source), name
+	}
+	return "", target
+}
@@ -0,0 +1,113 @@
+package netbox
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sort"
+	"strings"
+	"sync"
+)
+
+// Source binds a name to its own API client and credentials.
+type Source struct {
+	Name         string
+	URL          string
+	Token        string
+	TokenVersion int
+}
+
+func NewMultiClient(sources []Source) *Client {
+	c := &Client{sources: make(map[string]*Client)}
+	for _, source := range sources {
+		c.sources[strings.ToLower(source.Name)] = NewClient(source.URL, source.Token, source.TokenVersion)
+	}
+	return c
+}
+
+func (c *Client) searchSources(ctx context.Context, query string, opts SearchOptions, all bool) ([]HostEntry, error) {
+	requested, query := SplitTarget(query)
+	if requested != "" {
+		if _, ok := c.sources[requested]; !ok {
+			return nil, fmt.Errorf("unknown NetBox instance %q", requested)
+		}
+	}
+	var mu sync.Mutex
+	var wg sync.WaitGroup
+	var entries []HostEntry
+	var errs []error
+	started := 0
+	for name, client := range c.sources {
+		if requested != "" && name != requested {
+			continue
+		}
+		started++
+		wg.Add(1)
+		go func(name string, client *Client) {
+			defer wg.Done()
+			var found []HostEntry
+			var err error
+			if all {
+				found, err = client.SearchAll(ctx, query, opts)
+			} else {
+				found, err = client.Search(ctx, query, opts)
+			}
+			mu.Lock()
+			defer mu.Unlock()
+			if err != nil {
+				errs = append(errs, fmt.Errorf("NetBox %s: %w", name, err))
+				return
+			}
+			for i := range found {
+				found[i].Source = name
+			}
+			entries = append(entries, found...)
+		}(name, client)
+	}
+	wg.Wait()
+	// Full lookups must see every source to rule out ambiguous hostnames.
+	if len(errs) > 0 && (all || len(errs) == started) {
+		return nil, errors.Join(errs...)
+	}
+	sort.Slice(entries, func(i, j int) bool { return entries[i].Target() < entries[j].Target() })
+	return entries, errors.Join(errs...)
+}
+
+func (c *Client) sourceForHost(host HostEntry) (*Client, HostEntry, error) {
+	prefix, name := SplitTarget(host.Name)
+	if prefix != "" {
+		if host.Source != "" && !strings.EqualFold(prefix, host.Source) {
+			return nil, host, fmt.Errorf("conflicting NetBox instance for %q", host.Name)
+		}
+		host.Source, host.Name = prefix, name
+	}
+	if c.sources == nil {
+		return c, host, nil
+	}
+	source, ok := c.sources[strings.ToLower(host.Source)]
+	if !ok {
+		return nil, host, fmt.Errorf("specify a configured NetBox instance as instance/hostname (got %q)", host.Target())
+	}
+	host.Source = strings.ToLower(host.Source)
+	return source, host, nil
+}
+
+// HostScope identifies a local note by its original NetBox URL and canonical
+// hostname, so renaming an instance does not move or mix its notes.
+func (c *Client) HostScope(host HostEntry) (HostEntry, string, error) {
+	source, host, err := c.sourceForHost(host)
+	if err != nil {
+		return host, "", err
+	}
+	return host, source.baseURL, nil
+}
+
+func (c *Client) GetIPsWithFilterForHost(ctx context.Context, host HostEntry, filters string) ([]string, error) {
+	source, _, err := c.sourceForHost(host)
+	if err != nil {
+		return nil, err
+	}
+	return source.GetIPsWithFilter(ctx, filters)
+}
+
+func (c *Client) MultipleSources() bool { return c.sources != nil }
@@ -0,0 +1,192 @@
+package netbox
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync"
+	"testing"
+	"time"
+)
+
+func sourceTestServer(t *testing.T, token, ip, comment string) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if got := r.Header.Get("Authorization"); got != "Token "+token {
+			t.Errorf("credentials crossed sources: %q, want %q", got, token)
+			http.Error(w, "bad token", 401)
+			return
+		}
+		w.Header().Set("Content-Type", "application/json")
+		switch r.URL.Path {
+		case "/api/dcim/devices/":
+			json.NewEncoder(w).Encode(deviceListResponse(netboxDevice{ID: 1, Name: "web", PrimaryIP4: &netboxIP{Address: ip + "/24"}}))
+		case "/api/virtualization/virtual-machines/":
+			json.NewEncoder(w).Encode(vmListResponse())
+		case "/api/ipam/ip-addresses/":
+			if r.URL.Query().Get("device_id") != "1" {
+				t.Errorf("wrong host ID: %s", r.URL.RawQuery)
+			}
+			json.NewEncoder(w).Encode(ipListResponse(ip + "/24"))
+		case "/api/dcim/devices/1/":
+			if r.Method == http.MethodPatch {
+				var data map[string]string
+				if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
+					t.Error(err)
+				}
+				if data["comments"] != comment+" updated" {
+					t.Errorf("comment write crossed sources: %v", data)
+				}
+			}
+			json.NewEncoder(w).Encode(map[string]string{"comments": comment})
+		default:
+			http.NotFound(w, r)
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func TestMultiSourceSearchAndRouting(t *testing.T) {
+	office := sourceTestServer(t, "office-token", "10.0.0.1", "office note")
+	lab := sourceTestServer(t, "lab-token", "10.0.0.2", "lab note")
+	c := NewMultiClient([]Source{{Name: "office", URL: office.URL, Token: "office-token"}, {Name: "lab", URL: lab.URL, Token: "lab-token"}})
+	ctx := context.Background()
+	entries, err := c.Search(ctx, "web", SearchOptions{})
+	if err != nil || len(entries) != 2 {
+		t.Fatalf("combined search: %+v, %v", entries, err)
+	}
+	if entries[0].Target() != "lab/web" || entries[1].Target() != "office/web" {
+		t.Fatalf("source lost: %+v", entries)
+	}
+	if _, err := c.FindHost(ctx, "web", ""); err == nil {
+		t.Fatal("ambiguous unqualified hostname accepted")
+	}
+	for _, tc := range []struct{ name, ip, note string }{{"office", "10.0.0.1", "office note"}, {"lab", "10.0.0.2", "lab note"}} {
+		t.Run(tc.name, func(t *testing.T) {
+			host, err := c.FindHost(ctx, tc.name+"/web", "")
+			if err != nil {
+				t.Fatal(err)
+			}
+			if host.Source != tc.name || host.Name != "web" || host.ID != 1 {
+				t.Fatalf("wrong host identity: %+v", host)
+			}
+			ips, err := c.GetIPs(ctx, host)
+			if err != nil || len(ips) != 1 || ips[0] != tc.ip {
+				t.Fatalf("wrong IP source: %v, %v", ips, err)
+			}
+			ips, err = c.GetIPsWithFilterForHost(ctx, host, "device_id=1&interface_name=mgmt0")
+			if err != nil || len(ips) != 1 || ips[0] != tc.ip {
+				t.Fatalf("wrong interface source: %v, %v", ips, err)
+			}
+			text, _, err := c.GetComments(ctx, host)
+			if err != nil || text != tc.note {
+				t.Fatalf("wrong comment source: %q, %v", text, err)
+			}
+			// TUI cache items carry a qualified name rather than a separate source.
+			qualified := HostEntry{Name: tc.name + "/web", ID: 1, Kind: "device"}
+			if err := c.SetComments(ctx, qualified, tc.note+" updated", ""); err != nil {
+				t.Fatal(err)
+			}
+		})
+	}
+	if _, err := c.Search(ctx, "missing/web", SearchOptions{}); err == nil {
+		t.Fatal("unknown source accepted")
+	}
+	if _, _, err := c.HostScope(HostEntry{Name: "web", Kind: "device"}); err == nil {
+		t.Fatal("unqualified note namespace accepted")
+	}
+	if _, err := c.GetIPsWithFilter(ctx, "device_id=1"); err == nil {
+		t.Fatal("IP filter without instance accepted")
+	}
+}
+
+func TestMultiSourceSearchRunsConcurrently(t *testing.T) {
+	var wg sync.WaitGroup
+	wg.Add(2)
+	makeServer := func() *httptest.Server {
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			wg.Done()
+			both := make(chan struct{})
+			go func() { wg.Wait(); close(both) }()
+			select {
+			case <-both:
+				fmt.Fprint(w, `{"count":0,"results":[]}`)
+			case <-r.Context().Done():
+			}
+		}))
+		t.Cleanup(srv.Close)
+		return srv
+	}
+	a, b := makeServer(), makeServer()
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+	defer cancel()
+	c := NewMultiClient([]Source{{Name: "a", URL: a.URL}, {Name: "b", URL: b.URL}})
+	if _, err := c.Search(ctx, "web", SearchOptions{Kind: "device"}); err != nil {
+		t.Fatalf("sources did not run concurrently: %v", err)
+	}
+}
+
+func TestMultiSourcePartialFailure(t *testing.T) {
+	good := sourceTestServer(t, "good-token", "10.0.0.1", "note")
+	failed := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.Error(w, "unavailable", 503) }))
+	defer failed.Close()
+	c := NewMultiClient([]Source{{Name: "good", URL: good.URL, Token: "good-token"}, {Name: "offline", URL: failed.URL}})
+	ctx := context.Background()
+	entries, err := c.Search(ctx, "web", SearchOptions{})
+	if len(entries) != 1 || err == nil || !strings.Contains(err.Error(), "offline") {
+		t.Fatalf("partial search lost results or warning: %+v, %v", entries, err)
+	}
+	if _, err := c.SearchAll(ctx, "web", SearchOptions{}); err == nil {
+		t.Fatal("incomplete full search succeeded")
+	}
+	if _, err := c.FindHost(ctx, "web", ""); err == nil {
+		t.Fatal("incomplete search claimed unique host")
+	}
+	if host, err := c.FindHost(ctx, "good/web", ""); err != nil || host.Source != "good" {
+		t.Fatalf("qualified lookup blocked by unrelated outage: %+v, %v", host, err)
+	}
+}
+
+func TestMultiSourceKeepsIndependentDeviceAndVMInventories(t *testing.T) {
+	office := newTestServer(t, map[string]any{
+		"/api/dcim/devices/":                    deviceListResponse(netboxDevice{ID: 1, Name: "office-router", PrimaryIP4: &netboxIP{Address: "10.1.0.1/24"}}),
+		"/api/virtualization/virtual-machines/": vmListResponse(netboxVM{ID: 2, Name: "office-vm", PrimaryIP4: &netboxIP{Address: "10.1.0.2/24"}}),
+	})
+	defer office.Close()
+	lab := newTestServer(t, map[string]any{
+		"/api/dcim/devices/":                    deviceListResponse(netboxDevice{ID: 1, Name: "lab-switch", PrimaryIP4: &netboxIP{Address: "10.2.0.1/24"}}),
+		"/api/virtualization/virtual-machines/": vmListResponse(netboxVM{ID: 2, Name: "lab-vm", PrimaryIP4: &netboxIP{Address: "10.2.0.2/24"}}),
+	})
+	defer lab.Close()
+	client := NewMultiClient([]Source{{Name: "office", URL: office.URL}, {Name: "lab", URL: lab.URL}})
+	entries, err := client.SearchAll(context.Background(), "", SearchOptions{})
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := map[string]string{"office/office-router": "device", "office/office-vm": "vm", "lab/lab-switch": "device", "lab/lab-vm": "vm"}
+	if len(entries) != len(want) {
+		t.Fatalf("independent inventories missing: %+v", entries)
+	}
+	for _, entry := range entries {
+		if want[entry.Target()] != entry.Kind {
+			t.Fatalf("inventory crossed source or kind: %+v", entry)
+		}
+	}
+	vms, err := client.Search(context.Background(), "", SearchOptions{Kind: "vm"})
+	if err != nil || len(vms) != 2 {
+		t.Fatalf("VM filter across instances: %+v, %v", vms, err)
+	}
+	for _, host := range vms {
+		if host.Kind != "vm" {
+			t.Fatalf("device leaked into VM filter: %+v", host)
+		}
+	}
+	host, err := client.FindHost(context.Background(), "lab/lab-vm", "vm")
+	if err != nil || host.Source != "lab" || host.ID != 2 || host.PrimaryIP4 != "10.2.0.2" {
+		t.Fatalf("wrong VM selected: %+v, %v", host, err)
+	}
+}