feat: create new telemetry handle that supports connection strings (#3729)

BeegiiK · web-flow · commit 3e86bca5fc24 · 2025-07-24T15:43:54.000Z
* feat: create new telemetry handle with connection strings

* feat: Add application insights source code and remove dependency

* feat: Create connection string helper function and update telemetry handle

* feat: Polishing

* feat: Address PR comments

* feat: Update test telemetry handle initialization

* feat: Re-initialise telemetry handles to fix race conditions during test execution

* feat: revert string formatting changes

* feat: fix lint
diff --git a/aitelemetry/connection_string_parser.go b/aitelemetry/connection_string_parser.go
@@ -0,0 +1,52 @@
+package aitelemetry
+
+import (
+	"strings"
+
+	"github.com/pkg/errors"
+)
+
+type connectionVars struct {
+	instrumentationKey string
+	ingestionURL       string
+}
+
+func (c *connectionVars) String() string {
+	return "InstrumentationKey=" + c.instrumentationKey + ";IngestionEndpoint=" + c.ingestionURL
+}
+
+func parseConnectionString(connectionString string) (*connectionVars, error) {
+	connectionVars := &connectionVars{}
+
+	if connectionString == "" {
+		return nil, errors.New("connection string cannot be empty")
+	}
+
+	pairs := strings.Split(connectionString, ";")
+	for _, pair := range pairs {
+		kv := strings.SplitN(pair, "=", 2)
+		if len(kv) != 2 {
+			return nil, errors.Errorf("invalid connection string format: %s", pair)
+		}
+		key, value := strings.TrimSpace(kv[0]), strings.TrimSpace(kv[1])
+
+		if key == "" {
+			return nil, errors.Errorf("key in connection string cannot be empty")
+		}
+
+		switch strings.ToLower(key) {
+		case "instrumentationkey":
+			connectionVars.instrumentationKey = value
+		case "ingestionendpoint":
+			if value != "" {
+				connectionVars.ingestionURL = value + "v2.1/track"
+			}
+		}
+	}
+
+	if connectionVars.instrumentationKey == "" || connectionVars.ingestionURL == "" {
+		return nil, errors.Errorf("missing required fields in connection string: %s", connectionVars)
+	}
+
+	return connectionVars, nil
+}
diff --git a/aitelemetry/connection_string_parser_test.go b/aitelemetry/connection_string_parser_test.go
@@ -0,0 +1,66 @@
+package aitelemetry
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+const connectionString = "InstrumentationKey=0000-0000-0000-0000-0000;IngestionEndpoint=https://ingestion.endpoint.com/;LiveEndpoint=https://live.endpoint.com/;ApplicationId=1111-1111-1111-1111-1111"
+
+func TestParseConnectionString(t *testing.T) {
+	tests := []struct {
+		name             string
+		connectionString string
+		want             *connectionVars
+		wantErr          bool
+	}{
+		{
+			name:             "Valid connection string and instrumentation key",
+			connectionString: connectionString,
+			want: &connectionVars{
+				instrumentationKey: "0000-0000-0000-0000-0000",
+				ingestionURL:       "https://ingestion.endpoint.com/v2.1/track",
+			},
+			wantErr: false,
+		},
+		{
+			name:             "Invalid connection string format",
+			connectionString: "InvalidConnectionString",
+			want:             nil,
+			wantErr:          true,
+		},
+		{
+			name:             "Valid instrumentation key with missing ingestion endpoint",
+			connectionString: "InstrumentationKey=0000-0000-0000-0000-0000;IngestionEndpoint=",
+			want:             nil,
+			wantErr:          true,
+		},
+		{
+			name:             "Missing instrumentation key with valid ingestion endpoint",
+			connectionString: "InstrumentationKey=;IngestionEndpoint=https://ingestion.endpoint.com/",
+			want:             nil,
+			wantErr:          true,
+		},
+		{
+			name:             "Empty connection string",
+			connectionString: "",
+			want:             nil,
+			wantErr:          true,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			got, err := parseConnectionString(tt.connectionString)
+			if tt.wantErr {
+				require.Error(t, err, "Expected error but got none")
+			} else {
+				require.NoError(t, err, "Expected no error but got one")
+				require.NotNil(t, got, "Expected a non-nil result")
+				require.Equal(t, tt.want.instrumentationKey, got.instrumentationKey, "Instrumentation Key does not match")
+				require.Equal(t, tt.want.ingestionURL, got.ingestionURL, "Ingestion URL does not match")
+			}
+		})
+	}
+}
diff --git a/aitelemetry/telemetrywrapper.go b/aitelemetry/telemetrywrapper.go
@@ -13,6 +13,7 @@ import (
 	"github.com/Azure/azure-container-networking/store"
 	"github.com/microsoft/ApplicationInsights-Go/appinsights"
 	"github.com/microsoft/ApplicationInsights-Go/appinsights/contracts"
+	"github.com/pkg/errors"
 )
 
 const (
@@ -161,7 +162,7 @@ func isPublicEnvironment(url string, retryCount, waitTimeInSecs int) (bool, erro
 			return true, nil
 		} else if err == nil {
 			debugLog("[AppInsights] This is not azure public cloud:%s", cloudName)
-			return false, fmt.Errorf("Not an azure public cloud: %s", cloudName)
+			return false, errors.Errorf("not an azure public cloud: %s", cloudName)
 		}
 
 		debugLog("GetAzureCloud returned err :%v", err)
@@ -214,6 +215,46 @@ func NewAITelemetry(
 	return th, nil
 }
 
+// NewWithConnectionString creates telemetry handle with user specified appinsights connection string.
+func NewWithConnectionString(connectionString string, aiConfig AIConfig) (TelemetryHandle, error) {
+	debugMode = aiConfig.DebugMode
+
+	if connectionString == "" {
+		debugLog("Empty connection string")
+		return nil, errors.New("AI connection string is empty")
+	}
+
+	setAIConfigDefaults(&aiConfig)
+
+	connectionVars, err := parseConnectionString(connectionString)
+	if err != nil {
+		debugLog("Error parsing connection string: %v", err)
+		return nil, err
+	}
+
+	telemetryConfig := appinsights.NewTelemetryConfiguration(connectionVars.instrumentationKey)
+	telemetryConfig.EndpointUrl = connectionVars.ingestionURL
+	telemetryConfig.MaxBatchSize = aiConfig.BatchSize
+	telemetryConfig.MaxBatchInterval = time.Duration(aiConfig.BatchInterval) * time.Second
+
+	th := &telemetryHandle{
+		client:                       appinsights.NewTelemetryClientFromConfig(telemetryConfig),
+		appName:                      aiConfig.AppName,
+		appVersion:                   aiConfig.AppVersion,
+		diagListener:                 messageListener(),
+		disableMetadataRefreshThread: aiConfig.DisableMetadataRefreshThread,
+		refreshTimeout:               aiConfig.RefreshTimeout,
+	}
+
+	if th.disableMetadataRefreshThread {
+		getMetadata(th)
+	} else {
+		go getMetadata(th)
+	}
+
+	return th, nil
+}
+
 // TrackLog function sends report (trace) to appinsights resource. It overrides few of the existing columns with app information
 // and for rest it uses custom dimesion
 func (th *telemetryHandle) TrackLog(report Report) {
diff --git a/aitelemetry/telemetrywrapper_test.go b/aitelemetry/telemetrywrapper_test.go
@@ -15,7 +15,7 @@ import (
 )
 
 var (
-	th               TelemetryHandle
+	aiConfig         AIConfig
 	hostAgentUrl     = "localhost:3501"
 	getCloudResponse = "AzurePublicCloud"
 	httpURL          = "http://" + hostAgentUrl
@@ -54,6 +54,18 @@ func TestMain(m *testing.M) {
 		return
 	}
 
+	aiConfig = AIConfig{
+		AppName:                      "testapp",
+		AppVersion:                   "v1.0.26",
+		BatchSize:                    4096,
+		BatchInterval:                2,
+		RefreshTimeout:               10,
+		GetEnvRetryCount:             1,
+		GetEnvRetryWaitTimeInSecs:    2,
+		DebugMode:                    true,
+		DisableMetadataRefreshThread: true,
+	}
+
 	exitCode := m.Run()
 
 	if runtime.GOOS == "linux" {
@@ -75,67 +87,78 @@ func handleGetCloud(w http.ResponseWriter, req *http.Request) {
 	w.Write([]byte(getCloudResponse))
 }
 
+func initTelemetry(_ *testing.T) (th1, th2 TelemetryHandle) {
+	th1, err1 := NewAITelemetry(httpURL, "00ca2a73-c8d6-4929-a0c2-cf84545ec225", aiConfig)
+	if err1 != nil {
+		fmt.Printf("Error initializing AI telemetry: %v", err1)
+	}
+
+	th2, err2 := NewWithConnectionString(connectionString, aiConfig)
+	if err2 != nil {
+		fmt.Printf("Error initializing AI telemetry with connection string: %v", err2)
+	}
+
+	return
+}
+
 func TestEmptyAIKey(t *testing.T) {
 	var err error
 
-	aiConfig := AIConfig{
-		AppName:                      "testapp",
-		AppVersion:                   "v1.0.26",
-		BatchSize:                    4096,
-		BatchInterval:                2,
-		RefreshTimeout:               10,
-		DebugMode:                    true,
-		DisableMetadataRefreshThread: true,
-	}
 	_, err = NewAITelemetry(httpURL, "", aiConfig)
 	if err == nil {
-		t.Errorf("Error intializing AI telemetry:%v", err)
+		t.Errorf("Error initializing AI telemetry:%v", err)
+	}
+
+	_, err = NewWithConnectionString("", aiConfig)
+	if err == nil {
+		t.Errorf("Error initializing AI telemetry with connection string:%v", err)
 	}
 }
 
 func TestNewAITelemetry(t *testing.T) {
 	var err error
 
-	aiConfig := AIConfig{
-		AppName:                      "testapp",
-		AppVersion:                   "v1.0.26",
-		BatchSize:                    4096,
-		BatchInterval:                2,
-		RefreshTimeout:               10,
-		GetEnvRetryCount:             1,
-		GetEnvRetryWaitTimeInSecs:    2,
-		DebugMode:                    true,
-		DisableMetadataRefreshThread: true,
+	th1, th2 := initTelemetry(t)
+	if th1 == nil {
+		t.Errorf("Error initializing AI telemetry: %v", err)
 	}
-	th, err = NewAITelemetry(httpURL, "00ca2a73-c8d6-4929-a0c2-cf84545ec225", aiConfig)
-	if th == nil {
-		t.Errorf("Error intializing AI telemetry: %v", err)
+
+	if th2 == nil {
+		t.Errorf("Error initializing AI telemetry with connection string: %v", err)
 	}
 }
 
 func TestTrackMetric(t *testing.T) {
+	th1, th2 := initTelemetry(t)
+
 	metric := Metric{
 		Name:             "test",
 		Value:            1.0,
 		CustomDimensions: make(map[string]string),
 	}
 
 	metric.CustomDimensions["dim1"] = "col1"
-	th.TrackMetric(metric)
+	th1.TrackMetric(metric)
+	th2.TrackMetric(metric)
 }
 
 func TestTrackLog(t *testing.T) {
+	th1, th2 := initTelemetry(t)
+
 	report := Report{
 		Message:          "test",
 		Context:          "10a",
 		CustomDimensions: make(map[string]string),
 	}
 
 	report.CustomDimensions["dim1"] = "col1"
-	th.TrackLog(report)
+	th1.TrackLog(report)
+	th2.TrackLog(report)
 }
 
 func TestTrackEvent(t *testing.T) {
+	th1, th2 := initTelemetry(t)
+
 	event := Event{
 		EventName:  "testEvent",
 		ResourceID: "SomeResourceId",
@@ -144,35 +167,20 @@ func TestTrackEvent(t *testing.T) {
 
 	event.Properties["P1"] = "V1"
 	event.Properties["P2"] = "V2"
-	th.TrackEvent(event)
+	th1.TrackEvent(event)
+	th2.TrackEvent(event)
 }
 
 func TestFlush(t *testing.T) {
-	th.Flush()
-}
+	th1, th2 := initTelemetry(t)
 
-func TestClose(t *testing.T) {
-	th.Close(10)
+	th1.Flush()
+	th2.Flush()
 }
 
-func TestClosewithoutSend(t *testing.T) {
-	var err error
-
-	aiConfig := AIConfig{
-		AppName:                      "testapp",
-		AppVersion:                   "v1.0.26",
-		BatchSize:                    4096,
-		BatchInterval:                2,
-		DisableMetadataRefreshThread: true,
-		RefreshTimeout:               10,
-		GetEnvRetryCount:             1,
-		GetEnvRetryWaitTimeInSecs:    2,
-	}
-
-	thtest, err := NewAITelemetry(httpURL, "00ca2a73-c8d6-4929-a0c2-cf84545ec225", aiConfig)
-	if thtest == nil {
-		t.Errorf("Error intializing AI telemetry:%v", err)
-	}
+func TestClose(t *testing.T) {
+	th1, th2 := initTelemetry(t)
 
-	thtest.Close(10)
+	th1.Close(10)
+	th2.Close(10)
 }
diff --git a/go.mod b/go.mod
diff --git a/go.sum b/go.sum
