Storage tweaks (#1373)

* Sync API tweaks

* User API tweaks

4 files changed, 15 insertions, 22 deletions

diff --git a/syncapi/storage/postgres/invites_table.go b/syncapi/storage/postgres/invites_table.go
index eed58c15..c0dd42c5 100644
--- a/syncapi/storage/postgres/invites_table.go
+++ b/syncapi/storage/postgres/invites_table.go
@@ -110,9 +110,10 @@ func (s *inviteEventsStatements) InsertInviteEvent(
 }
 
 func (s *inviteEventsStatements) DeleteInviteEvent(
-	ctx context.Context, inviteEventID string,
+	ctx context.Context, txn *sql.Tx, inviteEventID string,
 ) (sp types.StreamPosition, err error) {
-	err = s.deleteInviteEventStmt.QueryRowContext(ctx, inviteEventID).Scan(&sp)
+	stmt := sqlutil.TxStmt(txn, s.deleteInviteEventStmt)
+	err = stmt.QueryRowContext(ctx, inviteEventID).Scan(&sp)
 	return
 }
 
diff --git a/syncapi/storage/shared/syncserver.go b/syncapi/storage/shared/syncserver.go
index 37793ba2..3388473a 100644
--- a/syncapi/storage/shared/syncserver.go
+++ b/syncapi/storage/shared/syncserver.go
@@ -139,8 +139,8 @@ func (d *Database) GetStateEventsForRoom(
 func (d *Database) AddInviteEvent(
 	ctx context.Context, inviteEvent gomatrixserverlib.HeaderedEvent,
 ) (sp types.StreamPosition, err error) {
-	_ = d.Writer.Do(nil, nil, func(_ *sql.Tx) error {
-		sp, err = d.Invites.InsertInviteEvent(ctx, nil, inviteEvent)
+	_ = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		sp, err = d.Invites.InsertInviteEvent(ctx, txn, inviteEvent)
 		return nil
 	})
 	return
@@ -151,8 +151,8 @@ func (d *Database) AddInviteEvent(
 func (d *Database) RetireInviteEvent(
 	ctx context.Context, inviteEventID string,
 ) (sp types.StreamPosition, err error) {
-	_ = d.Writer.Do(nil, nil, func(_ *sql.Tx) error {
-		sp, err = d.Invites.DeleteInviteEvent(ctx, inviteEventID)
+	_ = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		sp, err = d.Invites.DeleteInviteEvent(ctx, txn, inviteEventID)
 		return nil
 	})
 	return
@@ -422,7 +422,7 @@ func (d *Database) addPDUDeltaToResponse(
 	wantFullState bool,
 	res *types.Response,
 ) (joinedRoomIDs []string, err error) {
-	txn, err := d.DB.BeginTx(context.TODO(), &txReadOnlySnapshot) // TODO: check mattn/go-sqlite3#764
+	txn, err := d.DB.BeginTx(ctx, &txReadOnlySnapshot)
 	if err != nil {
 		return nil, err
 	}
@@ -606,7 +606,7 @@ func (d *Database) getResponseWithPDUsForCompleteSync(
 	// a consistent view of the database throughout. This includes extracting the sync position.
 	// This does have the unfortunate side-effect that all the matrixy logic resides in this function,
 	// but it's better to not hide the fact that this is being done in a transaction.
-	txn, err := d.DB.BeginTx(context.TODO(), &txReadOnlySnapshot) // TODO: check mattn/go-sqlite3#764
+	txn, err := d.DB.BeginTx(ctx, &txReadOnlySnapshot)
 	if err != nil {
 		return
 	}
@@ -1063,15 +1063,6 @@ func (d *Database) SendToDeviceUpdatesWaiting(
 	return count > 0, nil
 }
 
-func (d *Database) AddSendToDeviceEvent(
-	ctx context.Context, txn *sql.Tx,
-	userID, deviceID, content string,
-) error {
-	return d.SendToDevice.InsertSendToDeviceMessage(
-		ctx, txn, userID, deviceID, content,
-	)
-}
-
 func (d *Database) StoreNewSendForDeviceMessage(
 	ctx context.Context, streamPos types.StreamPosition, userID, deviceID string, event gomatrixserverlib.SendToDeviceEvent,
 ) (types.StreamPosition, error) {
@@ -1082,7 +1073,7 @@ func (d *Database) StoreNewSendForDeviceMessage(
 	// Delegate the database write task to the SendToDeviceWriter. It'll guarantee
 	// that we don't lock the table for writes in more than one place.
 	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
-		return d.AddSendToDeviceEvent(
+		return d.SendToDevice.InsertSendToDeviceMessage(
 			ctx, txn, userID, deviceID, string(j),
 		)
 	})
diff --git a/syncapi/storage/sqlite3/invites_table.go b/syncapi/storage/sqlite3/invites_table.go
index 7da86683..1a36ad40 100644
--- a/syncapi/storage/sqlite3/invites_table.go
+++ b/syncapi/storage/sqlite3/invites_table.go
@@ -117,13 +117,14 @@ func (s *inviteEventsStatements) InsertInviteEvent(
 }
 
 func (s *inviteEventsStatements) DeleteInviteEvent(
-	ctx context.Context, inviteEventID string,
+	ctx context.Context, txn *sql.Tx, inviteEventID string,
 ) (types.StreamPosition, error) {
-	streamPos, err := s.streamIDStatements.nextStreamID(ctx, nil)
+	streamPos, err := s.streamIDStatements.nextStreamID(ctx, txn)
 	if err != nil {
 		return streamPos, err
 	}
-	_, err = s.deleteInviteEventStmt.ExecContext(ctx, streamPos, inviteEventID)
+	stmt := sqlutil.TxStmt(txn, s.deleteInviteEventStmt)
+	_, err = stmt.ExecContext(ctx, streamPos, inviteEventID)
 	return streamPos, err
 }
 
diff --git a/syncapi/storage/tables/interface.go b/syncapi/storage/tables/interface.go
index 2ff229cb..38f6d848 100644
--- a/syncapi/storage/tables/interface.go
+++ b/syncapi/storage/tables/interface.go
@@ -32,7 +32,7 @@ type AccountData interface {
 
 type Invites interface {
 	InsertInviteEvent(ctx context.Context, txn *sql.Tx, inviteEvent gomatrixserverlib.HeaderedEvent) (streamPos types.StreamPosition, err error)
-	DeleteInviteEvent(ctx context.Context, inviteEventID string) (types.StreamPosition, error)
+	DeleteInviteEvent(ctx context.Context, txn *sql.Tx, inviteEventID string) (types.StreamPosition, error)
 	// SelectInviteEventsInRange returns a map of room ID to invite events. If multiple invite/retired invites exist in the given range, return the latest value
 	// for the room.
 	SelectInviteEventsInRange(ctx context.Context, txn *sql.Tx, targetUserID string, r types.Range) (invites map[string]gomatrixserverlib.HeaderedEvent, retired map[string]gomatrixserverlib.HeaderedEvent, err error)