3 files changed, 31 insertions, 55 deletions

diff --git a/services/asymkey/deploy_key.go b/services/asymkey/deploy_key.go
index 9e5a6d6292..04023f9ffb 100644
--- a/services/asymkey/deploy_key.go
+++ b/services/asymkey/deploy_key.go
@@ -49,28 +49,21 @@ func deleteDeployKeyFromDB(ctx context.Context, key *asymkey_model.DeployKey) er
 // DeleteDeployKey deletes deploy key from its repository authorized_keys file if needed.
 // Permissions check should be done outside.
 func DeleteDeployKey(ctx context.Context, repo *repo_model.Repository, id int64) error {
-	dbCtx, committer, err := db.TxContext(ctx)
-	if err != nil {
-		return err
-	}
-	defer committer.Close()
-
-	key, err := asymkey_model.GetDeployKeyByID(ctx, id)
-	if err != nil {
-		if asymkey_model.IsErrDeployKeyNotExist(err) {
-			return nil
+	if err := db.WithTx(ctx, func(ctx context.Context) error {
+		key, err := asymkey_model.GetDeployKeyByID(ctx, id)
+		if err != nil {
+			if asymkey_model.IsErrDeployKeyNotExist(err) {
+				return nil
+			}
+			return fmt.Errorf("GetDeployKeyByID: %w", err)
 		}
-		return fmt.Errorf("GetDeployKeyByID: %w", err)
-	}
 
-	if key.RepoID != repo.ID {
-		return fmt.Errorf("deploy key %d does not belong to repository %d", id, repo.ID)
-	}
+		if key.RepoID != repo.ID {
+			return fmt.Errorf("deploy key %d does not belong to repository %d", id, repo.ID)
+		}
 
-	if err := deleteDeployKeyFromDB(dbCtx, key); err != nil {
-		return err
-	}
-	if err := committer.Commit(); err != nil {
+		return deleteDeployKeyFromDB(ctx, key)
+	}); err != nil {
 		return err
 	}
 
diff --git a/services/asymkey/ssh_key.go b/services/asymkey/ssh_key.go
index da57059d4b..01fa7ff15f 100644
--- a/services/asymkey/ssh_key.go
+++ b/services/asymkey/ssh_key.go
@@ -27,20 +27,9 @@ func DeletePublicKey(ctx context.Context, doer *user_model.User, id int64) (err
 		}
 	}
 
-	dbCtx, committer, err := db.TxContext(ctx)
-	if err != nil {
-		return err
-	}
-	defer committer.Close()
-
-	if _, err = db.DeleteByID[asymkey_model.PublicKey](dbCtx, id); err != nil {
-		return err
-	}
-
-	if err = committer.Commit(); err != nil {
+	if _, err = db.DeleteByID[asymkey_model.PublicKey](ctx, id); err != nil {
 		return err
 	}
-	committer.Close()
 
 	if key.Type == asymkey_model.KeyTypePrincipal {
 		return RewriteAllPrincipalKeys(ctx)
diff --git a/services/asymkey/ssh_key_principals.go b/services/asymkey/ssh_key_principals.go
index 5ed5cfa782..6493c1cc51 100644
--- a/services/asymkey/ssh_key_principals.go
+++ b/services/asymkey/ssh_key_principals.go
@@ -14,24 +14,6 @@ import (
 
 // AddPrincipalKey adds new principal to database and authorized_principals file.
 func AddPrincipalKey(ctx context.Context, ownerID int64, content string, authSourceID int64) (*asymkey_model.PublicKey, error) {
-	dbCtx, committer, err := db.TxContext(ctx)
-	if err != nil {
-		return nil, err
-	}
-	defer committer.Close()
-
-	// Principals cannot be duplicated.
-	has, err := db.GetEngine(dbCtx).
-		Where("content = ? AND type = ?", content, asymkey_model.KeyTypePrincipal).
-		Get(new(asymkey_model.PublicKey))
-	if err != nil {
-		return nil, err
-	} else if has {
-		return nil, asymkey_model.ErrKeyAlreadyExist{
-			Content: content,
-		}
-	}
-
 	key := &asymkey_model.PublicKey{
 		OwnerID:       ownerID,
 		Name:          content,
@@ -40,15 +22,27 @@ func AddPrincipalKey(ctx context.Context, ownerID int64, content string, authSou
 		Type:          asymkey_model.KeyTypePrincipal,
 		LoginSourceID: authSourceID,
 	}
-	if err = db.Insert(dbCtx, key); err != nil {
-		return nil, fmt.Errorf("addKey: %w", err)
-	}
 
-	if err = committer.Commit(); err != nil {
+	if err := db.WithTx(ctx, func(ctx context.Context) error {
+		// Principals cannot be duplicated.
+		has, err := db.GetEngine(ctx).
+			Where("content = ? AND type = ?", content, asymkey_model.KeyTypePrincipal).
+			Get(new(asymkey_model.PublicKey))
+		if err != nil {
+			return err
+		} else if has {
+			return asymkey_model.ErrKeyAlreadyExist{
+				Content: content,
+			}
+		}
+
+		if err = db.Insert(ctx, key); err != nil {
+			return fmt.Errorf("addKey: %w", err)
+		}
+		return nil
+	}); err != nil {
 		return nil, err
 	}
 
-	committer.Close()
-
 	return key, RewriteAllPrincipalKeys(ctx)
 }