Chore: Collect elasticsearch version usage stats (#31787)

* Chore: Collect elasticsearch version usage stats

* Fix lint error

* use GetDataSources from sqlstore

* Apply review suggestions

* Return error if datasource type is not specified

* Update pkg/services/sqlstore/datasource.go

* fix undefined var

pkg/infra/usagestats/usage_stats.go

@@ -115,6 +115,25 @@ func (uss *UsageStatsService) GetUsageReport(ctx context.Context) (UsageReport,
 	}
 	metrics["stats.ds.other.count"] = dsOtherCount
 
+	esDataSourcesQuery := models.GetDataSourcesByTypeQuery{Type: models.DS_ES}
+	if err := uss.Bus.Dispatch(&esDataSourcesQuery); err != nil {
+		metricsLogger.Error("Failed to get elasticsearch json data", "error", err)
+		return report, err
+	}
+
+	for _, data := range esDataSourcesQuery.Result {
+		esVersion, err := data.JsonData.Get("esVersion").Int()
+		if err != nil {
+			continue
+		}
+
+		statName := fmt.Sprintf("stats.ds.elasticsearch.v%d.count", esVersion)
+
+		count, _ := metrics[statName].(int64)
+
+		metrics[statName] = count + 1
+	}
+
 	metrics["stats.packaging."+uss.Cfg.Packaging+".count"] = 1
 	metrics["stats.distributor."+uss.Cfg.ReportingDistributor+".count"] = 1
 

pkg/infra/usagestats/usage_stats_test.go

@@ -91,6 +91,29 @@ func TestMetrics(t *testing.T) {
 			return nil
 		})
 
+		var getESDatasSourcesQuery *models.GetDataSourcesByTypeQuery
+		uss.Bus.AddHandler(func(query *models.GetDataSourcesByTypeQuery) error {
+			query.Result = []*models.DataSource{
+				{
+					JsonData: simplejson.NewFromAny(map[string]interface{}{
+						"esVersion": 2,
+					}),
+				},
+				{
+					JsonData: simplejson.NewFromAny(map[string]interface{}{
+						"esVersion": 2,
+					}),
+				},
+				{
+					JsonData: simplejson.NewFromAny(map[string]interface{}{
+						"esVersion": 70,
+					}),
+				},
+			}
+			getESDatasSourcesQuery = query
+			return nil
+		})
+
 		var getDataSourceAccessStatsQuery *models.GetDataSourceAccessStatsQuery
 		uss.Bus.AddHandler(func(query *models.GetDataSourceAccessStatsQuery) error {
 			query.Result = []*models.DataSourceAccessStats{
@@ -190,6 +213,7 @@ func TestMetrics(t *testing.T) {
 			assert.Nil(t, getSystemStatsQuery)
 			assert.Nil(t, getDataSourceStatsQuery)
 			assert.Nil(t, getDataSourceAccessStatsQuery)
+			assert.Nil(t, getESDatasSourcesQuery)
 		})
 
 		t.Run("Given reporting enabled, stats should be gathered and sent to HTTP endpoint", func(t *testing.T) {
@@ -248,6 +272,7 @@ func TestMetrics(t *testing.T) {
 
 			assert.NotNil(t, getSystemStatsQuery)
 			assert.NotNil(t, getDataSourceStatsQuery)
+			assert.NotNil(t, getESDatasSourcesQuery)
 			assert.NotNil(t, getDataSourceAccessStatsQuery)
 			assert.NotNil(t, getAlertNotifierUsageStatsQuery)
 			assert.NotNil(t, resp.req)
@@ -289,6 +314,10 @@ func TestMetrics(t *testing.T) {
 
 			assert.Equal(t, 9, metrics.Get("stats.ds."+models.DS_ES+".count").MustInt())
 			assert.Equal(t, 10, metrics.Get("stats.ds."+models.DS_PROMETHEUS+".count").MustInt())
+
+			assert.Equal(t, 2, metrics.Get("stats.ds."+models.DS_ES+".v2.count").MustInt())
+			assert.Equal(t, 1, metrics.Get("stats.ds."+models.DS_ES+".v70.count").MustInt())
+
 			assert.Equal(t, 11+12, metrics.Get("stats.ds.other.count").MustInt())
 
 			assert.Equal(t, 1, metrics.Get("stats.ds_access."+models.DS_ES+".direct.count").MustInt())
@@ -421,6 +450,11 @@ func TestMetrics(t *testing.T) {
 			return nil
 		})
 
+		uss.Bus.AddHandler(func(query *models.GetDataSourcesByTypeQuery) error {
+			query.Result = []*models.DataSource{}
+			return nil
+		})
+
 		uss.Bus.AddHandler(func(query *models.GetDataSourceAccessStatsQuery) error {
 			query.Result = []*models.DataSourceAccessStats{}
 			return nil

pkg/models/datasource.go

@@ -154,6 +154,11 @@ type GetDataSourcesQuery struct {
 	Result          []*DataSource
 }
 
+type GetDataSourcesByTypeQuery struct {
+	Type   string
+	Result []*DataSource
+}
+
 type GetDefaultDataSourceQuery struct {
 	OrgId  int64
 	User   *SignedInUser

pkg/services/sqlstore/datasource.go

@@ -1,6 +1,7 @@
 package sqlstore
 
 import (
+	"fmt"
 	"strings"
 	"time"
 
@@ -18,6 +19,7 @@ import (
 
 func init() {
 	bus.AddHandler("sql", GetDataSources)
+	bus.AddHandler("sql", GetDataSourcesByType)
 	bus.AddHandler("sql", GetDataSource)
 	bus.AddHandler("sql", AddDataSource)
 	bus.AddHandler("sql", DeleteDataSource)
@@ -71,10 +73,21 @@ func GetDataSources(query *models.GetDataSourcesQuery) error {
 	} else {
 		sess = x.Limit(query.DataSourceLimit, 0).Where("org_id=?", query.OrgId).Asc("name")
 	}
+
 	query.Result = make([]*models.DataSource, 0)
 	return sess.Find(&query.Result)
 }
 
+// GetDataSourcesByType returns all datasources for a given type or an error if the specified type is an empty string
+func GetDataSourcesByType(query *models.GetDataSourcesByTypeQuery) error {
+	if query.Type == "" {
+		return fmt.Errorf("datasource type cannot be empty")
+	}
+
+	query.Result = make([]*models.DataSource, 0)
+	return x.Where("type=?", query.Type).Asc("id").Find(&query.Result)
+}
+
 // GetDefaultDataSource is used to get the default datasource of organization
 func GetDefaultDataSource(query *models.GetDefaultDataSourceQuery) error {
 	datasource := models.DataSource{}

pkg/services/sqlstore/datasource_test.go

@@ -218,7 +218,7 @@ func TestDataAccess(t *testing.T) {
 		require.Equal(t, 0, len(query.Result))
 	})
 
-	t.Run("GetDataSource", func(t *testing.T) {
+	t.Run("GetDataSources", func(t *testing.T) {
 		t.Run("Number of data sources returned limited to 6 per organization", func(t *testing.T) {
 			InitTestDB(t)
 			datasourceLimit := 6
@@ -288,6 +288,49 @@ func TestDataAccess(t *testing.T) {
 			require.Equal(t, numberOfDatasource, len(query.Result))
 		})
 	})
+
+	t.Run("GetDataSourcesByType", func(t *testing.T) {
+		t.Run("Only returns datasources of specified type", func(t *testing.T) {
+			InitTestDB(t)
+
+			err := AddDataSource(&models.AddDataSourceCommand{
+				OrgId:    10,
+				Name:     "Elasticsearch",
+				Type:     models.DS_ES,
+				Access:   models.DS_ACCESS_DIRECT,
+				Url:      "http://test",
+				Database: "site",
+				ReadOnly: true,
+			})
+			require.NoError(t, err)
+
+			err = AddDataSource(&models.AddDataSourceCommand{
+				OrgId:    10,
+				Name:     "Graphite",
+				Type:     models.DS_GRAPHITE,
+				Access:   models.DS_ACCESS_DIRECT,
+				Url:      "http://test",
+				Database: "site",
+				ReadOnly: true,
+			})
+			require.NoError(t, err)
+
+			query := models.GetDataSourcesByTypeQuery{Type: models.DS_ES}
+
+			err = GetDataSourcesByType(&query)
+
+			require.NoError(t, err)
+			require.Equal(t, 1, len(query.Result))
+		})
+
+		t.Run("Returns an error if no type specified", func(t *testing.T) {
+			query := models.GetDataSourcesByTypeQuery{}
+
+			err := GetDataSourcesByType(&query)
+
+			require.Error(t, err)
+		})
+	})
 }
 
 func TestGetDefaultDataSource(t *testing.T) {