SQLStore: Ensure that sessions are always closed (#55864)

* SQLStore: Ensure that sessions are always closed

Delete `NewSession()` in favour of `WithDbSession()`

* Add WithDbSessionForceNewSession to the interface

* Apply suggestions from code review
This commit is contained in:
Sofia Papagiannaki
2022-09-29 15:55:47 +03:00
committed by GitHub
parent c9e957a44e
commit 8b77ee2734
19 changed files with 289 additions and 210 deletions

View File

@@ -566,40 +566,54 @@ type testContext struct {
}
func (c *testContext) getAuthTokenByID(id int64) (*userAuthToken, error) {
sess := c.sqlstore.NewSession(context.Background())
var t userAuthToken
found, err := sess.ID(id).Get(&t)
if err != nil || !found {
return nil, err
}
var res *userAuthToken
err := c.sqlstore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
var t userAuthToken
found, err := sess.ID(id).Get(&t)
if err != nil || !found {
return err
}
return &t, nil
res = &t
return nil
})
return res, err
}
func (c *testContext) markAuthTokenAsSeen(id int64) (bool, error) {
sess := c.sqlstore.NewSession(context.Background())
res, err := sess.Exec("UPDATE user_auth_token SET auth_token_seen = ? WHERE id = ?", c.sqlstore.Dialect.BooleanStr(true), id)
if err != nil {
return false, err
}
hasRowsAffected := false
err := c.sqlstore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
res, err := sess.Exec("UPDATE user_auth_token SET auth_token_seen = ? WHERE id = ?", c.sqlstore.Dialect.BooleanStr(true), id)
if err != nil {
return err
}
rowsAffected, err := res.RowsAffected()
if err != nil {
return false, err
}
return rowsAffected == 1, nil
rowsAffected, err := res.RowsAffected()
if err != nil {
return err
}
hasRowsAffected = rowsAffected == 1
return nil
})
return hasRowsAffected, err
}
func (c *testContext) updateRotatedAt(id, rotatedAt int64) (bool, error) {
sess := c.sqlstore.NewSession(context.Background())
res, err := sess.Exec("UPDATE user_auth_token SET rotated_at = ? WHERE id = ?", rotatedAt, id)
if err != nil {
return false, err
}
hasRowsAffected := false
err := c.sqlstore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
res, err := sess.Exec("UPDATE user_auth_token SET rotated_at = ? WHERE id = ?", rotatedAt, id)
if err != nil {
return err
}
rowsAffected, err := res.RowsAffected()
if err != nil {
return false, err
}
return rowsAffected == 1, nil
rowsAffected, err := res.RowsAffected()
if err != nil {
return err
}
hasRowsAffected = rowsAffected == 1
return nil
})
return hasRowsAffected, err
}

View File

@@ -6,6 +6,7 @@ import (
"testing"
"time"
"github.com/grafana/grafana/pkg/services/sqlstore"
"github.com/stretchr/testify/require"
)
@@ -21,8 +22,12 @@ func TestUserAuthTokenCleanup(t *testing.T) {
insertToken := func(ctx *testContext, token string, prev string, createdAt, rotatedAt int64) {
ut := userAuthToken{AuthToken: token, PrevAuthToken: prev, CreatedAt: createdAt, RotatedAt: rotatedAt, UserAgent: "", ClientIp: ""}
_, err := ctx.sqlstore.NewSession(context.Background()).Insert(&ut)
require.Nil(t, err)
err := ctx.sqlstore.WithDbSession(context.Background(), func(sess *sqlstore.DBSession) error {
_, err := sess.Insert(&ut)
require.Nil(t, err)
return nil
})
require.NoError(t, err)
}
now := time.Date(2018, 12, 13, 13, 45, 0, 0, time.UTC)