Merge pull request #288 from newrelic/host_id

feat: add host_id argument

Author: rogercoll

args/args.go

@@ -21,6 +21,7 @@ type DefaultArgumentList struct {
 	Metadata   bool   `default:"false" help:"Add customer defined key-value attributes to the samples."`
 	NriCluster string `default:"" help:"Optional. Cluster name"`
 	NriService string `default:"" help:"Optional. Service name"`
+	HostID     string `default:"" help:"Optional. Host ID to be set in entity or/and in the payload"`
 }
 
 // All returns if all data should be published

integration/integration.go

@@ -7,6 +7,7 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"regexp"
 	"strings"
 	"sync"
 	"time"
@@ -27,6 +28,10 @@ const (
 	protocolVersion = "4"
 )
 
+// A different regex is needed for replacing because `localhostRE` matches
+// IPV6 by using extra `:` that don't belong to the IP but are separators.
+var localhostReplaceRE = regexp.MustCompile(`(localhost|LOCALHOST|127(?:\.[0-9]+){0,2}\.[0-9]+|::1)`)
+
 // Metadata describes the integration
 type Metadata struct {
 	Name    string `json:"name"`
@@ -125,6 +130,15 @@ func (i *Integration) AddEntity(e *Entity) {
 func (i *Integration) Publish() error {
 	defer i.Clear()
 
+	hostID := i.GetHostID()
+	if hostID != "" {
+		for k := range i.Entities {
+			if i.Entities[k].Metadata != nil {
+				i.Entities[k].Metadata.Name = replaceLocalhost(i.Entities[k].Metadata.Name, hostID)
+			}
+		}
+	}
+
 	// add the host entity to the list of entities to be serialized, if not empty
 	if notEmpty(i.HostEntity) {
 		i.Entities = append(i.Entities, i.HostEntity)
@@ -159,6 +173,12 @@ func (i *Integration) MarshalJSON() (output []byte, err error) {
 	return
 }
 
+// GetHostID returns HostID or an empty string if not set
+func (i *Integration) GetHostID() string {
+	defaultArgs := args.GetDefaultArgs(i.args)
+	return defaultArgs.HostID
+}
+
 // toJSON serializes integration as JSON. If the pretty attribute is
 // set to true, the JSON will be indented for easy reading.
 func (i *Integration) toJSON(pretty bool) (output []byte, err error) {
@@ -272,3 +292,9 @@ func (i *Integration) addDefaultAttributes(e *Entity) error {
 
 	return nil
 }
+
+// ReplaceLocalhost replaces the occurrence of a localhost address with
+// the given hostname. This is done to avoid entity identifier collision.
+func replaceLocalhost(source, with string) string {
+	return localhostReplaceRE.ReplaceAllString(source, with)
+}

integration/integration_test.go

@@ -10,13 +10,15 @@ import (
 
 	"github.com/stretchr/testify/assert"
 
+	"github.com/newrelic/infra-integrations-sdk/v4/args"
 	"github.com/newrelic/infra-integrations-sdk/v4/data/event"
 	"github.com/newrelic/infra-integrations-sdk/v4/log"
 )
 
 var (
 	integrationName    = "TestIntegration"
 	integrationVersion = "1.0"
+	arguments          args.DefaultArgumentList
 )
 
 func Test_Integration_CreateIntegrationInitializesCorrectly(t *testing.T) {
@@ -463,6 +465,98 @@ func Test_Integration_PublishThrowsNoError(t *testing.T) {
 	assert.Empty(t, i.Entities)
 }
 
+func Test_Integration_HostId(t *testing.T) {
+	w := testWriter{
+		func(integrationBytes []byte) {
+			expectedOutputRaw := []byte(`
+{
+  "protocol_version": "4",
+  "integration": {
+    "name": "TestIntegration",
+    "version": "1.0"
+  },
+  "data": [
+    {
+      "common": {},
+      "entity": {
+        "name": "id-1234567890abc:8080",
+        "displayName": "",
+        "type": "test",
+        "metadata": {
+          "tags.env": "prod"
+        }
+      },
+      "metrics": [],
+      "inventory": {},
+      "events": []
+    },
+    {
+      "common": {
+        "attributes": {
+          "targetName": "id-1234567890abc"
+        }
+      },
+      "entity": {
+        "name": "EntityOne",
+        "displayName": "",
+        "type": "test",
+        "metadata": {}
+      },
+      "metrics": [
+        {
+          "timestamp": 10000000,
+          "name": "metricOne",
+          "type": "gauge",
+          "attributes": {
+            "processName": "java"
+          },
+          "value": 2
+        }
+      ],
+      "inventory": {},
+      "events": []
+    }
+  ]
+}`)
+
+			// awful but cannot compare with json.Unmarshal as it's not supported by Integration
+			assert.Equal(t, stripBlanks(expectedOutputRaw), stripBlanks(integrationBytes))
+		},
+	}
+
+	_ = os.Setenv("HOST_ID", "id-1234567890abc")
+
+	// os.ClearEnv breaks tests in Windows that use the fileStorer because clears the user env vars
+	defer func() {
+		_ = os.Unsetenv("HOST_ID")
+	}()
+
+	i, err := New("TestIntegration", "1.0", Logger(log.Discard), Writer(w), Args(&arguments))
+	assert.NoError(t, err)
+
+	e, err := i.NewEntity("localhost:8080", "test", "")
+	assert.NoError(t, err)
+	_ = e.AddTag("env", "prod")
+	i.AddEntity(e)
+
+	// add entity
+	e2, err := i.NewEntity("EntityOne", "test", "")
+	assert.NoError(t, err)
+	// add common attributes
+	e2.AddCommonDimension("targetName", i.GetHostID())
+	// add metric to entity 2
+	gauge, _ := Gauge(time.Unix(10000000, 0), "metricOne", 2)
+	_ = gauge.AddDimension("processName", "java")
+	e2.AddMetric(gauge)
+	// add entity 2 to integration
+	i.AddEntity(e2)
+
+	assert.NoError(t, i.Publish())
+
+	// check integration  was reset
+	assert.Empty(t, i.Entities)
+}
+
 func Test_Integration_FindEntity(t *testing.T) {
 	i := newTestIntegration(t)