feat: add branch support (#202)

Co-authored-by: Michal Fiedorowicz <mfiedorowicz@netboxlabs.com>
netboxlabs · Dec 16, 2024 · a8b8ef7 · a8b8ef7
1 parent 61b7d25
commit a8b8ef7
Show file tree

Hide file tree

Showing 9 changed files with 87 additions and 10 deletions.
diff --git a/diode-server/netboxdiodeplugin/client.go b/diode-server/netboxdiodeplugin/client.go
@@ -15,6 +15,7 @@ import (
 	"os"
 	"reflect"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/mitchellh/mapstructure"
@@ -42,6 +43,11 @@ const (
 	defaultBaseURL = "http://127.0.0.1:8080/api/plugins/diode"
 
 	defaultHTTPTimeoutSeconds = 5
+
+	// NetBoxBranchHeader is an HTTP header that indicates the NetBox branch to target
+	NetBoxBranchHeader = "X-NetBox-Branch"
+	// NetBoxBranchParam is a query parameter that indicates the NetBox branch to target
+	NetBoxBranchParam = "_branch"
 )
 
 var (
@@ -265,6 +271,7 @@ type ObjectState struct {
 type RetrieveObjectStateQueryParams struct {
 	ObjectType string
 	ObjectID   int
+	BranchID   string
 	Params     map[string]string
 }
 
@@ -280,6 +287,10 @@ func (c *Client) RetrieveObjectState(ctx context.Context, params RetrieveObjectS
 	if params.ObjectID > 0 {
 		queryParams.Set("object_id", strconv.Itoa(params.ObjectID))
 	}
+	branchID := strings.TrimSpace(params.BranchID)
+	if branchID != "" {
+		queryParams.Set(NetBoxBranchParam, branchID)
+	}
 	for k, v := range params.Params {
 		queryParams.Set(k, v)
 	}
@@ -391,6 +402,7 @@ func statusMapToStringHookFunc() mapstructure.DecodeHookFunc {
 type ChangeSetRequest struct {
 	ChangeSetID string   `json:"change_set_id"`
 	ChangeSet   []Change `json:"change_set"`
+	BranchID    string   `json:"-"` // Supplied as header
 }
 
 // Change represents a change
@@ -430,6 +442,11 @@ func (c *Client) ApplyChangeSet(ctx context.Context, payload ChangeSetRequest) (
 	}
 	req.Header.Set("Content-Type", "application/json")
 
+	branchID := strings.TrimSpace(payload.BranchID)
+	if branchID != "" {
+		req.Header.Set(NetBoxBranchHeader, branchID)
+	}
+
 	resp, err := c.httpClient.Do(req)
 	if err != nil {
 		return nil, err

diff --git a/diode-server/netboxdiodeplugin/client_test.go b/diode-server/netboxdiodeplugin/client_test.go
@@ -202,6 +202,25 @@ func TestRetrieveObjectState(t *testing.T) {
 			tlsSkipVerify: true,
 			shouldError:   false,
 		},
+		{
+			name:               "valid response for DCIM site with branch",
+			params:             netboxdiodeplugin.RetrieveObjectStateQueryParams{ObjectType: netbox.DcimSiteObjectType, ObjectID: 1, BranchID: "branch_id"},
+			mockServerResponse: `{"object_type":"dcim.site","object_change_id":1,"object":{"id":1,"name":"site 01", "slug": "site-01"}}`,
+			apiKey:             "foobar",
+			response: &netboxdiodeplugin.ObjectState{
+				ObjectType:     netbox.DcimSiteObjectType,
+				ObjectChangeID: 1,
+				Object: &netbox.DcimSiteDataWrapper{
+					Site: &netbox.DcimSite{
+						ID:   1,
+						Name: "site 01",
+						Slug: "site-01",
+					},
+				},
+			},
+			tlsSkipVerify: true,
+			shouldError:   false,
+		},
 		{
 			name:               "valid response for DCIM DeviceRole",
 			params:             netboxdiodeplugin.RetrieveObjectStateQueryParams{ObjectType: netbox.DcimDeviceRoleObjectType, ObjectID: 1},
@@ -552,6 +571,11 @@ func TestRetrieveObjectState(t *testing.T) {
 				assert.Equal(t, r.URL.Query().Get("object_id"), objectID)
 				assert.Equal(t, r.Header.Get("Authorization"), fmt.Sprintf("Token %s", tt.apiKey))
 				assert.Equal(t, r.Header.Get("User-Agent"), fmt.Sprintf("%s/%s", netboxdiodeplugin.SDKName, netboxdiodeplugin.SDKVersion))
+				if tt.params.BranchID != "" {
+					assert.Equal(t, r.URL.Query().Get(netboxdiodeplugin.NetBoxBranchParam), tt.params.BranchID)
+				} else {
+					assert.False(t, r.URL.Query().Has(netboxdiodeplugin.NetBoxBranchParam))
+				}
 				_, _ = w.Write([]byte(tt.mockServerResponse))
 			}
 
@@ -624,6 +648,33 @@ func TestApplyChangeSet(t *testing.T) {
 			},
 			shouldError: false,
 		},
+		{
+			name:   "valid apply change set response with branch",
+			apiKey: "foobar",
+			changeSetRequest: netboxdiodeplugin.ChangeSetRequest{
+				ChangeSetID: "00000000-0000-0000-0000-000000000000",
+				BranchID:    "test-branch",
+				ChangeSet: []netboxdiodeplugin.Change{
+					{
+						ChangeID:      "00000000-0000-0000-0000-000000000001",
+						ChangeType:    "create",
+						ObjectType:    "dcim.device",
+						ObjectID:      nil,
+						ObjectVersion: nil,
+						Data: &netbox.DcimDevice{
+							Name: "test",
+						},
+					},
+				},
+			},
+			mockServerResponse: `{"change_set_id":"00000000-0000-0000-0000-000000000000","result":"success"}`,
+			mockStatusCode:     http.StatusOK,
+			response: &netboxdiodeplugin.ChangeSetResponse{
+				ChangeSetID: "00000000-0000-0000-0000-000000000000",
+				Result:      "success",
+			},
+			shouldError: false,
+		},
 		{
 			name:   "invalid request",
 			apiKey: "foobar",
@@ -722,6 +773,11 @@ func TestApplyChangeSet(t *testing.T) {
 				assert.Equal(t, r.Header.Get("Authorization"), fmt.Sprintf("Token %s", tt.apiKey))
 				assert.Equal(t, r.Header.Get("User-Agent"), fmt.Sprintf("%s/%s", netboxdiodeplugin.SDKName, netboxdiodeplugin.SDKVersion))
 				assert.Equal(t, r.Header.Get("Content-Type"), "application/json")
+				if tt.changeSetRequest.BranchID != "" {
+					assert.Equal(t, r.Header.Get(netboxdiodeplugin.NetBoxBranchHeader), tt.changeSetRequest.BranchID)
+				} else {
+					assert.Len(t, r.Header.Values(netboxdiodeplugin.NetBoxBranchHeader), 0)
+				}
 				w.WriteHeader(tt.mockStatusCode)
 				_, _ = w.Write([]byte(tt.mockServerResponse))
 			}

diff --git a/diode-server/reconciler/applier/applier.go b/diode-server/reconciler/applier/applier.go
@@ -9,7 +9,7 @@ import (
 )
 
 // ApplyChangeSet applies a change set to NetBox
-func ApplyChangeSet(ctx context.Context, logger *slog.Logger, cs changeset.ChangeSet, nbClient netboxdiodeplugin.NetBoxAPI) error {
+func ApplyChangeSet(ctx context.Context, logger *slog.Logger, cs changeset.ChangeSet, branchID string, nbClient netboxdiodeplugin.NetBoxAPI) error {
 	changes := make([]netboxdiodeplugin.Change, 0)
 	for _, change := range cs.ChangeSet {
 		changes = append(changes, netboxdiodeplugin.Change{
@@ -25,6 +25,8 @@ func ApplyChangeSet(ctx context.Context, logger *slog.Logger, cs changeset.Chang
 	req := netboxdiodeplugin.ChangeSetRequest{
 		ChangeSetID: cs.ChangeSetID,
 		ChangeSet:   changes,
+		// TODO(mfiedorowicz): take branch from ChangeSet, remove parameter
+		BranchID: branchID,
 	}
 
 	resp, err := nbClient.ApplyChangeSet(ctx, req)

diff --git a/diode-server/reconciler/applier/applier_test.go b/diode-server/reconciler/applier/applier_test.go
@@ -61,7 +61,7 @@ func TestApplyChangeSet(t *testing.T) {
 
 	mockNetBoxAPI.On("ApplyChangeSet", ctx, req).Return(resp, nil)
 
-	err := applier.ApplyChangeSet(ctx, logger, cs, mockNetBoxAPI)
+	err := applier.ApplyChangeSet(ctx, logger, cs, "", mockNetBoxAPI)
 	assert.NoError(t, err)
 	mockNetBoxAPI.AssertExpectations(t)
 }

diff --git a/diode-server/reconciler/differ/differ.go b/diode-server/reconciler/differ/differ.go
@@ -30,7 +30,7 @@ type ObjectState struct {
 }
 
 // Diff compares ingested entity with the intended state in NetBox and returns a change set
-func Diff(ctx context.Context, entity IngestEntity, netboxAPI netboxdiodeplugin.NetBoxAPI) (*changeset.ChangeSet, error) {
+func Diff(ctx context.Context, entity IngestEntity, branchID string, netboxAPI netboxdiodeplugin.NetBoxAPI) (*changeset.ChangeSet, error) {
 	// extract ingested entity (actual)
 	actual, err := extractIngestEntityData(entity)
 	if err != nil {
@@ -52,7 +52,7 @@ func Diff(ctx context.Context, entity IngestEntity, netboxAPI netboxdiodeplugin.
 	// retrieve root object all its nested objects from NetBox (intended)
 	intendedNestedObjectsMap := make(map[string]netbox.ComparableData)
 	for _, obj := range actualNestedObjects {
-		intended, err := retrieveObjectState(ctx, netboxAPI, obj)
+		intended, err := retrieveObjectState(ctx, netboxAPI, obj, branchID)
 		if err != nil {
 			return nil, err
 		}
@@ -93,16 +93,18 @@ func Diff(ctx context.Context, entity IngestEntity, netboxAPI netboxdiodeplugin.
 			ObjectID:      objectID,
 			ObjectVersion: nil,
 			Data:          obj.Data(),
+			// TODO(mfiedorowicz): include branchID
 		})
 	}
 
 	return &changeset.ChangeSet{ChangeSetID: uuid.NewString(), ChangeSet: changes}, nil
 }
 
-func retrieveObjectState(ctx context.Context, netboxAPI netboxdiodeplugin.NetBoxAPI, change netbox.ComparableData) (netbox.ComparableData, error) {
+func retrieveObjectState(ctx context.Context, netboxAPI netboxdiodeplugin.NetBoxAPI, change netbox.ComparableData, branchID string) (netbox.ComparableData, error) {
 	params := netboxdiodeplugin.RetrieveObjectStateQueryParams{
 		ObjectID:   0,
 		ObjectType: change.DataType(),
+		BranchID:   branchID,
 		Params:     change.ObjectStateQueryParams(),
 	}
 	resp, err := netboxAPI.RetrieveObjectState(ctx, params)

diff --git a/diode-server/reconciler/differ/differ_dcim_test.go b/diode-server/reconciler/differ/differ_dcim_test.go
@@ -4348,7 +4348,7 @@ func TestDcimPrepare(t *testing.T) {
 				}, nil)
 			}
 
-			cs, err := differ.Diff(ctx, tt.ingestEntity, mockClient)
+			cs, err := differ.Diff(ctx, tt.ingestEntity, "", mockClient)
 			if tt.wantErr {
 				require.Error(t, err)
 				return

diff --git a/diode-server/reconciler/differ/differ_ipam_test.go b/diode-server/reconciler/differ/differ_ipam_test.go
@@ -1858,7 +1858,7 @@ func TestIpamPrepare(t *testing.T) {
 				}, nil)
 			}
 
-			cs, err := differ.Diff(ctx, tt.ingestEntity, mockClient)
+			cs, err := differ.Diff(ctx, tt.ingestEntity, "", mockClient)
 			if tt.wantErr {
 				require.Error(t, err)
 				return

diff --git a/diode-server/reconciler/differ/differ_virt_test.go b/diode-server/reconciler/differ/differ_virt_test.go
@@ -1644,7 +1644,7 @@ func TestVirtualizationPrepare(t *testing.T) {
 				}, nil)
 			}
 
-			cs, err := differ.Diff(ctx, tt.ingestEntity, mockClient)
+			cs, err := differ.Diff(ctx, tt.ingestEntity, "", mockClient)
 			if tt.wantErr {
 				require.Error(t, err)
 				return

diff --git a/diode-server/reconciler/ingestion_processor.go b/diode-server/reconciler/ingestion_processor.go
@@ -283,7 +283,7 @@ func (p *IngestionProcessor) GenerateChangeSet(ctx context.Context, generateChan
 					State:     int(ingestionLog.ingestionLog.GetState()),
 				}
 
-				changeSet, err := differ.Diff(ctx, ingestEntity, p.nbClient)
+				changeSet, err := differ.Diff(ctx, ingestEntity, "", p.nbClient)
 				if err != nil {
 					tags := map[string]string{
 						"request_id": ingestEntity.RequestID,
@@ -375,7 +375,7 @@ func (p *IngestionProcessor) ApplyChangeSet(ctx context.Context, applyChan <-cha
 
 				p.logger.Debug("applying change set", "ingestionLogID", ingestionLog.ingestionLog.GetId(), "changeSetID", ingestionLog.changeSet.ChangeSetID)
 
-				if err := applier.ApplyChangeSet(ctx, p.logger, *ingestionLog.changeSet, p.nbClient); err != nil {
+				if err := applier.ApplyChangeSet(ctx, p.logger, *ingestionLog.changeSet, "", p.nbClient); err != nil {
 					p.logger.Debug("failed to apply change set", "ingestionLogID", ingestionLog.ingestionLog.GetId(), "changeSetID", ingestionLog.changeSet.ChangeSetID, "error", err)
 					ingestionLog.errors = append(ingestionLog.errors, fmt.Errorf("failed to apply chang eset: %v", err))