add context to api crud calls (#40047)

Signed-off-by: bergquist <carl.bergquist@gmail.com>

pkg/api/apikey.go

@@ -11,10 +11,11 @@ import (
 	"github.com/grafana/grafana/pkg/models"
 )
 
+// GetAPIKeys returns a list of API keys
 func GetAPIKeys(c *models.ReqContext) response.Response {
 	query := models.GetApiKeysQuery{OrgId: c.OrgId, IncludeExpired: c.QueryBool("includeExpired")}
 
-	if err := bus.Dispatch(&query); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &query); err != nil {
 		return response.Error(500, "Failed to list api keys", err)
 	}
 
@@ -36,12 +37,13 @@ func GetAPIKeys(c *models.ReqContext) response.Response {
 	return response.JSON(200, result)
 }
 
+// DeleteAPIKey deletes an API key
 func DeleteAPIKey(c *models.ReqContext) response.Response {
 	id := c.ParamsInt64(":id")
 
 	cmd := &models.DeleteApiKeyCommand{Id: id, OrgId: c.OrgId}
 
-	err := bus.Dispatch(cmd)
+	err := bus.DispatchCtx(c.Req.Context(), cmd)
 	if err != nil {
 		var status int
 		if errors.Is(err, models.ErrApiKeyNotFound) {
@@ -55,6 +57,7 @@ func DeleteAPIKey(c *models.ReqContext) response.Response {
 	return response.Success("API key deleted")
 }
 
+// AddAPIKey adds an API key
 func (hs *HTTPServer) AddAPIKey(c *models.ReqContext, cmd models.AddApiKeyCommand) response.Response {
 	if !cmd.Role.IsValid() {
 		return response.Error(400, "Invalid role specified", nil)
@@ -77,7 +80,7 @@ func (hs *HTTPServer) AddAPIKey(c *models.ReqContext, cmd models.AddApiKeyComman
 
 	cmd.Key = newKeyInfo.HashedKey
 
-	if err := bus.Dispatch(&cmd); err != nil {
+	if err := bus.DispatchCtx(c.Req.Context(), &cmd); err != nil {
 		if errors.Is(err, models.ErrInvalidApiKeyExpiration) {
 			return response.Error(400, err.Error(), nil)
 		}

pkg/services/sqlstore/apikey.go

@@ -6,25 +6,36 @@ import (
 
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/models"
+	"xorm.io/xorm"
 )
 
 func init() {
-	bus.AddHandler("sql", GetApiKeys)
+	bus.AddHandlerCtx("sql", GetAPIKeys)
 	bus.AddHandler("sql", GetApiKeyById)
 	bus.AddHandler("sql", GetApiKeyByName)
 	bus.AddHandlerCtx("sql", DeleteApiKeyCtx)
-	bus.AddHandler("sql", AddApiKey)
+	bus.AddHandlerCtx("sql", AddAPIKey)
 }
 
-func GetApiKeys(query *models.GetApiKeysQuery) error {
-	sess := x.Limit(100, 0).Where("org_id=? and ( expires IS NULL or expires >= ?)",
-		query.OrgId, timeNow().Unix()).Asc("name")
-	if query.IncludeExpired {
-		sess = x.Limit(100, 0).Where("org_id=?", query.OrgId).Asc("name")
-	}
+// GetAPIKeys queries the database based
+// on input on GetApiKeysQuery
+func GetAPIKeys(ctx context.Context, query *models.GetApiKeysQuery) error {
+	return withDbSession(ctx, x, func(dbSession *DBSession) error {
+		var sess *xorm.Session
+
+		if query.IncludeExpired {
+			sess = dbSession.Limit(100, 0).
+				Where("org_id=?", query.OrgId).
+				Asc("name")
+		} else {
+			sess = dbSession.Limit(100, 0).
+				Where("org_id=? and ( expires IS NULL or expires >= ?)", query.OrgId, timeNow().Unix()).
+				Asc("name")
+		}
 
-	query.Result = make([]*models.ApiKey, 0)
-	return sess.Find(&query.Result)
+		query.Result = make([]*models.ApiKey, 0)
+		return sess.Find(&query.Result)
+	})
 }
 
 func DeleteApiKeyCtx(ctx context.Context, cmd *models.DeleteApiKeyCommand) error {
@@ -48,8 +59,9 @@ func deleteAPIKey(sess *DBSession, id, orgID int64) error {
 	return nil
 }
 
-func AddApiKey(cmd *models.AddApiKeyCommand) error {
-	return inTransaction(func(sess *DBSession) error {
+// AddAPIKey adds the API key to the database.
+func AddAPIKey(ctx context.Context, cmd *models.AddApiKeyCommand) error {
+	return inTransactionCtx(ctx, func(sess *DBSession) error {
 		key := models.ApiKey{OrgId: cmd.OrgId, Name: cmd.Name}
 		exists, _ := sess.Get(&key)
 		if exists {

pkg/services/sqlstore/apikey_test.go

@@ -21,7 +21,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 
 		t.Run("Given saved api key", func(t *testing.T) {
 			cmd := models.AddApiKeyCommand{OrgId: 1, Name: "hello", Key: "asd"}
-			err := AddApiKey(&cmd)
+			err := AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			t.Run("Should be able to get key by name", func(t *testing.T) {
@@ -35,7 +35,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 
 		t.Run("Add non expiring key", func(t *testing.T) {
 			cmd := models.AddApiKeyCommand{OrgId: 1, Name: "non-expiring", Key: "asd1", SecondsToLive: 0}
-			err := AddApiKey(&cmd)
+			err := AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			query := models.GetApiKeyByNameQuery{KeyName: "non-expiring", OrgId: 1}
@@ -48,7 +48,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 		t.Run("Add an expiring key", func(t *testing.T) {
 			// expires in one hour
 			cmd := models.AddApiKeyCommand{OrgId: 1, Name: "expiring-in-an-hour", Key: "asd2", SecondsToLive: 3600}
-			err := AddApiKey(&cmd)
+			err := AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			query := models.GetApiKeyByNameQuery{KeyName: "expiring-in-an-hour", OrgId: 1}
@@ -57,7 +57,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 
 			assert.True(t, *query.Result.Expires >= timeNow().Unix())
 
-			// timeNow() has been called twice since creation; once by AddApiKey and once by GetApiKeyByName
+			// timeNow() has been called twice since creation; once by AddAPIKey and once by GetApiKeyByName
 			// therefore two seconds should be subtracted by next value returned by timeNow()
 			// that equals the number by which timeSeed has been advanced
 			then := timeNow().Add(-2 * time.Second)
@@ -68,7 +68,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 		t.Run("Add a key with negative lifespan", func(t *testing.T) {
 			// expires in one day
 			cmd := models.AddApiKeyCommand{OrgId: 1, Name: "key-with-negative-lifespan", Key: "asd3", SecondsToLive: -3600}
-			err := AddApiKey(&cmd)
+			err := AddAPIKey(context.Background(), &cmd)
 			assert.EqualError(t, err, models.ErrInvalidApiKeyExpiration.Error())
 
 			query := models.GetApiKeyByNameQuery{KeyName: "key-with-negative-lifespan", OrgId: 1}
@@ -79,24 +79,24 @@ func TestApiKeyDataAccess(t *testing.T) {
 		t.Run("Add keys", func(t *testing.T) {
 			// never expires
 			cmd := models.AddApiKeyCommand{OrgId: 1, Name: "key1", Key: "key1", SecondsToLive: 0}
-			err := AddApiKey(&cmd)
+			err := AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			// expires in 1s
 			cmd = models.AddApiKeyCommand{OrgId: 1, Name: "key2", Key: "key2", SecondsToLive: 1}
-			err = AddApiKey(&cmd)
+			err = AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			// expires in one hour
 			cmd = models.AddApiKeyCommand{OrgId: 1, Name: "key3", Key: "key3", SecondsToLive: 3600}
-			err = AddApiKey(&cmd)
+			err = AddAPIKey(context.Background(), &cmd)
 			assert.Nil(t, err)
 
 			// advance mocked getTime by 1s
 			timeNow()
 
 			query := models.GetApiKeysQuery{OrgId: 1, IncludeExpired: false}
-			err = GetApiKeys(&query)
+			err = GetAPIKeys(context.Background(), &query)
 			assert.Nil(t, err)
 
 			for _, k := range query.Result {
@@ -106,7 +106,7 @@ func TestApiKeyDataAccess(t *testing.T) {
 			}
 
 			query = models.GetApiKeysQuery{OrgId: 1, IncludeExpired: true}
-			err = GetApiKeys(&query)
+			err = GetAPIKeys(context.Background(), &query)
 			assert.Nil(t, err)
 
 			found := false
@@ -137,12 +137,12 @@ func TestApiKeyErrors(t *testing.T) {
 		t.Run("Testing API Duplicate Key Errors", func(t *testing.T) {
 			t.Run("Given saved api key", func(t *testing.T) {
 				cmd := models.AddApiKeyCommand{OrgId: 0, Name: "duplicate", Key: "asd"}
-				err := AddApiKey(&cmd)
+				err := AddAPIKey(context.Background(), &cmd)
 				assert.Nil(t, err)
 
 				t.Run("Add API Key with existing Org ID and Name", func(t *testing.T) {
 					cmd := models.AddApiKeyCommand{OrgId: 0, Name: "duplicate", Key: "asd"}
-					err = AddApiKey(&cmd)
+					err = AddAPIKey(context.Background(), &cmd)
 					assert.EqualError(t, err, models.ErrDuplicateApiKey.Error())
 				})
 			})

pkg/services/sqlstore/transactions_test.go

@@ -21,7 +21,7 @@ func TestTransaction(t *testing.T) {
 	Convey("InTransaction", t, func() {
 		cmd := &models.AddApiKeyCommand{Key: "secret-key", Name: "key", OrgId: 1}
 
-		err := AddApiKey(cmd)
+		err := AddAPIKey(context.Background(), cmd)
 		So(err, ShouldBeNil)
 
 		Convey("can update key", func() {