1 files changed, 16 insertions, 14 deletions

diff --git a/models/ssh_key.go b/models/ssh_key.go
index baa0d2a54f..21d5b7d496 100644
--- a/models/ssh_key.go
+++ b/models/ssh_key.go
@@ -96,11 +96,12 @@ func AddPublicKey(ownerID int64, name, content string, loginSourceID int64) (*Pu
 		return nil, err
 	}
 
-	sess := db.NewSession(db.DefaultContext)
-	defer sess.Close()
-	if err = sess.Begin(); err != nil {
+	ctx, committer, err := db.TxContext()
+	if err != nil {
 		return nil, err
 	}
+	defer committer.Close()
+	sess := db.GetEngine(ctx)
 
 	if err := checkKeyFingerprint(sess, fingerprint); err != nil {
 		return nil, err
@@ -129,7 +130,7 @@ func AddPublicKey(ownerID int64, name, content string, loginSourceID int64) (*Pu
 		return nil, fmt.Errorf("addKey: %v", err)
 	}
 
-	return key, sess.Commit()
+	return key, committer.Commit()
 }
 
 // GetPublicKeyByID returns public key by given ID.
@@ -334,20 +335,20 @@ func DeletePublicKey(doer *User, id int64) (err error) {
 		return ErrKeyAccessDenied{doer.ID, key.ID, "public"}
 	}
 
-	sess := db.NewSession(db.DefaultContext)
-	defer sess.Close()
-	if err = sess.Begin(); err != nil {
+	ctx, committer, err := db.TxContext()
+	if err != nil {
 		return err
 	}
+	defer committer.Close()
 
-	if err = deletePublicKeys(sess, id); err != nil {
+	if err = deletePublicKeys(db.GetEngine(ctx), id); err != nil {
 		return err
 	}
 
-	if err = sess.Commit(); err != nil {
+	if err = committer.Commit(); err != nil {
 		return err
 	}
-	sess.Close()
+	committer.Close()
 
 	if key.Type == KeyTypePrincipal {
 		return RewriteAllPrincipalKeys()
@@ -359,11 +360,12 @@ func DeletePublicKey(doer *User, id int64) (err error) {
 // deleteKeysMarkedForDeletion returns true if ssh keys needs update
 func deleteKeysMarkedForDeletion(keys []string) (bool, error) {
 	// Start session
-	sess := db.NewSession(db.DefaultContext)
-	defer sess.Close()
-	if err := sess.Begin(); err != nil {
+	ctx, committer, err := db.TxContext()
+	if err != nil {
 		return false, err
 	}
+	defer committer.Close()
+	sess := db.GetEngine(ctx)
 
 	// Delete keys marked for deletion
 	var sshKeysNeedUpdate bool
@@ -380,7 +382,7 @@ func deleteKeysMarkedForDeletion(keys []string) (bool, error) {
 		sshKeysNeedUpdate = true
 	}
 
-	if err := sess.Commit(); err != nil {
+	if err := committer.Commit(); err != nil {
 		return false, err
 	}