Use Writer in shared package (#1296)

This commit is contained in:
Neil Alexander 2020-08-25 10:29:45 +01:00 committed by GitHub
parent 3b14119aff
commit 720ddce0a8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 98 additions and 109 deletions

View file

@ -63,7 +63,6 @@ const deleteAllDeviceKeysSQL = "" +
type deviceKeysStatements struct {
db *sql.DB
writer sqlutil.Writer
upsertDeviceKeysStmt *sql.Stmt
selectDeviceKeysStmt *sql.Stmt
selectBatchDeviceKeysStmt *sql.Stmt
@ -71,10 +70,9 @@ type deviceKeysStatements struct {
deleteAllDeviceKeysStmt *sql.Stmt
}
func NewSqliteDeviceKeysTable(db *sql.DB, writer sqlutil.Writer) (tables.DeviceKeys, error) {
func NewSqliteDeviceKeysTable(db *sql.DB) (tables.DeviceKeys, error) {
s := &deviceKeysStatements{
db: db,
writer: writer,
db: db,
}
_, err := db.Exec(deviceKeysSchema)
if err != nil {
@ -188,16 +186,14 @@ func (s *deviceKeysStatements) CountStreamIDsForUser(ctx context.Context, userID
}
func (s *deviceKeysStatements) InsertDeviceKeys(ctx context.Context, txn *sql.Tx, keys []api.DeviceMessage) error {
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, key.DisplayName,
)
if err != nil {
return err
}
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, key.DisplayName,
)
if err != nil {
return err
}
return nil
})
}
return nil
}

View file

@ -21,7 +21,6 @@ import (
"github.com/Shopify/sarama"
"github.com/matrix-org/dendrite/internal"
"github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/keyserver/storage/tables"
)
@ -52,15 +51,13 @@ const selectKeyChangesSQL = "" +
type keyChangesStatements struct {
db *sql.DB
writer sqlutil.Writer
upsertKeyChangeStmt *sql.Stmt
selectKeyChangesStmt *sql.Stmt
}
func NewSqliteKeyChangesTable(db *sql.DB, writer sqlutil.Writer) (tables.KeyChanges, error) {
func NewSqliteKeyChangesTable(db *sql.DB) (tables.KeyChanges, error) {
s := &keyChangesStatements{
db: db,
writer: writer,
db: db,
}
_, err := db.Exec(keyChangesSchema)
if err != nil {
@ -76,10 +73,8 @@ func NewSqliteKeyChangesTable(db *sql.DB, writer sqlutil.Writer) (tables.KeyChan
}
func (s *keyChangesStatements) InsertKeyChange(ctx context.Context, partition int32, offset int64, userID string) error {
return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
_, err := s.upsertKeyChangeStmt.ExecContext(ctx, partition, offset, userID)
return err
})
_, err := s.upsertKeyChangeStmt.ExecContext(ctx, partition, offset, userID)
return err
}
func (s *keyChangesStatements) SelectKeyChanges(

View file

@ -21,7 +21,6 @@ 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"
)
@ -60,7 +59,6 @@ const selectKeyByAlgorithmSQL = "" +
type oneTimeKeysStatements struct {
db *sql.DB
writer sqlutil.Writer
upsertKeysStmt *sql.Stmt
selectKeysStmt *sql.Stmt
selectKeysCountStmt *sql.Stmt
@ -68,10 +66,9 @@ type oneTimeKeysStatements struct {
deleteOneTimeKeyStmt *sql.Stmt
}
func NewSqliteOneTimeKeysTable(db *sql.DB, writer sqlutil.Writer) (tables.OneTimeKeys, error) {
func NewSqliteOneTimeKeysTable(db *sql.DB) (tables.OneTimeKeys, error) {
s := &oneTimeKeysStatements{
db: db,
writer: writer,
db: db,
}
_, err := db.Exec(oneTimeKeysSchema)
if err != nil {
@ -145,39 +142,39 @@ func (s *oneTimeKeysStatements) CountOneTimeKeys(ctx context.Context, userID, de
return counts, nil
}
func (s *oneTimeKeysStatements) InsertOneTimeKeys(ctx context.Context, keys api.OneTimeKeys) (*api.OneTimeKeysCount, error) {
func (s *oneTimeKeysStatements) InsertOneTimeKeys(
ctx context.Context, txn *sql.Tx, keys api.OneTimeKeys,
) (*api.OneTimeKeysCount, error) {
now := time.Now().Unix()
counts := &api.OneTimeKeysCount{
DeviceID: keys.DeviceID,
UserID: keys.UserID,
KeyCount: make(map[string]int),
}
return counts, s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
for keyIDWithAlgo, keyJSON := range keys.KeyJSON {
algo, keyID := keys.Split(keyIDWithAlgo)
_, err := txn.Stmt(s.upsertKeysStmt).ExecContext(
ctx, keys.UserID, keys.DeviceID, keyID, algo, now, string(keyJSON),
)
if err != nil {
return err
}
}
rows, err := txn.Stmt(s.selectKeysCountStmt).QueryContext(ctx, keys.UserID, keys.DeviceID)
for keyIDWithAlgo, keyJSON := range keys.KeyJSON {
algo, keyID := keys.Split(keyIDWithAlgo)
_, err := txn.Stmt(s.upsertKeysStmt).ExecContext(
ctx, keys.UserID, keys.DeviceID, keyID, algo, now, string(keyJSON),
)
if err != nil {
return err
return nil, err
}
defer internal.CloseAndLogIfError(ctx, rows, "selectKeysCountStmt: rows.close() failed")
for rows.Next() {
var algorithm string
var count int
if err = rows.Scan(&algorithm, &count); err != nil {
return err
}
counts.KeyCount[algorithm] = count
}
rows, err := txn.Stmt(s.selectKeysCountStmt).QueryContext(ctx, keys.UserID, keys.DeviceID)
if err != nil {
return nil, err
}
defer internal.CloseAndLogIfError(ctx, rows, "selectKeysCountStmt: rows.close() failed")
for rows.Next() {
var algorithm string
var count int
if err = rows.Scan(&algorithm, &count); err != nil {
return nil, err
}
counts.KeyCount[algorithm] = count
}
return rows.Err()
})
return counts, rows.Err()
}
func (s *oneTimeKeysStatements) SelectAndDeleteOneTimeKey(
@ -185,17 +182,17 @@ func (s *oneTimeKeysStatements) SelectAndDeleteOneTimeKey(
) (map[string]json.RawMessage, error) {
var keyID string
var keyJSON string
err := s.writer.Do(s.db, txn, func(txn *sql.Tx) error {
err := txn.StmtContext(ctx, s.selectKeyByAlgorithmStmt).QueryRowContext(ctx, userID, deviceID, algorithm).Scan(&keyID, &keyJSON)
if err != nil {
if err == sql.ErrNoRows {
return nil
}
return err
err := txn.StmtContext(ctx, s.selectKeyByAlgorithmStmt).QueryRowContext(ctx, userID, deviceID, algorithm).Scan(&keyID, &keyJSON)
if err != nil {
if err == sql.ErrNoRows {
return nil, nil
}
_, err = txn.StmtContext(ctx, s.deleteOneTimeKeyStmt).ExecContext(ctx, userID, deviceID, algorithm, keyID)
return err
})
return nil, err
}
_, err = txn.StmtContext(ctx, s.deleteOneTimeKeyStmt).ExecContext(ctx, userID, deviceID, algorithm, keyID)
if err != nil {
return nil, err
}
if keyJSON == "" {
return nil, nil
}

View file

@ -20,7 +20,6 @@ import (
"time"
"github.com/matrix-org/dendrite/internal"
"github.com/matrix-org/dendrite/internal/sqlutil"
"github.com/matrix-org/dendrite/keyserver/storage/tables"
"github.com/matrix-org/gomatrixserverlib"
)
@ -51,16 +50,14 @@ const selectStaleDeviceListsSQL = "" +
type staleDeviceListsStatements struct {
db *sql.DB
writer sqlutil.Writer
upsertStaleDeviceListStmt *sql.Stmt
selectStaleDeviceListsWithDomainsStmt *sql.Stmt
selectStaleDeviceListsStmt *sql.Stmt
}
func NewSqliteStaleDeviceListsTable(db *sql.DB, writer sqlutil.Writer) (tables.StaleDeviceLists, error) {
func NewSqliteStaleDeviceListsTable(db *sql.DB) (tables.StaleDeviceLists, error) {
s := &staleDeviceListsStatements{
db: db,
writer: writer,
db: db,
}
_, err := db.Exec(staleDeviceListsSchema)
if err != nil {
@ -83,11 +80,8 @@ func (s *staleDeviceListsStatements) InsertStaleDeviceList(ctx context.Context,
if err != nil {
return err
}
return s.writer.Do(s.db, nil, func(txn *sql.Tx) error {
stmt := sqlutil.TxStmt(txn, s.upsertStaleDeviceListStmt)
_, err = stmt.ExecContext(ctx, userID, string(domain), isStale, time.Now().Unix())
return err
})
_, err = s.upsertStaleDeviceListStmt.ExecContext(ctx, userID, string(domain), isStale, time.Now().Unix())
return err
}
func (s *staleDeviceListsStatements) SelectUserIDsWithStaleDeviceLists(ctx context.Context, domains []gomatrixserverlib.ServerName) ([]string, error) {

View file

@ -25,25 +25,25 @@ func NewDatabase(dbProperties *config.DatabaseOptions) (*shared.Database, error)
if err != nil {
return nil, err
}
writer := sqlutil.NewExclusiveWriter()
otk, err := NewSqliteOneTimeKeysTable(db, writer)
otk, err := NewSqliteOneTimeKeysTable(db)
if err != nil {
return nil, err
}
dk, err := NewSqliteDeviceKeysTable(db, writer)
dk, err := NewSqliteDeviceKeysTable(db)
if err != nil {
return nil, err
}
kc, err := NewSqliteKeyChangesTable(db, writer)
kc, err := NewSqliteKeyChangesTable(db)
if err != nil {
return nil, err
}
sdl, err := NewSqliteStaleDeviceListsTable(db, writer)
sdl, err := NewSqliteStaleDeviceListsTable(db)
if err != nil {
return nil, err
}
return &shared.Database{
DB: db,
Writer: sqlutil.NewExclusiveWriter(),
OneTimeKeysTable: otk,
DeviceKeysTable: dk,
KeyChangesTable: kc,