SQLStore: Close session in withDbSession (#31775)

* SQLStore: Close session in withDbSession

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

* SQLStore.WithDbSession: Never use session from context

Signed-off-by: Arve Knudsen <arve.knudsen@gmail.com>

pkg/services/sqlstore/alert_notification.go

@@ -510,7 +510,7 @@ func SetAlertNotificationStateToCompleteCommand(ctx context.Context, cmd *models
 }
 
 func SetAlertNotificationStateToPendingCommand(ctx context.Context, cmd *models.SetAlertNotificationStateToPendingCommand) error {
-	return withDbSession(ctx, func(sess *DBSession) error {
+	return withDbSession(ctx, x, func(sess *DBSession) error {
 		newVersion := cmd.Version + 1
 		sql := `UPDATE alert_notification_state SET
 			state = ?,

pkg/services/sqlstore/annotation_cleanup.go

@@ -91,7 +91,7 @@ func (acs *AnnotationCleanupService) executeUntilDoneOrCancelled(ctx context.Con
 			return totalAffected, ctx.Err()
 		default:
 			var affected int64
-			err := withDbSession(ctx, func(session *DBSession) error {
+			err := withDbSession(ctx, x, func(session *DBSession) error {
 				res, err := session.Exec(sql)
 				if err != nil {
 					return err

pkg/services/sqlstore/apikey.go

@@ -28,13 +28,17 @@ func GetApiKeys(query *models.GetApiKeysQuery) error {
 }
 
 func DeleteApiKeyCtx(ctx context.Context, cmd *models.DeleteApiKeyCommand) error {
-	return withDbSession(ctx, func(sess *DBSession) error {
-		var rawSQL = "DELETE FROM api_key WHERE id=? and org_id=?"
-		_, err := sess.Exec(rawSQL, cmd.Id, cmd.OrgId)
-		return err
+	return withDbSession(ctx, x, func(sess *DBSession) error {
+		return deleteAPIKey(sess, cmd.Id, cmd.OrgId)
 	})
 }
 
+func deleteAPIKey(sess *DBSession, id, orgID int64) error {
+	rawSQL := "DELETE FROM api_key WHERE id=? and org_id=?"
+	_, err := sess.Exec(rawSQL, id, orgID)
+	return err
+}
+
 func AddApiKey(cmd *models.AddApiKeyCommand) error {
 	return inTransaction(func(sess *DBSession) error {
 		key := models.ApiKey{OrgId: cmd.OrgId, Name: cmd.Name}

pkg/services/sqlstore/session.go

@@ -46,22 +46,14 @@ func startSession(ctx context.Context, engine *xorm.Engine, beginTran bool) (*DB
 	return newSess, nil
 }
 
-// WithDbSession calls the callback with an session attached to the context.
+// WithDbSession calls the callback with a session.
 func (ss *SQLStore) WithDbSession(ctx context.Context, callback dbTransactionFunc) error {
-	sess, err := startSession(ctx, ss.engine, false)
-	if err != nil {
-		return err
-	}
-	defer sess.Close()
-
-	return callback(sess)
+	return withDbSession(ctx, ss.engine, callback)
 }
 
-func withDbSession(ctx context.Context, callback dbTransactionFunc) error {
-	sess, err := startSession(ctx, x, false)
-	if err != nil {
-		return err
-	}
+func withDbSession(ctx context.Context, engine *xorm.Engine, callback dbTransactionFunc) error {
+	sess := &DBSession{Session: engine.NewSession()}
+	defer sess.Close()
 
 	return callback(sess)
 }

pkg/services/sqlstore/sqlstore.go

@@ -141,21 +141,15 @@ func (ss *SQLStore) Reset() error {
 }
 
 func (ss *SQLStore) ensureMainOrgAndAdminUser() error {
-	err := ss.InTransaction(context.Background(), func(ctx context.Context) error {
+	ctx := context.Background()
+	err := ss.WithTransactionalDbSession(ctx, func(sess *DBSession) error {
 		ss.log.Debug("Ensuring main org and admin user exist")
 		var stats models.SystemUserCountStats
-		err := ss.WithDbSession(ctx, func(sess *DBSession) error {
-			// TODO: Should be able to rename "Count" to "count", for more standard SQL style
-			// Just have to make sure it gets deserialized properly into models.SystemUserCountStats
-			rawSQL := `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
-			if _, err := sess.SQL(rawSQL).Get(&stats); err != nil {
-				return fmt.Errorf("could not determine if admin user exists: %w", err)
-			}
-
-			return nil
-		})
-		if err != nil {
-			return err
+		// TODO: Should be able to rename "Count" to "count", for more standard SQL style
+		// Just have to make sure it gets deserialized properly into models.SystemUserCountStats
+		rawSQL := `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
+		if _, err := sess.SQL(rawSQL).Get(&stats); err != nil {
+			return fmt.Errorf("could not determine if admin user exists: %w", err)
 		}
 
 		if stats.Count > 0 {
@@ -166,7 +160,7 @@ func (ss *SQLStore) ensureMainOrgAndAdminUser() error {
 		if !ss.Cfg.DisableInitAdminCreation {
 			ss.log.Debug("Creating default admin user")
 			ss.log.Debug("Creating default admin user")
-			if _, err := ss.createUser(ctx, userCreationArgs{
+			if _, err := ss.createUser(ctx, sess, userCreationArgs{
 				Login:    ss.Cfg.AdminUser,
 				Email:    ss.Cfg.AdminUser + "@localhost",
 				Password: ss.Cfg.AdminPassword,
@@ -181,11 +175,8 @@ func (ss *SQLStore) ensureMainOrgAndAdminUser() error {
 			// return nil
 		}
 
-		if err := inTransactionWithRetryCtx(ctx, ss.engine, func(sess *DBSession) error {
-			ss.log.Debug("Creating default org", "name", MainOrgName)
-			_, err := ss.getOrCreateOrg(sess, MainOrgName)
-			return err
-		}, 0); err != nil {
+		ss.log.Debug("Creating default org", "name", MainOrgName)
+		if _, err := ss.getOrCreateOrg(sess, MainOrgName); err != nil {
 			return fmt.Errorf("failed to create default organization: %w", err)
 		}
 

pkg/services/sqlstore/stats.go

@@ -170,7 +170,7 @@ func GetAdminStats(query *models.GetAdminStatsQuery) error {
 }
 
 func GetSystemUserCountStats(ctx context.Context, query *models.GetSystemUserCountStatsQuery) error {
-	return withDbSession(ctx, func(sess *DBSession) error {
+	return withDbSession(ctx, x, func(sess *DBSession) error {
 		var rawSQL = `SELECT COUNT(id) AS Count FROM ` + dialect.Quote("user")
 		var stats models.SystemUserCountStats
 		_, err := sess.SQL(rawSQL).Get(&stats)

pkg/services/sqlstore/transactions_test.go

@@ -23,11 +23,9 @@ func TestTransaction(t *testing.T) {
 		err := AddApiKey(cmd)
 		So(err, ShouldBeNil)
 
-		deleteApiKeyCmd := &models.DeleteApiKeyCommand{Id: cmd.Result.Id, OrgId: 1}
-
 		Convey("can update key", func() {
-			err := ss.InTransaction(context.Background(), func(ctx context.Context) error {
-				return DeleteApiKeyCtx(ctx, deleteApiKeyCmd)
+			err := ss.WithTransactionalDbSession(context.Background(), func(sess *DBSession) error {
+				return deleteAPIKey(sess, cmd.Result.Id, 1)
 			})
 
 			So(err, ShouldBeNil)
@@ -38,8 +36,8 @@ func TestTransaction(t *testing.T) {
 		})
 
 		Convey("won't update if one handler fails", func() {
-			err := ss.InTransaction(context.Background(), func(ctx context.Context) error {
-				err := DeleteApiKeyCtx(ctx, deleteApiKeyCmd)
+			err := ss.WithTransactionalDbSession(context.Background(), func(sess *DBSession) error {
+				err := deleteAPIKey(sess, cmd.Result.Id, 1)
 				if err != nil {
 					return err
 				}

pkg/services/sqlstore/user.go

@@ -86,104 +86,98 @@ func (ss *SQLStore) getOrgIDForNewUser(sess *DBSession, args userCreationArgs) (
 }
 
 // createUser creates a user in the database.
-func (ss *SQLStore) createUser(ctx context.Context, args userCreationArgs, skipOrgSetup bool) (models.User, error) {
+func (ss *SQLStore) createUser(ctx context.Context, sess *DBSession, args userCreationArgs, skipOrgSetup bool) (models.User, error) {
 	var user models.User
-	if err := inTransactionWithRetryCtx(ctx, ss.engine, func(sess *DBSession) error {
-		var orgID int64 = -1
-		if !skipOrgSetup {
-			var err error
-			orgID, err = ss.getOrgIDForNewUser(sess, args)
-			if err != nil {
-				return err
-			}
-		}
-
-		if args.Email == "" {
-			args.Email = args.Login
-		}
-
-		exists, err := sess.Where("email=? OR login=?", args.Email, args.Login).Get(&models.User{})
+	var orgID int64 = -1
+	if !skipOrgSetup {
+		var err error
+		orgID, err = ss.getOrgIDForNewUser(sess, args)
 		if err != nil {
-			return err
-		}
-		if exists {
-			return models.ErrUserAlreadyExists
+			return user, err
 		}
+	}
 
-		// create user
-		user = models.User{
-			Email:         args.Email,
-			Name:          args.Name,
-			Login:         args.Login,
-			Company:       args.Company,
-			IsAdmin:       args.IsAdmin,
-			IsDisabled:    args.IsDisabled,
-			OrgId:         orgID,
-			EmailVerified: args.EmailVerified,
-			Created:       time.Now(),
-			Updated:       time.Now(),
-			LastSeenAt:    time.Now().AddDate(-10, 0, 0),
-		}
+	if args.Email == "" {
+		args.Email = args.Login
+	}
 
-		salt, err := util.GetRandomString(10)
-		if err != nil {
-			return err
-		}
-		user.Salt = salt
-		rands, err := util.GetRandomString(10)
-		if err != nil {
-			return err
-		}
-		user.Rands = rands
+	exists, err := sess.Where("email=? OR login=?", args.Email, args.Login).Get(&models.User{})
+	if err != nil {
+		return user, err
+	}
+	if exists {
+		return user, models.ErrUserAlreadyExists
+	}
 
-		if len(args.Password) > 0 {
-			encodedPassword, err := util.EncodePassword(args.Password, user.Salt)
-			if err != nil {
-				return err
-			}
-			user.Password = encodedPassword
-		}
+	// create user
+	user = models.User{
+		Email:         args.Email,
+		Name:          args.Name,
+		Login:         args.Login,
+		Company:       args.Company,
+		IsAdmin:       args.IsAdmin,
+		IsDisabled:    args.IsDisabled,
+		OrgId:         orgID,
+		EmailVerified: args.EmailVerified,
+		Created:       time.Now(),
+		Updated:       time.Now(),
+		LastSeenAt:    time.Now().AddDate(-10, 0, 0),
+	}
 
-		sess.UseBool("is_admin")
+	salt, err := util.GetRandomString(10)
+	if err != nil {
+		return user, err
+	}
+	user.Salt = salt
+	rands, err := util.GetRandomString(10)
+	if err != nil {
+		return user, err
+	}
+	user.Rands = rands
 
-		if _, err := sess.Insert(&user); err != nil {
-			return err
+	if len(args.Password) > 0 {
+		encodedPassword, err := util.EncodePassword(args.Password, user.Salt)
+		if err != nil {
+			return user, err
 		}
+		user.Password = encodedPassword
+	}
 
-		sess.publishAfterCommit(&events.UserCreated{
-			Timestamp: user.Created,
-			Id:        user.Id,
-			Name:      user.Name,
-			Login:     user.Login,
-			Email:     user.Email,
-		})
+	sess.UseBool("is_admin")
 
-		// create org user link
-		if !skipOrgSetup {
-			orgUser := models.OrgUser{
-				OrgId:   orgID,
-				UserId:  user.Id,
-				Role:    models.ROLE_ADMIN,
-				Created: time.Now(),
-				Updated: time.Now(),
-			}
+	if _, err := sess.Insert(&user); err != nil {
+		return user, err
+	}
 
-			if ss.Cfg.AutoAssignOrg && !user.IsAdmin {
-				if len(args.DefaultOrgRole) > 0 {
-					orgUser.Role = models.RoleType(args.DefaultOrgRole)
-				} else {
-					orgUser.Role = models.RoleType(ss.Cfg.AutoAssignOrgRole)
-				}
-			}
+	sess.publishAfterCommit(&events.UserCreated{
+		Timestamp: user.Created,
+		Id:        user.Id,
+		Name:      user.Name,
+		Login:     user.Login,
+		Email:     user.Email,
+	})
 
-			if _, err = sess.Insert(&orgUser); err != nil {
-				return err
+	// create org user link
+	if !skipOrgSetup {
+		orgUser := models.OrgUser{
+			OrgId:   orgID,
+			UserId:  user.Id,
+			Role:    models.ROLE_ADMIN,
+			Created: time.Now(),
+			Updated: time.Now(),
+		}
+
+		if ss.Cfg.AutoAssignOrg && !user.IsAdmin {
+			if len(args.DefaultOrgRole) > 0 {
+				orgUser.Role = models.RoleType(args.DefaultOrgRole)
+			} else {
+				orgUser.Role = models.RoleType(ss.Cfg.AutoAssignOrgRole)
 			}
 		}
 
-		return nil
-	}, 0); err != nil {
-		return user, err
+		if _, err = sess.Insert(&orgUser); err != nil {
+			return user, err
+		}
 	}
 
 	return user, nil