Use TransactionWriter in SQLite keyserver (#1239)

* Use TransactionWriter in SQLite keyserver

* Fix keyserver storage tests
This commit is contained in:
Neil Alexander 2020-08-05 11:01:37 +01:00 committed by GitHub
parent 22f028e141
commit 15dc1f4d03
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 65 additions and 38 deletions

View file

@ -20,6 +20,7 @@ import (
"time"
"github.com/matrix-org/dendrite/internal"
"github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/keyserver/api"
"github.com/matrix-org/dendrite/keyserver/storage/tables"
)
@ -54,6 +55,7 @@ const selectMaxStreamForUserSQL = "" +
type deviceKeysStatements struct {
db *sql.DB
writer *sqlutil.TransactionWriter
upsertDeviceKeysStmt *sql.Stmt
selectDeviceKeysStmt *sql.Stmt
selectBatchDeviceKeysStmt *sql.Stmt
@ -62,7 +64,8 @@ type deviceKeysStatements struct {
func NewSqliteDeviceKeysTable(db *sql.DB) (tables.DeviceKeys, error) {
s := &deviceKeysStatements{
db: db,
db: db,
writer: sqlutil.NewTransactionWriter(),
}
_, err := db.Exec(deviceKeysSchema)
if err != nil {
@ -141,14 +144,16 @@ func (s *deviceKeysStatements) SelectMaxStreamIDForUser(ctx context.Context, txn
}
func (s *deviceKeysStatements) InsertDeviceKeys(ctx context.Context, txn *sql.Tx, keys []api.DeviceMessage) error {
for _, key := range keys {
now := time.Now().Unix()
_, err := txn.Stmt(s.upsertDeviceKeysStmt).ExecContext(
ctx, key.UserID, key.DeviceID, now, string(key.KeyJSON), key.StreamID,
)
if err != nil {
return err
return s.writer.Do(s.db, txn, func(txn *sql.Tx) error {
for _, key := range keys {
now := time.Now().Unix()
_, err := txn.Stmt(s.upsertDeviceKeysStmt).ExecContext(
ctx, key.UserID, key.DeviceID, now, string(key.KeyJSON), key.StreamID,
)
if err != nil {
return err
}
}
}
return nil
return nil
})
}