From aacb003ad80f74ea9237b5c29472f2119d13daf8 Mon Sep 17 00:00:00 2001 From: Kegan Dougal Date: Wed, 27 May 2020 16:25:32 +0100 Subject: [PATCH] Move Updater structs to shared and use it for postgres --- go.sum | 2 - roomserver/storage/postgres/storage.go | 390 ++---------------- .../storage/shared/membership_updater.go | 144 +++++++ roomserver/storage/shared/prepare.go | 16 + .../shared/room_recent_events_updater.go | 86 ++++ roomserver/storage/shared/storage.go | 67 +++ 6 files changed, 338 insertions(+), 367 deletions(-) create mode 100644 roomserver/storage/shared/membership_updater.go create mode 100644 roomserver/storage/shared/room_recent_events_updater.go diff --git a/go.sum b/go.sum index aa0c8dd57..3398195b6 100644 --- a/go.sum +++ b/go.sum @@ -352,8 +352,6 @@ github.com/matrix-org/dugong v0.0.0-20171220115018-ea0a4690a0d5/go.mod h1:NgPCr+ github.com/matrix-org/dugong v0.0.0-20171220115018-ea0a4690a0d5/go.mod h1:NgPCr+UavRGH6n5jmdX8DuqFZ4JiCWIJoZiuhTRLSUg= github.com/matrix-org/go-http-js-libp2p v0.0.0-20200518170932-783164aeeda4 h1:eqE5OnGx9ZMWmrRbD3KF/3KtTunw0iQulI7YxOIdxo4= github.com/matrix-org/go-http-js-libp2p v0.0.0-20200518170932-783164aeeda4/go.mod h1:3WluEZ9QXSwU30tWYqktnpC1x9mwZKx1r8uAv8Iq+a4= -github.com/matrix-org/go-sqlite3-js v0.0.0-20200326102434-98eda28055bd h1:C1FV4dRKF1uuGK8UH01+IoW6zZpfsTV1MvQimZvt418= -github.com/matrix-org/go-sqlite3-js v0.0.0-20200326102434-98eda28055bd/go.mod h1:e+cg2q7C7yE5QnAXgzo512tgFh1RbQLC0+jozuegKgo= github.com/matrix-org/go-sqlite3-js v0.0.0-20200522092705-bc8506ccbcf3 h1:Yb+Wlf/iHhWlLWd+kCgG+Fsg4Dc+xBl7hptfK7lD0zY= github.com/matrix-org/go-sqlite3-js v0.0.0-20200522092705-bc8506ccbcf3/go.mod h1:e+cg2q7C7yE5QnAXgzo512tgFh1RbQLC0+jozuegKgo= github.com/matrix-org/gomatrix v0.0.0-20190528120928-7df988a63f26 h1:Hr3zjRsq2bhrnp3Ky1qgx/fzCtCALOoGYylh2tpS9K4= diff --git a/roomserver/storage/postgres/storage.go b/roomserver/storage/postgres/storage.go index 53a58076a..971f2b9e6 100644 --- a/roomserver/storage/postgres/storage.go +++ b/roomserver/storage/postgres/storage.go @@ -9,14 +9,13 @@ // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, -// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implie // See the License for the specific language governing permissions and // limitations under the License. package postgres import ( - "context" "database/sql" "github.com/matrix-org/dendrite/internal" @@ -25,423 +24,84 @@ import ( // Import the postgres database driver. _ "github.com/lib/pq" "github.com/matrix-org/dendrite/roomserver/storage/shared" - "github.com/matrix-org/dendrite/roomserver/storage/tables" - "github.com/matrix-org/dendrite/roomserver/types" - "github.com/matrix-org/gomatrixserverlib" ) // A Database is used to store room events and stream offsets. type Database struct { shared.Database - events tables.Events - eventTypes tables.EventTypes - eventStateKeys tables.EventStateKeys - eventJSON tables.EventJSON - rooms tables.Rooms - transactions tables.Transactions - prevEvents tables.PreviousEvents - invites tables.Invites - membership tables.Membership - db *sql.DB } // Open a postgres database. // nolint: gocyclo func Open(dataSourceName string, dbProperties internal.DbProperties) (*Database, error) { var d Database + var db *sql.DB var err error - if d.db, err = sqlutil.Open("postgres", dataSourceName, dbProperties); err != nil { + if db, err = sqlutil.Open("postgres", dataSourceName, dbProperties); err != nil { return nil, err } - d.eventStateKeys, err = NewPostgresEventStateKeysTable(d.db) + eventStateKeys, err := NewPostgresEventStateKeysTable(db) if err != nil { return nil, err } - d.eventTypes, err = NewPostgresEventTypesTable(d.db) + eventTypes, err := NewPostgresEventTypesTable(db) if err != nil { return nil, err } - d.eventJSON, err = NewPostgresEventJSONTable(d.db) + eventJSON, err := NewPostgresEventJSONTable(db) if err != nil { return nil, err } - d.events, err = NewPostgresEventsTable(d.db) + events, err := NewPostgresEventsTable(db) if err != nil { return nil, err } - d.rooms, err = NewPostgresRoomsTable(d.db) + rooms, err := NewPostgresRoomsTable(db) if err != nil { return nil, err } - d.transactions, err = NewPostgresTransactionsTable(d.db) + transactions, err := NewPostgresTransactionsTable(db) if err != nil { return nil, err } - stateBlock, err := NewPostgresStateBlockTable(d.db) + stateBlock, err := NewPostgresStateBlockTable(db) if err != nil { return nil, err } - stateSnapshot, err := NewPostgresStateSnapshotTable(d.db) + stateSnapshot, err := NewPostgresStateSnapshotTable(db) if err != nil { return nil, err } - roomAliases, err := NewPostgresRoomAliasesTable(d.db) + roomAliases, err := NewPostgresRoomAliasesTable(db) if err != nil { return nil, err } - d.prevEvents, err = NewPostgresPreviousEventsTable(d.db) + prevEvents, err := NewPostgresPreviousEventsTable(db) if err != nil { return nil, err } - d.invites, err = NewPostgresInvitesTable(d.db) + invites, err := NewPostgresInvitesTable(db) if err != nil { return nil, err } - d.membership, err = NewPostgresMembershipTable(d.db) + membership, err := NewPostgresMembershipTable(db) if err != nil { return nil, err } d.Database = shared.Database{ - DB: d.db, - EventTypesTable: d.eventTypes, - EventStateKeysTable: d.eventStateKeys, - EventJSONTable: d.eventJSON, - EventsTable: d.events, - RoomsTable: d.rooms, - TransactionsTable: d.transactions, + DB: db, + EventTypesTable: eventTypes, + EventStateKeysTable: eventStateKeys, + EventJSONTable: eventJSON, + EventsTable: events, + RoomsTable: rooms, + TransactionsTable: transactions, StateBlockTable: stateBlock, StateSnapshotTable: stateSnapshot, - PrevEventsTable: d.prevEvents, + PrevEventsTable: prevEvents, RoomAliasesTable: roomAliases, - InvitesTable: d.invites, - MembershipTable: d.membership, + InvitesTable: invites, + MembershipTable: membership, } return &d, nil } - -func (d *Database) assignRoomNID( - ctx context.Context, txn *sql.Tx, - roomID string, roomVersion gomatrixserverlib.RoomVersion, -) (types.RoomNID, error) { - // Check if we already have a numeric ID in the database. - roomNID, err := d.rooms.SelectRoomNID(ctx, txn, roomID) - if err == sql.ErrNoRows { - // We don't have a numeric ID so insert one into the database. - roomNID, err = d.rooms.InsertRoomNID(ctx, txn, roomID, roomVersion) - if err == sql.ErrNoRows { - // We raced with another insert so run the select again. - roomNID, err = d.rooms.SelectRoomNID(ctx, txn, roomID) - } - } - return roomNID, err -} - -func (d *Database) assignStateKeyNID( - ctx context.Context, txn *sql.Tx, eventStateKey string, -) (types.EventStateKeyNID, error) { - // Check if we already have a numeric ID in the database. - eventStateKeyNID, err := d.eventStateKeys.SelectEventStateKeyNID(ctx, txn, eventStateKey) - if err == sql.ErrNoRows { - // We don't have a numeric ID so insert one into the database. - eventStateKeyNID, err = d.eventStateKeys.InsertEventStateKeyNID(ctx, txn, eventStateKey) - if err == sql.ErrNoRows { - // We raced with another insert so run the select again. - eventStateKeyNID, err = d.eventStateKeys.SelectEventStateKeyNID(ctx, txn, eventStateKey) - } - } - return eventStateKeyNID, err -} - -// GetLatestEventsForUpdate implements input.EventDatabase -func (d *Database) GetLatestEventsForUpdate( - ctx context.Context, roomNID types.RoomNID, -) (types.RoomRecentEventsUpdater, error) { - txn, err := d.db.Begin() - if err != nil { - return nil, err - } - eventNIDs, lastEventNIDSent, currentStateSnapshotNID, err := - d.rooms.SelectLatestEventsNIDsForUpdate(ctx, txn, roomNID) - if err != nil { - txn.Rollback() // nolint: errcheck - return nil, err - } - stateAndRefs, err := d.events.BulkSelectStateAtEventAndReference(ctx, txn, eventNIDs) - if err != nil { - txn.Rollback() // nolint: errcheck - return nil, err - } - var lastEventIDSent string - if lastEventNIDSent != 0 { - lastEventIDSent, err = d.events.SelectEventID(ctx, txn, lastEventNIDSent) - if err != nil { - txn.Rollback() // nolint: errcheck - return nil, err - } - } - return &roomRecentEventsUpdater{ - transaction{ctx, txn}, d, roomNID, stateAndRefs, lastEventIDSent, currentStateSnapshotNID, - }, nil -} - -type roomRecentEventsUpdater struct { - transaction - d *Database - roomNID types.RoomNID - latestEvents []types.StateAtEventAndReference - lastEventIDSent string - currentStateSnapshotNID types.StateSnapshotNID -} - -// RoomVersion implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) RoomVersion() (version gomatrixserverlib.RoomVersion) { - version, _ = u.d.GetRoomVersionForRoomNID(u.ctx, u.roomNID) - return -} - -// LatestEvents implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) LatestEvents() []types.StateAtEventAndReference { - return u.latestEvents -} - -// LastEventIDSent implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) LastEventIDSent() string { - return u.lastEventIDSent -} - -// CurrentStateSnapshotNID implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) CurrentStateSnapshotNID() types.StateSnapshotNID { - return u.currentStateSnapshotNID -} - -// StorePreviousEvents implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) StorePreviousEvents(eventNID types.EventNID, previousEventReferences []gomatrixserverlib.EventReference) error { - for _, ref := range previousEventReferences { - if err := u.d.prevEvents.InsertPreviousEvent(u.ctx, u.txn, ref.EventID, ref.EventSHA256, eventNID); err != nil { - return err - } - } - return nil -} - -// IsReferenced implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) IsReferenced(eventReference gomatrixserverlib.EventReference) (bool, error) { - err := u.d.prevEvents.SelectPreviousEventExists(u.ctx, u.txn, eventReference.EventID, eventReference.EventSHA256) - if err == nil { - return true, nil - } - if err == sql.ErrNoRows { - return false, nil - } - return false, err -} - -// SetLatestEvents implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) SetLatestEvents( - roomNID types.RoomNID, latest []types.StateAtEventAndReference, lastEventNIDSent types.EventNID, - currentStateSnapshotNID types.StateSnapshotNID, -) error { - eventNIDs := make([]types.EventNID, len(latest)) - for i := range latest { - eventNIDs[i] = latest[i].EventNID - } - return u.d.rooms.UpdateLatestEventNIDs(u.ctx, u.txn, roomNID, eventNIDs, lastEventNIDSent, currentStateSnapshotNID) -} - -// HasEventBeenSent implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) HasEventBeenSent(eventNID types.EventNID) (bool, error) { - return u.d.events.SelectEventSentToOutput(u.ctx, u.txn, eventNID) -} - -// MarkEventAsSent implements types.RoomRecentEventsUpdater -func (u *roomRecentEventsUpdater) MarkEventAsSent(eventNID types.EventNID) error { - return u.d.events.UpdateEventSentToOutput(u.ctx, u.txn, eventNID) -} - -func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (types.MembershipUpdater, error) { - return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomNID, targetUserNID, targetLocal) -} - -// MembershipUpdater implements input.RoomEventDatabase -func (d *Database) MembershipUpdater( - ctx context.Context, roomID, targetUserID string, - targetLocal bool, roomVersion gomatrixserverlib.RoomVersion, -) (types.MembershipUpdater, error) { - txn, err := d.db.Begin() - if err != nil { - return nil, err - } - succeeded := false - defer func() { - if !succeeded { - txn.Rollback() // nolint: errcheck - } - }() - - roomNID, err := d.assignRoomNID(ctx, txn, roomID, roomVersion) - if err != nil { - return nil, err - } - - targetUserNID, err := d.assignStateKeyNID(ctx, txn, targetUserID) - if err != nil { - return nil, err - } - - updater, err := d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID, targetLocal) - if err != nil { - return nil, err - } - - succeeded = true - return updater, nil -} - -type membershipUpdater struct { - transaction - d *Database - roomNID types.RoomNID - targetUserNID types.EventStateKeyNID - membership tables.MembershipState -} - -func (d *Database) membershipUpdaterTxn( - ctx context.Context, - txn *sql.Tx, - roomNID types.RoomNID, - targetUserNID types.EventStateKeyNID, - targetLocal bool, -) (types.MembershipUpdater, error) { - - if err := d.membership.InsertMembership(ctx, txn, roomNID, targetUserNID, targetLocal); err != nil { - return nil, err - } - - membership, err := d.membership.SelectMembershipForUpdate(ctx, txn, roomNID, targetUserNID) - if err != nil { - return nil, err - } - - return &membershipUpdater{ - transaction{ctx, txn}, d, roomNID, targetUserNID, membership, - }, nil -} - -// IsInvite implements types.MembershipUpdater -func (u *membershipUpdater) IsInvite() bool { - return u.membership == tables.MembershipStateInvite -} - -// IsJoin implements types.MembershipUpdater -func (u *membershipUpdater) IsJoin() bool { - return u.membership == tables.MembershipStateJoin -} - -// IsLeave implements types.MembershipUpdater -func (u *membershipUpdater) IsLeave() bool { - return u.membership == tables.MembershipStateLeaveOrBan -} - -// SetToInvite implements types.MembershipUpdater -func (u *membershipUpdater) SetToInvite(event gomatrixserverlib.Event) (bool, error) { - senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, event.Sender()) - if err != nil { - return false, err - } - inserted, err := u.d.invites.InsertInviteEvent( - u.ctx, u.txn, event.EventID(), u.roomNID, u.targetUserNID, senderUserNID, event.JSON(), - ) - if err != nil { - return false, err - } - if u.membership != tables.MembershipStateInvite { - if err = u.d.membership.UpdateMembership( - u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateInvite, 0, - ); err != nil { - return false, err - } - } - return inserted, nil -} - -// SetToJoin implements types.MembershipUpdater -func (u *membershipUpdater) SetToJoin(senderUserID string, eventID string, isUpdate bool) ([]string, error) { - var inviteEventIDs []string - - senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID) - if err != nil { - return nil, err - } - - // If this is a join event update, there is no invite to update - if !isUpdate { - inviteEventIDs, err = u.d.invites.UpdateInviteRetired( - u.ctx, u.txn, u.roomNID, u.targetUserNID, - ) - if err != nil { - return nil, err - } - } - - // Look up the NID of the new join event - nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID}) - if err != nil { - return nil, err - } - - if u.membership != tables.MembershipStateJoin || isUpdate { - if err = u.d.membership.UpdateMembership( - u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, - tables.MembershipStateJoin, nIDs[eventID], - ); err != nil { - return nil, err - } - } - - return inviteEventIDs, nil -} - -// SetToLeave implements types.MembershipUpdater -func (u *membershipUpdater) SetToLeave(senderUserID string, eventID string) ([]string, error) { - senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID) - if err != nil { - return nil, err - } - inviteEventIDs, err := u.d.invites.UpdateInviteRetired( - u.ctx, u.txn, u.roomNID, u.targetUserNID, - ) - if err != nil { - return nil, err - } - - // Look up the NID of the new leave event - nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID}) - if err != nil { - return nil, err - } - - if u.membership != tables.MembershipStateLeaveOrBan { - if err = u.d.membership.UpdateMembership( - u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, - tables.MembershipStateLeaveOrBan, nIDs[eventID], - ); err != nil { - return nil, err - } - } - return inviteEventIDs, nil -} - -type transaction struct { - ctx context.Context - txn *sql.Tx -} - -// Commit implements types.Transaction -func (t *transaction) Commit() error { - return t.txn.Commit() -} - -// Rollback implements types.Transaction -func (t *transaction) Rollback() error { - return t.txn.Rollback() -} diff --git a/roomserver/storage/shared/membership_updater.go b/roomserver/storage/shared/membership_updater.go new file mode 100644 index 000000000..8c777381a --- /dev/null +++ b/roomserver/storage/shared/membership_updater.go @@ -0,0 +1,144 @@ +package shared + +import ( + "context" + "database/sql" + + "github.com/matrix-org/dendrite/roomserver/storage/tables" + "github.com/matrix-org/dendrite/roomserver/types" + "github.com/matrix-org/gomatrixserverlib" +) + +type membershipUpdater struct { + transaction + d *Database + roomNID types.RoomNID + targetUserNID types.EventStateKeyNID + membership tables.MembershipState +} + +func (d *Database) membershipUpdaterTxn( + ctx context.Context, + txn *sql.Tx, + roomNID types.RoomNID, + targetUserNID types.EventStateKeyNID, + targetLocal bool, +) (types.MembershipUpdater, error) { + + if err := d.MembershipTable.InsertMembership(ctx, txn, roomNID, targetUserNID, targetLocal); err != nil { + return nil, err + } + + membership, err := d.MembershipTable.SelectMembershipForUpdate(ctx, txn, roomNID, targetUserNID) + if err != nil { + return nil, err + } + + return &membershipUpdater{ + transaction{ctx, txn}, d, roomNID, targetUserNID, membership, + }, nil +} + +// IsInvite implements types.MembershipUpdater +func (u *membershipUpdater) IsInvite() bool { + return u.membership == tables.MembershipStateInvite +} + +// IsJoin implements types.MembershipUpdater +func (u *membershipUpdater) IsJoin() bool { + return u.membership == tables.MembershipStateJoin +} + +// IsLeave implements types.MembershipUpdater +func (u *membershipUpdater) IsLeave() bool { + return u.membership == tables.MembershipStateLeaveOrBan +} + +// SetToInvite implements types.MembershipUpdater +func (u *membershipUpdater) SetToInvite(event gomatrixserverlib.Event) (bool, error) { + senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, event.Sender()) + if err != nil { + return false, err + } + inserted, err := u.d.InvitesTable.InsertInviteEvent( + u.ctx, u.txn, event.EventID(), u.roomNID, u.targetUserNID, senderUserNID, event.JSON(), + ) + if err != nil { + return false, err + } + if u.membership != tables.MembershipStateInvite { + if err = u.d.MembershipTable.UpdateMembership( + u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateInvite, 0, + ); err != nil { + return false, err + } + } + return inserted, nil +} + +// SetToJoin implements types.MembershipUpdater +func (u *membershipUpdater) SetToJoin(senderUserID string, eventID string, isUpdate bool) ([]string, error) { + var inviteEventIDs []string + + senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID) + if err != nil { + return nil, err + } + + // If this is a join event update, there is no invite to update + if !isUpdate { + inviteEventIDs, err = u.d.InvitesTable.UpdateInviteRetired( + u.ctx, u.txn, u.roomNID, u.targetUserNID, + ) + if err != nil { + return nil, err + } + } + + // Look up the NID of the new join event + nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID}) + if err != nil { + return nil, err + } + + if u.membership != tables.MembershipStateJoin || isUpdate { + if err = u.d.MembershipTable.UpdateMembership( + u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, + tables.MembershipStateJoin, nIDs[eventID], + ); err != nil { + return nil, err + } + } + + return inviteEventIDs, nil +} + +// SetToLeave implements types.MembershipUpdater +func (u *membershipUpdater) SetToLeave(senderUserID string, eventID string) ([]string, error) { + senderUserNID, err := u.d.assignStateKeyNID(u.ctx, u.txn, senderUserID) + if err != nil { + return nil, err + } + inviteEventIDs, err := u.d.InvitesTable.UpdateInviteRetired( + u.ctx, u.txn, u.roomNID, u.targetUserNID, + ) + if err != nil { + return nil, err + } + + // Look up the NID of the new leave event + nIDs, err := u.d.EventNIDs(u.ctx, []string{eventID}) + if err != nil { + return nil, err + } + + if u.membership != tables.MembershipStateLeaveOrBan { + if err = u.d.MembershipTable.UpdateMembership( + u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, + tables.MembershipStateLeaveOrBan, nIDs[eventID], + ); err != nil { + return nil, err + } + } + return inviteEventIDs, nil +} diff --git a/roomserver/storage/shared/prepare.go b/roomserver/storage/shared/prepare.go index 1b3497fd8..e048d43af 100644 --- a/roomserver/storage/shared/prepare.go +++ b/roomserver/storage/shared/prepare.go @@ -16,6 +16,7 @@ package shared import ( + "context" "database/sql" ) @@ -34,3 +35,18 @@ func (s StatementList) Prepare(db *sql.DB) (err error) { } return } + +type transaction struct { + ctx context.Context + txn *sql.Tx +} + +// Commit implements types.Transaction +func (t *transaction) Commit() error { + return t.txn.Commit() +} + +// Rollback implements types.Transaction +func (t *transaction) Rollback() error { + return t.txn.Rollback() +} diff --git a/roomserver/storage/shared/room_recent_events_updater.go b/roomserver/storage/shared/room_recent_events_updater.go new file mode 100644 index 000000000..04c254f92 --- /dev/null +++ b/roomserver/storage/shared/room_recent_events_updater.go @@ -0,0 +1,86 @@ +package shared + +import ( + "database/sql" + + "github.com/matrix-org/dendrite/roomserver/types" + "github.com/matrix-org/gomatrixserverlib" +) + +type roomRecentEventsUpdater struct { + transaction + d *Database + roomNID types.RoomNID + latestEvents []types.StateAtEventAndReference + lastEventIDSent string + currentStateSnapshotNID types.StateSnapshotNID +} + +// RoomVersion implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) RoomVersion() (version gomatrixserverlib.RoomVersion) { + version, _ = u.d.GetRoomVersionForRoomNID(u.ctx, u.roomNID) + return +} + +// LatestEvents implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) LatestEvents() []types.StateAtEventAndReference { + return u.latestEvents +} + +// LastEventIDSent implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) LastEventIDSent() string { + return u.lastEventIDSent +} + +// CurrentStateSnapshotNID implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) CurrentStateSnapshotNID() types.StateSnapshotNID { + return u.currentStateSnapshotNID +} + +// StorePreviousEvents implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) StorePreviousEvents(eventNID types.EventNID, previousEventReferences []gomatrixserverlib.EventReference) error { + for _, ref := range previousEventReferences { + if err := u.d.PrevEventsTable.InsertPreviousEvent(u.ctx, u.txn, ref.EventID, ref.EventSHA256, eventNID); err != nil { + return err + } + } + return nil +} + +// IsReferenced implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) IsReferenced(eventReference gomatrixserverlib.EventReference) (bool, error) { + err := u.d.PrevEventsTable.SelectPreviousEventExists(u.ctx, u.txn, eventReference.EventID, eventReference.EventSHA256) + if err == nil { + return true, nil + } + if err == sql.ErrNoRows { + return false, nil + } + return false, err +} + +// SetLatestEvents implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) SetLatestEvents( + roomNID types.RoomNID, latest []types.StateAtEventAndReference, lastEventNIDSent types.EventNID, + currentStateSnapshotNID types.StateSnapshotNID, +) error { + eventNIDs := make([]types.EventNID, len(latest)) + for i := range latest { + eventNIDs[i] = latest[i].EventNID + } + return u.d.RoomsTable.UpdateLatestEventNIDs(u.ctx, u.txn, roomNID, eventNIDs, lastEventNIDSent, currentStateSnapshotNID) +} + +// HasEventBeenSent implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) HasEventBeenSent(eventNID types.EventNID) (bool, error) { + return u.d.EventsTable.SelectEventSentToOutput(u.ctx, u.txn, eventNID) +} + +// MarkEventAsSent implements types.RoomRecentEventsUpdater +func (u *roomRecentEventsUpdater) MarkEventAsSent(eventNID types.EventNID) error { + return u.d.EventsTable.UpdateEventSentToOutput(u.ctx, u.txn, eventNID) +} + +func (u *roomRecentEventsUpdater) MembershipUpdater(targetUserNID types.EventStateKeyNID, targetLocal bool) (types.MembershipUpdater, error) { + return u.d.membershipUpdaterTxn(u.ctx, u.txn, u.roomNID, targetUserNID, targetLocal) +} diff --git a/roomserver/storage/shared/storage.go b/roomserver/storage/shared/storage.go index 17bcd6964..89ae4fc37 100644 --- a/roomserver/storage/shared/storage.go +++ b/roomserver/storage/shared/storage.go @@ -311,6 +311,73 @@ func (d *Database) GetTransactionEventID( return eventID, err } +// MembershipUpdater implements input.RoomEventDatabase +func (d *Database) MembershipUpdater( + ctx context.Context, roomID, targetUserID string, + targetLocal bool, roomVersion gomatrixserverlib.RoomVersion, +) (types.MembershipUpdater, error) { + txn, err := d.DB.Begin() + if err != nil { + return nil, err + } + succeeded := false + defer func() { + if !succeeded { + txn.Rollback() // nolint: errcheck + } + }() + + roomNID, err := d.assignRoomNID(ctx, txn, roomID, roomVersion) + if err != nil { + return nil, err + } + + targetUserNID, err := d.assignStateKeyNID(ctx, txn, targetUserID) + if err != nil { + return nil, err + } + + updater, err := d.membershipUpdaterTxn(ctx, txn, roomNID, targetUserNID, targetLocal) + if err != nil { + return nil, err + } + + succeeded = true + return updater, nil +} + +// GetLatestEventsForUpdate implements input.EventDatabase +func (d *Database) GetLatestEventsForUpdate( + ctx context.Context, roomNID types.RoomNID, +) (types.RoomRecentEventsUpdater, error) { + txn, err := d.DB.Begin() + if err != nil { + return nil, err + } + eventNIDs, lastEventNIDSent, currentStateSnapshotNID, err := + d.RoomsTable.SelectLatestEventsNIDsForUpdate(ctx, txn, roomNID) + if err != nil { + txn.Rollback() // nolint: errcheck + return nil, err + } + stateAndRefs, err := d.EventsTable.BulkSelectStateAtEventAndReference(ctx, txn, eventNIDs) + if err != nil { + txn.Rollback() // nolint: errcheck + return nil, err + } + var lastEventIDSent string + if lastEventNIDSent != 0 { + lastEventIDSent, err = d.EventsTable.SelectEventID(ctx, txn, lastEventNIDSent) + if err != nil { + txn.Rollback() // nolint: errcheck + return nil, err + } + } + return &roomRecentEventsUpdater{ + transaction{ctx, txn}, d, roomNID, stateAndRefs, lastEventIDSent, currentStateSnapshotNID, + }, nil +} + func (d *Database) StoreEvent( ctx context.Context, event gomatrixserverlib.Event, txnAndSessionID *api.TransactionID, authEventNIDs []types.EventNID,