From 386faf5958082d4aba519fcaa714157bbc21aa8a Mon Sep 17 00:00:00 2001
From: Sven Grossmann <sven.grossmann@grafana.com>
Date: Wed, 14 Dec 2022 20:22:26 +0100
Subject: [PATCH] OpenSearch: Use `aoss` servicename if OpenSearch is
 configured as `serverless` (#60344)

* Use `aoss` if opensearch is configured as `serverless`
---
 .../datasources/service/datasource.go         | 14 +++-
 .../datasources/service/datasource_test.go    | 66 +++++++++++++++++++
 2 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/pkg/services/datasources/service/datasource.go b/pkg/services/datasources/service/datasource.go
index 3b4bb78e01e..6d91b515e47 100644
--- a/pkg/services/datasources/service/datasource.go
+++ b/pkg/services/datasources/service/datasource.go
@@ -441,7 +441,7 @@ func (s *Service) httpClientOptions(ctx context.Context, ds *datasources.DataSou
 
 	if ds.JsonData != nil && ds.JsonData.Get("sigV4Auth").MustBool(false) && setting.SigV4AuthEnabled {
 		opts.SigV4 = &sdkhttpclient.SigV4Config{
-			Service:       awsServiceNamespace(ds.Type),
+			Service:       awsServiceNamespace(ds.Type, ds.JsonData),
 			Region:        ds.JsonData.Get("sigV4Region").MustString(),
 			AssumeRoleARN: ds.JsonData.Get("sigV4AssumeRoleArn").MustString(),
 			AuthType:      ds.JsonData.Get("sigV4AuthType").MustString(),
@@ -572,10 +572,18 @@ func (s *Service) getCustomHeaders(jsonData *simplejson.Json, decryptedValues ma
 	return headers
 }
 
-func awsServiceNamespace(dsType string) string {
+func awsServiceNamespace(dsType string, jsonData *simplejson.Json) string {
 	switch dsType {
-	case datasources.DS_ES, datasources.DS_ES_OPEN_DISTRO, datasources.DS_ES_OPENSEARCH:
+	case datasources.DS_ES, datasources.DS_ES_OPEN_DISTRO:
 		return "es"
+	case datasources.DS_ES_OPENSEARCH:
+		serverless := jsonData.Get("serverless").MustBool()
+
+		if serverless {
+			return "aoss"
+		} else {
+			return "es"
+		}
 	case datasources.DS_PROMETHEUS, datasources.DS_ALERTMANAGER:
 		return "aps"
 	default:
diff --git a/pkg/services/datasources/service/datasource_test.go b/pkg/services/datasources/service/datasource_test.go
index 2e3fdbe7068..440659001d4 100644
--- a/pkg/services/datasources/service/datasource_test.go
+++ b/pkg/services/datasources/service/datasource_test.go
@@ -180,6 +180,72 @@ func TestService_IDScopeResolver(t *testing.T) {
 	}
 }
 
+func TestService_awsServiceNamespace(t *testing.T) {
+	type testCaseResolver struct {
+		desc      string
+		givenDs   string
+		givenJson string
+		want      string
+		panic     bool
+	}
+
+	testCases := []testCaseResolver{
+		{
+			desc:      "elasticsearch",
+			givenDs:   datasources.DS_ES,
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "es",
+		}, {
+			desc:      "opendistro",
+			givenDs:   datasources.DS_ES_OPEN_DISTRO,
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "es",
+		}, {
+			desc:      "opensearch not serverless",
+			givenDs:   datasources.DS_ES_OPENSEARCH,
+			givenJson: `{ "sigV4Auth": true }`,
+			want:      "es",
+		}, {
+			desc:      "opensearch not serverless",
+			givenDs:   datasources.DS_ES_OPENSEARCH,
+			givenJson: `{ "sigV4Auth": true, "serverless": false }`,
+			want:      "es",
+		}, {
+			desc:      "opensearch serverless",
+			givenDs:   datasources.DS_ES_OPENSEARCH,
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "aoss",
+		}, {
+			desc:      "prometheus",
+			givenDs:   datasources.DS_PROMETHEUS,
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "aps",
+		}, {
+			desc:      "alertmanager",
+			givenDs:   datasources.DS_ALERTMANAGER,
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "aps",
+		}, {
+			desc:      "panic",
+			givenDs:   "panic",
+			givenJson: `{ "sigV4Auth": true, "serverless": true }`,
+			want:      "aps",
+			panic:     true,
+		},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			json, _ := simplejson.NewJson([]byte(tc.givenJson))
+			if tc.panic {
+				require.Panics(t, func() { awsServiceNamespace(tc.givenDs, json) })
+			} else {
+				resolved := awsServiceNamespace(tc.givenDs, json)
+				require.Equal(t, tc.want, resolved)
+			}
+		})
+	}
+}
+
 //nolint:goconst
 func TestService_GetHttpTransport(t *testing.T) {
 	cfg := &setting.Cfg{}