Chore: Propagate context for data source provisioning (#40235)

* context all the things

* apply feedback

* rollback some alerting changes

* rollback some alerting changes #2

* more rollbacks

* more rollbacks #2

* more rollbacks #3

* more rollbacks #4

* fix integration test

* add missing context

* add missing and remove incorrect dispatch

pkg/api/admin_provisioning.go

@@ -17,7 +17,7 @@ func (hs *HTTPServer) AdminProvisioningReloadDashboards(c *models.ReqContext) re
 }
 
 func (hs *HTTPServer) AdminProvisioningReloadDatasources(c *models.ReqContext) response.Response {
-	err := hs.ProvisioningService.ProvisionDatasources()
+	err := hs.ProvisioningService.ProvisionDatasources(c.Req.Context())
 	if err != nil {
 		return response.Error(500, "", err)
 	}

pkg/api/datasources.go

@@ -70,7 +70,7 @@ func GetDataSourceById(c *models.ReqContext) response.Response {
 		OrgId: c.OrgId,
 	}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
 		}
@@ -119,7 +119,7 @@ func (hs *HTTPServer) DeleteDataSourceById(c *models.ReqContext) response.Respon
 
 // GET /api/datasources/uid/:uid
 func GetDataSourceByUID(c *models.ReqContext) response.Response {
-	ds, err := getRawDataSourceByUID(web.Params(c.Req)[":uid"], c.OrgId)
+	ds, err := getRawDataSourceByUID(c.Req.Context(), web.Params(c.Req)[":uid"], c.OrgId)
 
 	if err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
@@ -140,7 +140,7 @@ func (hs *HTTPServer) DeleteDataSourceByUID(c *models.ReqContext) response.Respo
 		return response.Error(400, "Missing datasource uid", nil)
 	}
 
-	ds, err := getRawDataSourceByUID(uid, c.OrgId)
+	ds, err := getRawDataSourceByUID(c.Req.Context(), uid, c.OrgId)
 	if err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
@@ -154,7 +154,7 @@ func (hs *HTTPServer) DeleteDataSourceByUID(c *models.ReqContext) response.Respo
 
 	cmd := &models.DeleteDataSourceCommand{UID: uid, OrgID: c.OrgId}
 
-	err = bus.Dispatch(cmd)
+	err = bus.DispatchCtx(c.Req.Context(), cmd)
 	if err != nil {
 		return response.Error(500, "Failed to delete datasource", err)
 	}
@@ -172,7 +172,7 @@ func (hs *HTTPServer) DeleteDataSourceByName(c *models.ReqContext) response.Resp
 	}
 
 	getCmd := &models.GetDataSourceQuery{Name: name, OrgId: c.OrgId}
-	if err := bus.Dispatch(getCmd); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), getCmd); err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
 		}
@@ -184,7 +184,7 @@ func (hs *HTTPServer) DeleteDataSourceByName(c *models.ReqContext) response.Resp
 	}
 
 	cmd := &models.DeleteDataSourceCommand{Name: name, OrgID: c.OrgId}
-	err := bus.Dispatch(cmd)
+	err := bus.DispatchCtx(c.Req.Context(), cmd)
 	if err != nil {
 		return response.Error(500, "Failed to delete datasource", err)
 	}
@@ -216,7 +216,7 @@ func AddDataSource(c *models.ReqContext, cmd models.AddDataSourceCommand) respon
 		return resp
 	}
 
-	if err := bus.Dispatch(&cmd); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &cmd); err != nil {
 		if errors.Is(err, models.ErrDataSourceNameExists) || errors.Is(err, models.ErrDataSourceUidExists) {
 			return response.Error(409, err.Error(), err)
 		}
@@ -246,7 +246,7 @@ func (hs *HTTPServer) UpdateDataSource(c *models.ReqContext, cmd models.UpdateDa
 		return response.Error(500, "Failed to update datasource", err)
 	}
 
-	err = bus.Dispatch(&cmd)
+	err = bus.DispatchCtx(c.Req.Context(), &cmd)
 	if err != nil {
 		if errors.Is(err, models.ErrDataSourceUpdatingOldVersion) {
 			return response.Error(409, "Datasource has already been updated by someone else. Please reload and try again", err)
@@ -259,7 +259,7 @@ func (hs *HTTPServer) UpdateDataSource(c *models.ReqContext, cmd models.UpdateDa
 		OrgId: c.OrgId,
 	}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
 		}
@@ -319,13 +319,13 @@ func getRawDataSourceById(ctx context.Context, id int64, orgID int64) (*models.D
 	return query.Result, nil
 }
 
-func getRawDataSourceByUID(uid string, orgID int64) (*models.DataSource, error) {
+func getRawDataSourceByUID(ctx context.Context, uid string, orgID int64) (*models.DataSource, error) {
 	query := models.GetDataSourceQuery{
 		Uid:   uid,
 		OrgId: orgID,
 	}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(ctx, &query); err != nil {
 		return nil, err
 	}
 
@@ -336,7 +336,7 @@ func getRawDataSourceByUID(uid string, orgID int64) (*models.DataSource, error)
 func GetDataSourceByName(c *models.ReqContext) response.Response {
 	query := models.GetDataSourceQuery{Name: web.Params(c.Req)[":name"], OrgId: c.OrgId}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
 		}
@@ -351,7 +351,7 @@ func GetDataSourceByName(c *models.ReqContext) response.Response {
 func GetDataSourceIdByName(c *models.ReqContext) response.Response {
 	query := models.GetDataSourceQuery{Name: web.Params(c.Req)[":name"], OrgId: c.OrgId}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			return response.Error(404, "Data source not found", nil)
 		}

pkg/expr/transform.go

@@ -178,7 +178,7 @@ func (s *Service) queryData(ctx context.Context, req *backend.QueryDataRequest)
 		Uid:   datasourceUID,
 	}
 
-	if err := bus.Dispatch(getDsInfo); err != nil {
+	if err := bus.DispatchCtx(ctx, getDsInfo); err != nil {
 		return nil, fmt.Errorf("could not find datasource: %w", err)
 	}
 

pkg/expr/translate/translate.go

@@ -1,6 +1,7 @@
 package translate
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"sort"
@@ -188,7 +189,7 @@ func (dc *dashConditionsJSON) GetNew(orgID int64) (*ngmodels.Condition, error) {
 				Id:    dc.Conditions[condIdx].Query.DatasourceID,
 			}
 
-			if err := bus.Dispatch(getDsInfo); err != nil {
+			if err := bus.DispatchCtx(context.TODO(), getDsInfo); err != nil {
 				return nil, fmt.Errorf("could not find datasource: %w", err)
 			}
 

pkg/expr/translate/translate_test.go

@@ -1,6 +1,7 @@
 package translate
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -132,7 +133,7 @@ func alertRuleByRefId(cond *ngmodels.Condition, refID string) (ngmodels.AlertQue
 }
 
 func registerGetDsInfoHandler() {
-	bus.AddHandler("test", func(query *models.GetDataSourceQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDataSourceQuery) error {
 		switch {
 		case query.Id == 2:
 			query.Result = &models.DataSource{Id: 2, OrgId: 1, Uid: "000000002"}

pkg/infra/usagestats/service/usage_stats.go

@@ -129,7 +129,7 @@ func (uss *UsageStats) GetUsageReport(ctx context.Context) (usagestats.Report, e
 	metrics["stats.ds.other.count"] = dsOtherCount
 
 	esDataSourcesQuery := models.GetDataSourcesByTypeQuery{Type: models.DS_ES}
-	if err := uss.Bus.DispatchCtx(ctx, &esDataSourcesQuery); err != nil {
+	if err := uss.Bus.Dispatch(&esDataSourcesQuery); err != nil {
 		uss.log.Error("Failed to get elasticsearch json data", "error", err)
 		return report, err
 	}

pkg/server/server.go

@@ -124,7 +124,7 @@ func (s *Server) init() error {
 		return err
 	}
 
-	return s.provisioningService.RunInitProvisioners()
+	return s.provisioningService.RunInitProvisioners(s.context)
 }
 
 // Run initializes and starts services. This will block until all services have

pkg/services/alerting/alerting_usage.go

@@ -1,6 +1,7 @@
 package alerting
 
 import (
+	"context"
 	"encoding/json"
 
 	"github.com/grafana/grafana/pkg/models"
@@ -62,7 +63,7 @@ func (e *AlertEngine) mapRulesToUsageStats(rules []*models.Alert) (DatasourceAle
 	result := map[string]int{}
 	for k, v := range typeCount {
 		query := &models.GetDataSourceQuery{Id: k}
-		err := e.Bus.Dispatch(query)
+		err := e.Bus.DispatchCtx(context.TODO(), query)
 		if err != nil {
 			return map[string]int{}, nil
 		}

pkg/services/alerting/alerting_usage_test.go

@@ -1,6 +1,7 @@
 package alerting
 
 import (
+	"context"
 	"encoding/json"
 	"io/ioutil"
 	"testing"
@@ -38,7 +39,7 @@ func TestAlertingUsageStats(t *testing.T) {
 		return nil
 	})
 
-	ae.Bus.AddHandler(func(query *models.GetDataSourceQuery) error {
+	ae.Bus.AddHandlerCtx(func(ctx context.Context, query *models.GetDataSourceQuery) error {
 		ds := map[int64]*models.DataSource{
 			1: {Type: "influxdb"},
 			2: {Type: "graphite"},

pkg/services/alerting/conditions/query.go

@@ -140,7 +140,7 @@ func (c *QueryCondition) executeQuery(context *alerting.EvalContext, timeRange p
 		OrgId: context.Rule.OrgID,
 	}
 
-	if err := bus.Dispatch(getDsInfo); err != nil {
+	if err := bus.DispatchCtx(context.Ctx, getDsInfo); err != nil {
 		return nil, fmt.Errorf("could not find datasource: %w", err)
 	}
 

pkg/services/alerting/conditions/query_interval_test.go

@@ -137,13 +137,14 @@ func (rh fakeIntervalTestReqHandler) HandleRequest(ctx context.Context, dsInfo *
 //nolint: staticcheck // plugins.DataResponse deprecated
 func applyScenario(timeRange string, dataSourceJsonData *simplejson.Json, queryModel string, verifier func(query plugins.DataSubQuery)) {
 	Convey("desc", func() {
-		bus.AddHandler("test", func(query *models.GetDataSourceQuery) error {
+		bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDataSourceQuery) error {
 			query.Result = &models.DataSource{Id: 1, Type: "graphite", JsonData: dataSourceJsonData}
 			return nil
 		})
 
 		ctx := &queryIntervalTestContext{}
 		ctx.result = &alerting.EvalContext{
+			Ctx:              context.Background(),
 			Rule:             &alerting.Rule{},
 			RequestValidator: &validations.OSSPluginRequestValidator{},
 		}

pkg/services/alerting/conditions/query_test.go

@@ -241,13 +241,14 @@ func (rh fakeReqHandler) HandleRequest(context.Context, *models.DataSource, plug
 
 func queryConditionScenario(desc string, fn queryConditionScenarioFunc) {
 	Convey(desc, func() {
-		bus.AddHandler("test", func(query *models.GetDataSourceQuery) error {
+		bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDataSourceQuery) error {
 			query.Result = &models.DataSource{Id: 1, Type: "graphite"}
 			return nil
 		})
 
 		ctx := &queryConditionTestContext{}
 		ctx.result = &alerting.EvalContext{
+			Ctx:              context.Background(),
 			Rule:             &alerting.Rule{},
 			RequestValidator: &validations.OSSPluginRequestValidator{},
 		}

pkg/services/alerting/extractor.go

@@ -1,6 +1,7 @@
 package alerting
 
 import (
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -32,14 +33,14 @@ func NewDashAlertExtractor(dash *models.Dashboard, orgID int64, user *models.Sig
 func (e *DashAlertExtractor) lookupDatasourceID(dsName string) (*models.DataSource, error) {
 	if dsName == "" {
 		query := &models.GetDefaultDataSourceQuery{OrgId: e.OrgID}
-		if err := bus.Dispatch(query); err != nil {
+		if err := bus.DispatchCtx(context.TODO(), query); err != nil {
 			return nil, err
 		}
 		return query.Result, nil
 	}
 
 	query := &models.GetDataSourceQuery{Name: dsName, OrgId: e.OrgID}
-	if err := bus.Dispatch(query); err != nil {
+	if err := bus.DispatchCtx(context.TODO(), query); err != nil {
 		return nil, err
 	}
 

pkg/services/dashboards/dashboard_service.go

@@ -270,7 +270,8 @@ func (dr *dashboardServiceImpl) SaveFolderForProvisionedDashboards(dto *SaveDash
 	return dash, nil
 }
 
-func (dr *dashboardServiceImpl) SaveDashboard(dto *SaveDashboardDTO, allowUiUpdate bool) (*models.Dashboard, error) {
+func (dr *dashboardServiceImpl) SaveDashboard(dto *SaveDashboardDTO,
+	allowUiUpdate bool) (*models.Dashboard, error) {
 	if err := validateDashboardRefreshInterval(dto.Dashboard); err != nil {
 		dr.log.Warn("Changing refresh interval for imported dashboard to minimum refresh interval",
 			"dashboardUid", dto.Dashboard.Uid, "dashboardTitle", dto.Dashboard.Title, "minRefreshInterval",

pkg/services/datasources/cache.go

@@ -1,6 +1,7 @@
 package datasources
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -45,7 +46,7 @@ func (dc *CacheServiceImpl) GetDatasource(
 	plog.Debug("Querying for data source via SQL store", "id", datasourceID, "orgId", user.OrgId)
 
 	query := &models.GetDataSourceQuery{Id: datasourceID, OrgId: user.OrgId}
-	err := dc.SQLStore.GetDataSource(query)
+	err := dc.SQLStore.GetDataSource(context.TODO(), query)
 	if err != nil {
 		return nil, err
 	}
@@ -83,7 +84,7 @@ func (dc *CacheServiceImpl) GetDatasourceByUID(
 
 	plog.Debug("Querying for data source via SQL store", "uid", datasourceUID, "orgId", user.OrgId)
 	query := &models.GetDataSourceQuery{Uid: datasourceUID, OrgId: user.OrgId}
-	err := dc.SQLStore.GetDataSource(query)
+	err := dc.SQLStore.GetDataSource(context.TODO(), query)
 	if err != nil {
 		return nil, err
 	}

pkg/services/datasources/service.go

@@ -64,17 +64,17 @@ func ProvideService(bus bus.Bus, store *sqlstore.SQLStore, encryptionService enc
 
 	s.Bus.AddHandler(s.GetDataSources)
 	s.Bus.AddHandler(s.GetDataSourcesByType)
-	s.Bus.AddHandler(s.GetDataSource)
+	s.Bus.AddHandlerCtx(s.GetDataSource)
 	s.Bus.AddHandlerCtx(s.AddDataSource)
-	s.Bus.AddHandler(s.DeleteDataSource)
+	s.Bus.AddHandlerCtx(s.DeleteDataSource)
 	s.Bus.AddHandlerCtx(s.UpdateDataSource)
 	s.Bus.AddHandler(s.GetDefaultDataSource)
 
 	return s
 }
 
-func (s *Service) GetDataSource(query *models.GetDataSourceQuery) error {
+func (s *Service) GetDataSource(ctx context.Context, query *models.GetDataSourceQuery) error {
-	return s.SQLStore.GetDataSource(query)
+	return s.SQLStore.GetDataSource(ctx, query)
 }
 
 func (s *Service) GetDataSources(query *models.GetDataSourcesQuery) error {
@@ -92,11 +92,11 @@ func (s *Service) AddDataSource(ctx context.Context, cmd *models.AddDataSourceCo
 		return err
 	}
 
-	return s.SQLStore.AddDataSource(cmd)
+	return s.SQLStore.AddDataSource(ctx, cmd)
 }
 
-func (s *Service) DeleteDataSource(cmd *models.DeleteDataSourceCommand) error {
+func (s *Service) DeleteDataSource(ctx context.Context, cmd *models.DeleteDataSourceCommand) error {
-	return s.SQLStore.DeleteDataSource(cmd)
+	return s.SQLStore.DeleteDataSource(ctx, cmd)
 }
 
 func (s *Service) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSourceCommand) error {
@@ -106,7 +106,7 @@ func (s *Service) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSo
 		return err
 	}
 
-	return s.SQLStore.UpdateDataSource(cmd)
+	return s.SQLStore.UpdateDataSource(ctx, cmd)
 }
 
 func (s *Service) GetDefaultDataSource(query *models.GetDefaultDataSourceQuery) error {

pkg/services/provisioning/datasources/config_reader_test.go

@@ -1,6 +1,7 @@
 package datasources
 
 import (
+	"context"
 	"os"
 	"testing"
 
@@ -39,7 +40,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 		Convey("apply default values when missing", func() {
 			dc := newDatasourceProvisioner(logger)
-			err := dc.applyChanges(withoutDefaults)
+			err := dc.applyChanges(context.Background(), withoutDefaults)
 			if err != nil {
 				t.Fatalf("applyChanges return an error %v", err)
 			}
@@ -52,7 +53,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 		Convey("One configured datasource", func() {
 			Convey("no datasource in database", func() {
 				dc := newDatasourceProvisioner(logger)
-				err := dc.applyChanges(twoDatasourcesConfig)
+				err := dc.applyChanges(context.Background(), twoDatasourcesConfig)
 				if err != nil {
 					t.Fatalf("applyChanges return an error %v", err)
 				}
@@ -69,7 +70,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 				Convey("should update one datasource", func() {
 					dc := newDatasourceProvisioner(logger)
-					err := dc.applyChanges(twoDatasourcesConfig)
+					err := dc.applyChanges(context.Background(), twoDatasourcesConfig)
 					if err != nil {
 						t.Fatalf("applyChanges return an error %v", err)
 					}
@@ -82,7 +83,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 			Convey("Two datasources with is_default", func() {
 				dc := newDatasourceProvisioner(logger)
-				err := dc.applyChanges(doubleDatasourcesConfig)
+				err := dc.applyChanges(context.Background(), doubleDatasourcesConfig)
 				Convey("should raise error", func() {
 					So(err, ShouldEqual, ErrInvalidConfigToManyDefault)
 				})
@@ -91,7 +92,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 		Convey("Multiple datasources in different organizations with isDefault in each organization", func() {
 			dc := newDatasourceProvisioner(logger)
-			err := dc.applyChanges(multipleOrgsWithDefault)
+			err := dc.applyChanges(context.Background(), multipleOrgsWithDefault)
 			Convey("should not raise error", func() {
 				So(err, ShouldBeNil)
 				So(len(fakeRepo.inserted), ShouldEqual, 4)
@@ -111,7 +112,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 				Convey("should have two new datasources", func() {
 					dc := newDatasourceProvisioner(logger)
-					err := dc.applyChanges(twoDatasourcesConfigPurgeOthers)
+					err := dc.applyChanges(context.Background(), twoDatasourcesConfigPurgeOthers)
 					if err != nil {
 						t.Fatalf("applyChanges return an error %v", err)
 					}
@@ -132,7 +133,7 @@ func TestDatasourceAsConfig(t *testing.T) {
 
 				Convey("should have two new datasources", func() {
 					dc := newDatasourceProvisioner(logger)
-					err := dc.applyChanges(twoDatasourcesConfig)
+					err := dc.applyChanges(context.Background(), twoDatasourcesConfig)
 					if err != nil {
 						t.Fatalf("applyChanges return an error %v", err)
 					}

pkg/services/provisioning/datasources/datasources.go

@@ -1,6 +1,7 @@
 package datasources
 
 import (
+	"context"
 	"errors"
 
 	"github.com/grafana/grafana/pkg/bus"
@@ -18,9 +19,9 @@ var (
 
 // Provision scans a directory for provisioning config files
 // and provisions the datasource in those files.
-func Provision(configDirectory string) error {
+func Provision(ctx context.Context, configDirectory string) error {
 	dc := newDatasourceProvisioner(log.New("provisioning.datasources"))
-	return dc.applyChanges(configDirectory)
+	return dc.applyChanges(ctx, configDirectory)
 }
 
 // DatasourceProvisioner is responsible for provisioning datasources based on
@@ -37,14 +38,14 @@ func newDatasourceProvisioner(log log.Logger) DatasourceProvisioner {
 	}
 }
 
-func (dc *DatasourceProvisioner) apply(cfg *configs) error {
+func (dc *DatasourceProvisioner) apply(ctx context.Context, cfg *configs) error {
-	if err := dc.deleteDatasources(cfg.DeleteDatasources); err != nil {
+	if err := dc.deleteDatasources(ctx, cfg.DeleteDatasources); err != nil {
 		return err
 	}
 
 	for _, ds := range cfg.Datasources {
 		cmd := &models.GetDataSourceQuery{OrgId: ds.OrgID, Name: ds.Name}
-		err := bus.Dispatch(cmd)
+		err := bus.DispatchCtx(ctx, cmd)
 		if err != nil && !errors.Is(err, models.ErrDataSourceNotFound) {
 			return err
 		}
@@ -52,13 +53,13 @@ func (dc *DatasourceProvisioner) apply(cfg *configs) error {
 		if errors.Is(err, models.ErrDataSourceNotFound) {
 			dc.log.Info("inserting datasource from configuration ", "name", ds.Name, "uid", ds.UID)
 			insertCmd := createInsertCommand(ds)
-			if err := bus.Dispatch(insertCmd); err != nil {
+			if err := bus.DispatchCtx(ctx, insertCmd); err != nil {
 				return err
 			}
 		} else {
 			dc.log.Debug("updating datasource from configuration", "name", ds.Name, "uid", ds.UID)
 			updateCmd := createUpdateCommand(ds, cmd.Result.Id)
-			if err := bus.Dispatch(updateCmd); err != nil {
+			if err := bus.DispatchCtx(ctx, updateCmd); err != nil {
 				return err
 			}
 		}
@@ -67,14 +68,14 @@ func (dc *DatasourceProvisioner) apply(cfg *configs) error {
 	return nil
 }
 
-func (dc *DatasourceProvisioner) applyChanges(configPath string) error {
+func (dc *DatasourceProvisioner) applyChanges(ctx context.Context, configPath string) error {
 	configs, err := dc.cfgProvider.readConfig(configPath)
 	if err != nil {
 		return err
 	}
 
 	for _, cfg := range configs {
-		if err := dc.apply(cfg); err != nil {
+		if err := dc.apply(ctx, cfg); err != nil {
 			return err
 		}
 	}
@@ -82,10 +83,10 @@ func (dc *DatasourceProvisioner) applyChanges(configPath string) error {
 	return nil
 }
 
-func (dc *DatasourceProvisioner) deleteDatasources(dsToDelete []*deleteDatasourceConfig) error {
+func (dc *DatasourceProvisioner) deleteDatasources(ctx context.Context, dsToDelete []*deleteDatasourceConfig) error {
 	for _, ds := range dsToDelete {
 		cmd := &models.DeleteDataSourceCommand{OrgID: ds.OrgID, Name: ds.Name}
-		if err := bus.Dispatch(cmd); err != nil {
+		if err := bus.DispatchCtx(ctx, cmd); err != nil {
 			return err
 		}
 

pkg/services/provisioning/provisioning.go

@@ -36,8 +36,8 @@ func ProvideService(cfg *setting.Cfg, sqlStore *sqlstore.SQLStore, pluginManager
 
 type ProvisioningService interface {
 	registry.BackgroundService
-	RunInitProvisioners() error
+	RunInitProvisioners(ctx context.Context) error
-	ProvisionDatasources() error
+	ProvisionDatasources(ctx context.Context) error
 	ProvisionPlugins() error
 	ProvisionNotifications() error
 	ProvisionDashboards(ctx context.Context) error
@@ -60,7 +60,7 @@ func NewProvisioningServiceImpl() *ProvisioningServiceImpl {
 func newProvisioningServiceImpl(
 	newDashboardProvisioner dashboards.DashboardProvisionerFactory,
 	provisionNotifiers func(string, encryption.Service) error,
-	provisionDatasources func(string) error,
+	provisionDatasources func(context.Context, string) error,
 	provisionPlugins func(string, plugifaces.Manager) error,
 ) *ProvisioningServiceImpl {
 	return &ProvisioningServiceImpl{
@@ -82,13 +82,13 @@ type ProvisioningServiceImpl struct {
 	newDashboardProvisioner dashboards.DashboardProvisionerFactory
 	dashboardProvisioner    dashboards.DashboardProvisioner
 	provisionNotifiers      func(string, encryption.Service) error
-	provisionDatasources    func(string) error
+	provisionDatasources    func(context.Context, string) error
 	provisionPlugins        func(string, plugifaces.Manager) error
 	mutex                   sync.Mutex
 }
 
-func (ps *ProvisioningServiceImpl) RunInitProvisioners() error {
+func (ps *ProvisioningServiceImpl) RunInitProvisioners(ctx context.Context) error {
-	err := ps.ProvisionDatasources()
+	err := ps.ProvisionDatasources(ctx)
 	if err != nil {
 		return err
 	}
@@ -135,9 +135,9 @@ func (ps *ProvisioningServiceImpl) Run(ctx context.Context) error {
 	}
 }
 
-func (ps *ProvisioningServiceImpl) ProvisionDatasources() error {
+func (ps *ProvisioningServiceImpl) ProvisionDatasources(ctx context.Context) error {
 	datasourcePath := filepath.Join(ps.Cfg.ProvisioningPath, "datasources")
-	err := ps.provisionDatasources(datasourcePath)
+	err := ps.provisionDatasources(ctx, datasourcePath)
 	return errutil.Wrap("Datasource provisioning error", err)
 }
 

pkg/services/provisioning/provisioning_mock.go

@@ -15,8 +15,8 @@ type Calls struct {
 
 type ProvisioningServiceMock struct {
 	Calls                                   *Calls
-	RunInitProvisionersFunc                 func() error
+	RunInitProvisionersFunc                 func(ctx context.Context) error
-	ProvisionDatasourcesFunc                func() error
+	ProvisionDatasourcesFunc                func(ctx context.Context) error
 	ProvisionPluginsFunc                    func() error
 	ProvisionNotificationsFunc              func() error
 	ProvisionDashboardsFunc                 func() error
@@ -31,18 +31,18 @@ func NewProvisioningServiceMock() *ProvisioningServiceMock {
 	}
 }
 
-func (mock *ProvisioningServiceMock) RunInitProvisioners() error {
+func (mock *ProvisioningServiceMock) RunInitProvisioners(ctx context.Context) error {
 	mock.Calls.RunInitProvisioners = append(mock.Calls.RunInitProvisioners, nil)
 	if mock.RunInitProvisionersFunc != nil {
-		return mock.RunInitProvisionersFunc()
+		return mock.RunInitProvisionersFunc(ctx)
 	}
 	return nil
 }
 
-func (mock *ProvisioningServiceMock) ProvisionDatasources() error {
+func (mock *ProvisioningServiceMock) ProvisionDatasources(ctx context.Context) error {
 	mock.Calls.ProvisionDatasources = append(mock.Calls.ProvisionDatasources, nil)
 	if mock.ProvisionDatasourcesFunc != nil {
-		return mock.ProvisionDatasourcesFunc()
+		return mock.ProvisionDatasourcesFunc(ctx)
 	}
 	return nil
 }

pkg/services/sqlstore/datasource.go

@@ -1,6 +1,7 @@
 package sqlstore
 
 import (
+	"context"
 	"fmt"
 	"strings"
 	"time"
@@ -15,24 +16,28 @@ import (
 
 // GetDataSource adds a datasource to the query model by querying by org_id as well as
 // either uid (preferred), id, or name and is added to the bus.
-func (ss *SQLStore) GetDataSource(query *models.GetDataSourceQuery) error {
+func (ss *SQLStore) GetDataSource(ctx context.Context, query *models.GetDataSourceQuery) error {
 	metrics.MDBDataSourceQueryByID.Inc()
-	if query.OrgId == 0 || (query.Id == 0 && len(query.Name) == 0 && len(query.Uid) == 0) {
-		return models.ErrDataSourceIdentifierNotSet
-	}
 
-	datasource := models.DataSource{Name: query.Name, OrgId: query.OrgId, Id: query.Id, Uid: query.Uid}
+	return ss.WithDbSession(ctx, func(sess *DBSession) error {
-	has, err := x.Get(&datasource)
+		if query.OrgId == 0 || (query.Id == 0 && len(query.Name) == 0 && len(query.Uid) == 0) {
+			return models.ErrDataSourceIdentifierNotSet
+		}
 
-	if err != nil {
+		datasource := &models.DataSource{Name: query.Name, OrgId: query.OrgId, Id: query.Id, Uid: query.Uid}
-		sqlog.Error("Failed getting data source", "err", err, "uid", query.Uid, "id", query.Id, "name", query.Name, "orgId", query.OrgId)
+		has, err := sess.Get(datasource)
-		return err
-	} else if !has {
-		return models.ErrDataSourceNotFound
-	}
 
-	query.Result = &datasource
+		if err != nil {
-	return nil
+			sqlog.Error("Failed getting data source", "err", err, "uid", query.Uid, "id", query.Id, "name", query.Name, "orgId", query.OrgId)
+			return err
+		} else if !has {
+			return models.ErrDataSourceNotFound
+		}
+
+		query.Result = datasource
+
+		return nil
+	})
 }
 
 func (ss *SQLStore) GetDataSources(query *models.GetDataSourcesQuery) error {
@@ -73,7 +78,7 @@ func (ss *SQLStore) GetDefaultDataSource(query *models.GetDefaultDataSourceQuery
 
 // DeleteDataSource removes a datasource by org_id as well as either uid (preferred), id, or name
 // and is added to the bus.
-func (ss *SQLStore) DeleteDataSource(cmd *models.DeleteDataSourceCommand) error {
+func (ss *SQLStore) DeleteDataSource(ctx context.Context, cmd *models.DeleteDataSourceCommand) error {
 	params := make([]interface{}, 0)
 
 	makeQuery := func(sql string, p ...interface{}) {
@@ -94,7 +99,7 @@ func (ss *SQLStore) DeleteDataSource(cmd *models.DeleteDataSourceCommand) error
 		return models.ErrDataSourceIdentifierNotSet
 	}
 
-	return inTransaction(func(sess *DBSession) error {
+	return ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
 		result, err := sess.Exec(params...)
 		cmd.DeletedDatasourcesCount, _ = result.RowsAffected()
 
@@ -110,8 +115,8 @@ func (ss *SQLStore) DeleteDataSource(cmd *models.DeleteDataSourceCommand) error
 	})
 }
 
-func (ss *SQLStore) AddDataSource(cmd *models.AddDataSourceCommand) error {
+func (ss *SQLStore) AddDataSource(ctx context.Context, cmd *models.AddDataSourceCommand) error {
-	return inTransaction(func(sess *DBSession) error {
+	return ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
 		existing := models.DataSource{OrgId: cmd.OrgId, Name: cmd.Name}
 		has, _ := sess.Get(&existing)
 
@@ -188,8 +193,8 @@ func updateIsDefaultFlag(ds *models.DataSource, sess *DBSession) error {
 	return nil
 }
 
-func (ss *SQLStore) UpdateDataSource(cmd *models.UpdateDataSourceCommand) error {
+func (ss *SQLStore) UpdateDataSource(ctx context.Context, cmd *models.UpdateDataSourceCommand) error {
-	return inTransaction(func(sess *DBSession) error {
+	return inTransactionCtx(ctx, func(sess *DBSession) error {
 		if cmd.JsonData == nil {
 			cmd.JsonData = simplejson.New()
 		}

pkg/services/sqlstore/datasource_test.go

@@ -4,6 +4,7 @@
 package sqlstore
 
 import (
+	"context"
 	"errors"
 	"strconv"
 	"testing"
@@ -35,7 +36,7 @@ func TestDataAccess(t *testing.T) {
 
 	initDatasource := func(sqlStore *SQLStore) *models.DataSource {
 		cmd := defaultAddDatasourceCommand
-		err := sqlStore.AddDataSource(&cmd)
+		err := sqlStore.AddDataSource(context.Background(), &cmd)
 		require.NoError(t, err)
 
 		query := models.GetDataSourcesQuery{OrgId: 10}
@@ -50,7 +51,7 @@ func TestDataAccess(t *testing.T) {
 		t.Run("Can add datasource", func(t *testing.T) {
 			sqlStore := InitTestDB(t)
 
-			err := sqlStore.AddDataSource(&models.AddDataSourceCommand{
+			err := sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 				OrgId:    10,
 				Name:     "laban",
 				Type:     models.DS_GRAPHITE,
@@ -85,9 +86,9 @@ func TestDataAccess(t *testing.T) {
 			cmd2 := defaultAddDatasourceCommand
 			cmd1.Uid = "test"
 			cmd2.Uid = "test"
-			err := sqlStore.AddDataSource(&cmd1)
+			err := sqlStore.AddDataSource(context.Background(), &cmd1)
 			require.NoError(t, err)
-			err = sqlStore.AddDataSource(&cmd2)
+			err = sqlStore.AddDataSource(context.Background(), &cmd2)
 			require.Error(t, err)
 			require.IsType(t, models.ErrDataSourceUidExists, err)
 		})
@@ -101,7 +102,7 @@ func TestDataAccess(t *testing.T) {
 				return nil
 			})
 
-			err := sqlStore.AddDataSource(&defaultAddDatasourceCommand)
+			err := sqlStore.AddDataSource(context.Background(), &defaultAddDatasourceCommand)
 			require.NoError(t, err)
 
 			require.Eventually(t, func() bool {
@@ -127,7 +128,7 @@ func TestDataAccess(t *testing.T) {
 			cmd := defaultUpdateDatasourceCommand
 			cmd.Id = ds.Id
 			cmd.Version = ds.Version
-			err := sqlStore.UpdateDataSource(&cmd)
+			err := sqlStore.UpdateDataSource(context.Background(), &cmd)
 			require.NoError(t, err)
 		})
 
@@ -138,11 +139,11 @@ func TestDataAccess(t *testing.T) {
 
 			cmd := defaultUpdateDatasourceCommand
 			cmd.Id = ds.Id
-			err := sqlStore.UpdateDataSource(&cmd)
+			err := sqlStore.UpdateDataSource(context.Background(), &cmd)
 			require.NoError(t, err)
 
 			query := models.GetDataSourceQuery{Id: ds.Id, OrgId: 10}
-			err = sqlStore.GetDataSource(&query)
+			err = sqlStore.GetDataSource(context.Background(), &query)
 			require.NoError(t, err)
 			require.Equal(t, ds.Uid, query.Result.Uid)
 		})
@@ -163,10 +164,10 @@ func TestDataAccess(t *testing.T) {
 			// Make a copy as UpdateDataSource modifies it
 			cmd2 := cmd
 
-			err := sqlStore.UpdateDataSource(&cmd)
+			err := sqlStore.UpdateDataSource(context.Background(), &cmd)
 			require.NoError(t, err)
 
-			err = sqlStore.UpdateDataSource(&cmd2)
+			err = sqlStore.UpdateDataSource(context.Background(), &cmd2)
 			require.Error(t, err)
 		})
 
@@ -183,7 +184,7 @@ func TestDataAccess(t *testing.T) {
 				Url:    "http://test",
 			}
 
-			err := sqlStore.UpdateDataSource(cmd)
+			err := sqlStore.UpdateDataSource(context.Background(), cmd)
 			require.NoError(t, err)
 		})
 
@@ -201,7 +202,7 @@ func TestDataAccess(t *testing.T) {
 				Version: 90000,
 			}
 
-			err := sqlStore.UpdateDataSource(cmd)
+			err := sqlStore.UpdateDataSource(context.Background(), cmd)
 			require.NoError(t, err)
 		})
 	})
@@ -211,7 +212,7 @@ func TestDataAccess(t *testing.T) {
 			sqlStore := InitTestDB(t)
 			ds := initDatasource(sqlStore)
 
-			err := sqlStore.DeleteDataSource(&models.DeleteDataSourceCommand{ID: ds.Id, OrgID: ds.OrgId})
+			err := sqlStore.DeleteDataSource(context.Background(), &models.DeleteDataSourceCommand{ID: ds.Id, OrgID: ds.OrgId})
 			require.NoError(t, err)
 
 			query := models.GetDataSourcesQuery{OrgId: 10}
@@ -225,7 +226,7 @@ func TestDataAccess(t *testing.T) {
 			sqlStore := InitTestDB(t)
 			ds := initDatasource(sqlStore)
 
-			err := sqlStore.DeleteDataSource(&models.DeleteDataSourceCommand{ID: ds.Id, OrgID: 123123})
+			err := sqlStore.DeleteDataSource(context.Background(), &models.DeleteDataSourceCommand{ID: ds.Id, OrgID: 123123})
 			require.NoError(t, err)
 			query := models.GetDataSourcesQuery{OrgId: 10}
 			err = sqlStore.GetDataSources(&query)
@@ -245,7 +246,7 @@ func TestDataAccess(t *testing.T) {
 			return nil
 		})
 
-		err := sqlStore.DeleteDataSource(&models.DeleteDataSourceCommand{ID: ds.Id, UID: "nisse-uid", Name: "nisse", OrgID: 123123})
+		err := sqlStore.DeleteDataSource(context.Background(), &models.DeleteDataSourceCommand{ID: ds.Id, UID: "nisse-uid", Name: "nisse", OrgID: 123123})
 		require.NoError(t, err)
 
 		require.Eventually(t, func() bool {
@@ -263,7 +264,7 @@ func TestDataAccess(t *testing.T) {
 		ds := initDatasource(sqlStore)
 		query := models.GetDataSourcesQuery{OrgId: 10}
 
-		err := sqlStore.DeleteDataSource(&models.DeleteDataSourceCommand{Name: ds.Name, OrgID: ds.OrgId})
+		err := sqlStore.DeleteDataSource(context.Background(), &models.DeleteDataSourceCommand{Name: ds.Name, OrgID: ds.OrgId})
 		require.NoError(t, err)
 
 		err = sqlStore.GetDataSources(&query)
@@ -277,7 +278,7 @@ func TestDataAccess(t *testing.T) {
 			sqlStore := InitTestDB(t)
 			datasourceLimit := 6
 			for i := 0; i < datasourceLimit+1; i++ {
-				err := sqlStore.AddDataSource(&models.AddDataSourceCommand{
+				err := sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 					OrgId:    10,
 					Name:     "laban" + strconv.Itoa(i),
 					Type:     models.DS_GRAPHITE,
@@ -300,7 +301,7 @@ func TestDataAccess(t *testing.T) {
 			sqlStore := InitTestDB(t)
 			numberOfDatasource := 5100
 			for i := 0; i < numberOfDatasource; i++ {
-				err := sqlStore.AddDataSource(&models.AddDataSourceCommand{
+				err := sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 					OrgId:    10,
 					Name:     "laban" + strconv.Itoa(i),
 					Type:     models.DS_GRAPHITE,
@@ -323,7 +324,7 @@ func TestDataAccess(t *testing.T) {
 			sqlStore := InitTestDB(t)
 			numberOfDatasource := 5100
 			for i := 0; i < numberOfDatasource; i++ {
-				err := sqlStore.AddDataSource(&models.AddDataSourceCommand{
+				err := sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 					OrgId:    10,
 					Name:     "laban" + strconv.Itoa(i),
 					Type:     models.DS_GRAPHITE,
@@ -347,7 +348,7 @@ func TestDataAccess(t *testing.T) {
 		t.Run("Only returns datasources of specified type", func(t *testing.T) {
 			sqlStore := InitTestDB(t)
 
-			err := sqlStore.AddDataSource(&models.AddDataSourceCommand{
+			err := sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 				OrgId:    10,
 				Name:     "Elasticsearch",
 				Type:     models.DS_ES,
@@ -358,7 +359,7 @@ func TestDataAccess(t *testing.T) {
 			})
 			require.NoError(t, err)
 
-			err = sqlStore.AddDataSource(&models.AddDataSourceCommand{
+			err = sqlStore.AddDataSource(context.Background(), &models.AddDataSourceCommand{
 				OrgId:    10,
 				Name:     "Graphite",
 				Type:     models.DS_GRAPHITE,
@@ -403,7 +404,7 @@ func TestGetDefaultDataSource(t *testing.T) {
 			Url:    "http://test",
 		}
 
-		err := sqlStore.AddDataSource(&cmd)
+		err := sqlStore.AddDataSource(context.Background(), &cmd)
 		require.NoError(t, err)
 
 		query := models.GetDefaultDataSourceQuery{OrgId: 10}
@@ -424,7 +425,7 @@ func TestGetDefaultDataSource(t *testing.T) {
 			IsDefault: true,
 		}
 
-		err := sqlStore.AddDataSource(&cmd)
+		err := sqlStore.AddDataSource(context.Background(), &cmd)
 		require.NoError(t, err)
 
 		query := models.GetDefaultDataSourceQuery{OrgId: 10}