Add custom headers config for OpenSearch (#6248)

* Add custom headers config for opensearch

* Change custom headers config to a map

common/config/elasticsearch.go

@@ -55,6 +55,8 @@ type (
 		AWSSigning AWSSigning `yaml:"awsSigning"`
 		// optional to use Signed Certificates over https
 		TLS TLS `yaml:"tls"`
+		// optional to add custom headers
+		CustomHeaders map[string]string `yaml:"customHeaders,omitempty"`
 	}
 
 	// AWSSigning contains config to enable signing,
@@ -139,3 +141,11 @@ func (a AWSSigning) GetCredentials() (*credentials.Credentials, *string, error)
 
 	return awsCredentials, &a.StaticCredential.Region, nil
 }
+
+// GetCustomHeader returns the header for the specified key
+func (cfg *ElasticSearchConfig) GetCustomHeader(headerKey string) string {
+	if headerValue, ok := cfg.CustomHeaders[headerKey]; ok {
+		return headerValue
+	}
+	return ""
+}

common/config/elasticsearch_test.go

@@ -96,3 +96,37 @@ func TestAWSSigning_ValidateEmpty(t *testing.T) {
 	}
 
 }
+
+func TestGetCustomHeader(t *testing.T) {
+
+	tests := []struct {
+		config   ElasticSearchConfig
+		header   string
+		expected string
+	}{
+		{
+			config: ElasticSearchConfig{
+				CustomHeaders: map[string]string{
+					"key1": "value1",
+				},
+			},
+			header:   "key1",
+			expected: "value1",
+		},
+		{
+			config: ElasticSearchConfig{
+				CustomHeaders: map[string]string{
+					"key1": "value1",
+				},
+			},
+			header:   "key2",
+			expected: "",
+		},
+	}
+
+	for _, tc := range tests {
+		val := tc.config.GetCustomHeader(tc.header)
+		assert.Equal(t, val, tc.expected)
+	}
+
+}

common/elasticsearch/client/os2/client.go

@@ -103,6 +103,14 @@ func NewClient(
 		RetryBackoff: func(i int) time.Duration { return time.Duration(i) * 100 * time.Millisecond },
 	}
 
+	if len(connectConfig.CustomHeaders) > 0 {
+		osconfig.Header = http.Header{}
+
+		for key, value := range connectConfig.CustomHeaders {
+			osconfig.Header.Set(key, value)
+		}
+	}
+
 	// DiscoverNodesOnStart is false by default. Turn it on only when disable sniff is set to False in ES config
 	if !connectConfig.DisableSniff {
 		osconfig.DiscoverNodesOnStart = true

common/elasticsearch/client/os2/client_test.go

@@ -79,6 +79,9 @@ func TestNewClient(t *testing.T) {
 			config: &config.ElasticSearchConfig{
 				URL:          *url,
 				DisableSniff: false,
+				CustomHeaders: map[string]string{
+					"key": "value",
+				},
 			},
 			expectedErr: false,
 		},