SQLStore: customise the limit of retrieved datasources per organisation (#29358)

* SQLStore: customise the limit of retrieved datasources per organisation

* update all suggestions regarding nil or 0 as default

* Apply suggestions from code review

Co-authored-by: Emil Tullstedt <emil.tullstedt@grafana.com>

* correct default.ini description + adding unittest

* Apply suggestions from code review

Co-authored-by: Sofia Papagiannaki <papagian@users.noreply.github.com>

* modify unittest name

Co-authored-by: Emil Tullstedt <emil.tullstedt@grafana.com>
Co-authored-by: Sofia Papagiannaki <papagian@users.noreply.github.com>

conf/defaults.ini

@@ -255,6 +255,11 @@ min_refresh_interval = 5s
 # Path to the default home dashboard. If this value is empty, then Grafana uses StaticRootPath + "dashboards/home.json"
 default_home_dashboard_path =
 
+################################### Data sources #########################
+[datasources]
+# Upper limit of data sources that Grafana will return. This limit is a temporary configuration and it will be deprecated when pagination will be introduced on the list data sources API.
+datasource_limit = 5000
+
 #################################### Users ###############################
 [users]
 # disable user signup / registration

conf/sample.ini

@@ -110,6 +110,11 @@
 # For "sqlite3" only. cache mode setting used for connecting to the database. (private, shared)
 ;cache_mode = private
 
+################################### Data sources #########################
+[datasources]
+# Upper limit of data sources that Grafana will return. This limit is a temporary configuration and it will be deprecated when pagination will be introduced on the list data sources API.
+;datasource_limit = 5000
+
 #################################### Cache server #############################
 [remote_cache]
 # Either "redis", "memcached" or "database" default is "database"

pkg/api/api.go

@@ -249,7 +249,7 @@ func (hs *HTTPServer) registerRoutes() {
 
 		// Data sources
 		apiRoute.Group("/datasources", func(datasourceRoute routing.RouteRegister) {
-			datasourceRoute.Get("/", Wrap(GetDataSources))
+			datasourceRoute.Get("/", Wrap(hs.GetDataSources))
 			datasourceRoute.Post("/", quota("data_source"), bind(models.AddDataSourceCommand{}), Wrap(AddDataSource))
 			datasourceRoute.Put("/:id", bind(models.UpdateDataSourceCommand{}), Wrap(UpdateDataSource))
 			datasourceRoute.Delete("/:id", Wrap(DeleteDataSourceById))

pkg/api/datasources.go

@@ -19,8 +19,8 @@ import (
 
 var datasourcesLogger = log.New("datasources")
 
-func GetDataSources(c *models.ReqContext) Response {
-	query := models.GetDataSourcesQuery{OrgId: c.OrgId}
+func (hs *HTTPServer) GetDataSources(c *models.ReqContext) Response {
+	query := models.GetDataSourcesQuery{OrgId: c.OrgId, DataSourceLimit: hs.Cfg.DataSourceLimit}
 
 	if err := bus.Dispatch(&query); err != nil {
 		return Error(500, "Failed to query datasources", err)

pkg/api/datasources_test.go

@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/setting"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
@@ -32,7 +33,8 @@ func TestDataSourcesProxy_userLoggedIn(t *testing.T) {
 		})
 
 		// handler func being tested
-		sc.handlerFunc = GetDataSources
+		hs := &HTTPServer{Bus: bus.GetBus(), Cfg: setting.NewCfg()}
+		sc.handlerFunc = hs.GetDataSources
 		sc.fakeReq("GET", "/api/datasources").exec()
 
 		respJSON := []map[string]interface{}{}

pkg/api/frontendsettings.go

@@ -15,11 +15,11 @@ import (
 	"github.com/grafana/grafana/pkg/setting"
 )
 
-func getFSDataSources(c *models.ReqContext, enabledPlugins *plugins.EnabledPlugins) (map[string]interface{}, error) {
+func (hs *HTTPServer) getFSDataSources(c *models.ReqContext, enabledPlugins *plugins.EnabledPlugins) (map[string]interface{}, error) {
 	orgDataSources := make([]*models.DataSource, 0)
 
 	if c.OrgId != 0 {
-		query := models.GetDataSourcesQuery{OrgId: c.OrgId}
+		query := models.GetDataSourcesQuery{OrgId: c.OrgId, DataSourceLimit: hs.Cfg.DataSourceLimit}
 		err := bus.Dispatch(&query)
 
 		if err != nil {
@@ -135,7 +135,7 @@ func (hs *HTTPServer) getFrontendSettingsMap(c *models.ReqContext) (map[string]i
 		}
 	}
 
-	dataSources, err := getFSDataSources(c, enabledPlugins)
+	dataSources, err := hs.getFSDataSources(c, enabledPlugins)
 	if err != nil {
 		return nil, err
 	}

pkg/models/datasource.go

@@ -202,9 +202,10 @@ type DeleteDataSourceByNameCommand struct {
 // QUERIES
 
 type GetDataSourcesQuery struct {
-	OrgId  int64
-	User   *SignedInUser
-	Result []*DataSource
+	OrgId           int64
+	DataSourceLimit int
+	User            *SignedInUser
+	Result          []*DataSource
 }
 
 type GetDataSourceByIdQuery struct {

pkg/services/sqlstore/datasource.go

@@ -67,8 +67,12 @@ func GetDataSourceByName(query *models.GetDataSourceByNameQuery) error {
 }
 
 func GetDataSources(query *models.GetDataSourcesQuery) error {
-	sess := x.Limit(5000, 0).Where("org_id=?", query.OrgId).Asc("name")
-
+	var sess *xorm.Session
+	if query.DataSourceLimit <= 0 {
+		sess = x.Where("org_id=?", query.OrgId).Asc("name")
+	} 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)
 }

pkg/services/sqlstore/datasource_test.go

@@ -3,6 +3,7 @@
 package sqlstore
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/grafana/grafana/pkg/models"
@@ -214,4 +215,75 @@ func TestDataAccess(t *testing.T) {
 
 		require.Equal(t, 0, len(query.Result))
 	})
+
+	t.Run("GetDataSource", func(t *testing.T) {
+		t.Run("Number of data sources returned limited to 6 per organization", func(t *testing.T) {
+			InitTestDB(t)
+			datasourceLimit := 6
+			for i := 0; i < datasourceLimit+1; i++ {
+				err := AddDataSource(&models.AddDataSourceCommand{
+					OrgId:    10,
+					Name:     "laban" + strconv.Itoa(i),
+					Type:     models.DS_GRAPHITE,
+					Access:   models.DS_ACCESS_DIRECT,
+					Url:      "http://test",
+					Database: "site",
+					ReadOnly: true,
+				})
+				require.NoError(t, err)
+			}
+			query := models.GetDataSourcesQuery{OrgId: 10, DataSourceLimit: datasourceLimit}
+
+			err := GetDataSources(&query)
+
+			require.NoError(t, err)
+			require.Equal(t, datasourceLimit, len(query.Result))
+		})
+
+		t.Run("No limit should be applied on the returned data sources if the limit is not set", func(t *testing.T) {
+			InitTestDB(t)
+			numberOfDatasource := 5100
+			for i := 0; i < numberOfDatasource; i++ {
+				err := AddDataSource(&models.AddDataSourceCommand{
+					OrgId:    10,
+					Name:     "laban" + strconv.Itoa(i),
+					Type:     models.DS_GRAPHITE,
+					Access:   models.DS_ACCESS_DIRECT,
+					Url:      "http://test",
+					Database: "site",
+					ReadOnly: true,
+				})
+				require.NoError(t, err)
+			}
+			query := models.GetDataSourcesQuery{OrgId: 10}
+
+			err := GetDataSources(&query)
+
+			require.NoError(t, err)
+			require.Equal(t, numberOfDatasource, len(query.Result))
+		})
+
+		t.Run("No limit should be applied on the returned data sources if the limit is negative", func(t *testing.T) {
+			InitTestDB(t)
+			numberOfDatasource := 5100
+			for i := 0; i < numberOfDatasource; i++ {
+				err := AddDataSource(&models.AddDataSourceCommand{
+					OrgId:    10,
+					Name:     "laban" + strconv.Itoa(i),
+					Type:     models.DS_GRAPHITE,
+					Access:   models.DS_ACCESS_DIRECT,
+					Url:      "http://test",
+					Database: "site",
+					ReadOnly: true,
+				})
+				require.NoError(t, err)
+			}
+			query := models.GetDataSourcesQuery{OrgId: 10, DataSourceLimit: -1}
+
+			err := GetDataSources(&query)
+
+			require.NoError(t, err)
+			require.Equal(t, numberOfDatasource, len(query.Result))
+		})
+	})
 }

pkg/services/sqlstore/sqlstore.go

@@ -354,7 +354,7 @@ func (ss *SQLStore) readConfig() {
 	ss.dbCfg.CacheMode = sec.Key("cache_mode").MustString("private")
 }
 
-// Interface of arguments for testing db
+// ITestDB is an interface of arguments for testing db
 type ITestDB interface {
 	Helper()
 	Fatalf(format string, args ...interface{})

pkg/setting/setting.go

@@ -314,6 +314,9 @@ type Cfg struct {
 	// Sentry config
 	Sentry Sentry
 
+	// Data sources
+	DataSourceLimit int
+
 	// Snapshots
 	SnapshotPublicMode bool
 
@@ -826,6 +829,8 @@ func (cfg *Cfg) Load(args *CommandLineArgs) error {
 		return err
 	}
 
+	cfg.readDataSourcesSettings()
+
 	if VerifyEmailEnabled && !cfg.Smtp.Enabled {
 		log.Warnf("require_email_validation is enabled but smtp is disabled")
 	}
@@ -1255,3 +1260,8 @@ func readServerSettings(iniFile *ini.File, cfg *Cfg) error {
 
 	return nil
 }
+
+func (cfg *Cfg) readDataSourcesSettings() {
+	datasources := cfg.Raw.Section("datasources")
+	cfg.DataSourceLimit = datasources.Key("datasource_limit").MustInt(5000)
+}