Add basic storage methods

roomserver/storage/interface.go

@@ -126,7 +126,7 @@ type Database interface {
 	// in this room, along a boolean set to true if the user is still in this room,
 	// false if not.
 	// Returns an error if there was a problem talking to the database.
-	GetMembership(ctx context.Context, roomNID types.RoomNID, requestSenderUserID string) (membershipEventNID types.EventNID, stillInRoom bool, err error)
+	GetMembership(ctx context.Context, roomNID types.RoomNID, requestSenderUserID string) (membershipEventNID types.EventNID, stillInRoom, isRoomForgotten bool, err error)
 	// Lookup the membership event numeric IDs for all user that are or have
 	// been members of a given room. Only lookup events of "join" membership if
 	// joinOnly is set to true.
@@ -158,4 +158,6 @@ type Database interface {
 	GetKnownUsers(ctx context.Context, userID, searchString string, limit int) ([]string, error)
 	// GetKnownRooms returns a list of all rooms we know about.
 	GetKnownRooms(ctx context.Context) ([]string, error)
+	// ForgetRoom sets a flag in the membership table, that the user wishes to forget a specific room
+	ForgetRoom(ctx context.Context, userID, roomID string, forget bool) error
 }

roomserver/storage/postgres/membership_table.go

@@ -60,13 +60,15 @@ CREATE TABLE IF NOT EXISTS roomserver_membership (
 	-- a federated one. This is an optimisation for resetting state on federated
 	-- room joins.
 	target_local BOOLEAN NOT NULL DEFAULT false,
+	forgotten BOOLEAN NOT NULL DEFAULT FALSE,
 	UNIQUE (room_nid, target_nid)
 );
 `
 
 var selectJoinedUsersSetForRoomsSQL = "" +
 	"SELECT target_nid, COUNT(room_nid) FROM roomserver_membership WHERE room_nid = ANY($1) AND" +
-	" membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " GROUP BY target_nid"
+	" membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " and forgotten = false" +
+	" GROUP BY target_nid"
 
 // Insert a row in to membership table so that it can be locked by the
 // SELECT FOR UPDATE
@@ -76,37 +78,41 @@ const insertMembershipSQL = "" +
 	" ON CONFLICT DO NOTHING"
 
 const selectMembershipFromRoomAndTargetSQL = "" +
-	"SELECT membership_nid, event_nid FROM roomserver_membership" +
+	"SELECT membership_nid, event_nid, forgotten FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND target_nid = $2"
 
 const selectMembershipsFromRoomAndMembershipSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
-	" WHERE room_nid = $1 AND membership_nid = $2"
+	" WHERE room_nid = $1 AND membership_nid = $2 and forgotten = false"
 
 const selectLocalMembershipsFromRoomAndMembershipSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND membership_nid = $2" +
-	" AND target_local = true"
+	" AND target_local = true and forgotten = false"
 
 const selectMembershipsFromRoomSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
-	" WHERE room_nid = $1"
+	" WHERE room_nid = $1 and forgotten = false"
 
 const selectLocalMembershipsFromRoomSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1" +
-	" AND target_local = true"
+	" AND target_local = true and forgotten = false"
 
 const selectMembershipForUpdateSQL = "" +
 	"SELECT membership_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND target_nid = $2 FOR UPDATE"
 
 const updateMembershipSQL = "" +
-	"UPDATE roomserver_membership SET sender_nid = $3, membership_nid = $4, event_nid = $5" +
+	"UPDATE roomserver_membership SET sender_nid = $3, membership_nid = $4, event_nid = $5, forgotten = $6" +
+	" WHERE room_nid = $1 AND target_nid = $2"
+
+const updateMembershipForgetRoom = "" +
+	"UPDATE roomserver_membership SET forgotten = $3" +
 	" WHERE room_nid = $1 AND target_nid = $2"
 
 const selectRoomsWithMembershipSQL = "" +
-	"SELECT room_nid FROM roomserver_membership WHERE membership_nid = $1 AND target_nid = $2"
+	"SELECT room_nid FROM roomserver_membership WHERE membership_nid = $1 AND target_nid = $2 and forgotten = false"
 
 // selectKnownUsersSQL uses a sub-select statement here to find rooms that the user is
 // joined to. Since this information is used to populate the user directory, we will
@@ -130,6 +136,7 @@ type membershipStatements struct {
 	selectRoomsWithMembershipStmt                   *sql.Stmt
 	selectJoinedUsersSetForRoomsStmt                *sql.Stmt
 	selectKnownUsersStmt                            *sql.Stmt
+	updateMembershipForgetRoomStmt                  *sql.Stmt
 }
 
 func NewPostgresMembershipTable(db *sql.DB) (tables.Membership, error) {
@@ -151,6 +158,7 @@ func NewPostgresMembershipTable(db *sql.DB) (tables.Membership, error) {
 		{&s.selectRoomsWithMembershipStmt, selectRoomsWithMembershipSQL},
 		{&s.selectJoinedUsersSetForRoomsStmt, selectJoinedUsersSetForRoomsSQL},
 		{&s.selectKnownUsersStmt, selectKnownUsersSQL},
+		{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
 	}.Prepare(db)
 }
 
@@ -177,10 +185,10 @@ func (s *membershipStatements) SelectMembershipForUpdate(
 func (s *membershipStatements) SelectMembershipFromRoomAndTarget(
 	ctx context.Context,
 	roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
-) (eventNID types.EventNID, membership tables.MembershipState, err error) {
+) (eventNID types.EventNID, membership tables.MembershipState, forgotten bool, err error) {
 	err = s.selectMembershipFromRoomAndTargetStmt.QueryRowContext(
 		ctx, roomNID, targetUserNID,
-	).Scan(&membership, &eventNID)
+	).Scan(&membership, &eventNID, &forgotten)
 	return
 }
 
@@ -238,12 +246,11 @@ func (s *membershipStatements) SelectMembershipsFromRoomAndMembership(
 
 func (s *membershipStatements) UpdateMembership(
 	ctx context.Context,
-	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
+	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, senderUserNID types.EventStateKeyNID, membership tables.MembershipState,
-	senderUserNID types.EventStateKeyNID, membership tables.MembershipState,
+	eventNID types.EventNID, forgotten bool,
-	eventNID types.EventNID,
 ) error {
 	_, err := sqlutil.TxStmt(txn, s.updateMembershipStmt).ExecContext(
-		ctx, roomNID, targetUserNID, senderUserNID, membership, eventNID,
+		ctx, roomNID, targetUserNID, senderUserNID, membership, eventNID, forgotten,
 	)
 	return err
 }
@@ -305,3 +312,14 @@ func (s *membershipStatements) SelectKnownUsers(ctx context.Context, userID type
 	}
 	return result, rows.Err()
 }
+
+func (s *membershipStatements) UpdateForgetMembership(
+	ctx context.Context,
+	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
+	forget bool,
+) error {
+	_, err := sqlutil.TxStmt(txn, s.updateMembershipForgetRoomStmt).ExecContext(
+		ctx, roomNID, targetUserNID, forget,
+	)
+	return err
+}

roomserver/storage/shared/membership_updater.go

@@ -101,9 +101,7 @@ func (u *MembershipUpdater) SetToInvite(event gomatrixserverlib.Event) (bool, er
 			return fmt.Errorf("u.d.InvitesTable.InsertInviteEvent: %w", err)
 		}
 		if u.membership != tables.MembershipStateInvite {
-			if err = u.d.MembershipTable.UpdateMembership(
+			if err = u.d.MembershipTable.UpdateMembership(u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateInvite, 0, false); err != nil {
-				u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateInvite, 0,
-			); err != nil {
 				return fmt.Errorf("u.d.MembershipTable.UpdateMembership: %w", err)
 			}
 		}
@@ -139,10 +137,7 @@ func (u *MembershipUpdater) SetToJoin(senderUserID string, eventID string, isUpd
 		}
 
 		if u.membership != tables.MembershipStateJoin || isUpdate {
-			if err = u.d.MembershipTable.UpdateMembership(
+			if err = u.d.MembershipTable.UpdateMembership(u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateJoin, nIDs[eventID], false); err != nil {
-				u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID,
-				tables.MembershipStateJoin, nIDs[eventID],
-			); err != nil {
 				return fmt.Errorf("u.d.MembershipTable.UpdateMembership: %w", err)
 			}
 		}
@@ -176,10 +171,7 @@ func (u *MembershipUpdater) SetToLeave(senderUserID string, eventID string) ([]s
 		}
 
 		if u.membership != tables.MembershipStateLeaveOrBan {
-			if err = u.d.MembershipTable.UpdateMembership(
+			if err = u.d.MembershipTable.UpdateMembership(u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID, tables.MembershipStateLeaveOrBan, nIDs[eventID], false); err != nil {
-				u.ctx, u.txn, u.roomNID, u.targetUserNID, senderUserNID,
-				tables.MembershipStateLeaveOrBan, nIDs[eventID],
-			); err != nil {
 				return fmt.Errorf("u.d.MembershipTable.UpdateMembership: %w", err)
 			}
 		}

roomserver/storage/shared/storage.go

@@ -258,30 +258,28 @@ func (d *Database) RemoveRoomAlias(ctx context.Context, alias string) error {
 	})
 }
 
-func (d *Database) GetMembership(
+func (d *Database) GetMembership(ctx context.Context, roomNID types.RoomNID, requestSenderUserID string) (membershipEventNID types.EventNID, stillInRoom, isRoomforgotten bool, err error) {
-	ctx context.Context, roomNID types.RoomNID, requestSenderUserID string,
-) (membershipEventNID types.EventNID, stillInRoom bool, err error) {
 	var requestSenderUserNID types.EventStateKeyNID
 	err = d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
 		requestSenderUserNID, err = d.assignStateKeyNID(ctx, txn, requestSenderUserID)
 		return err
 	})
 	if err != nil {
-		return 0, false, fmt.Errorf("d.assignStateKeyNID: %w", err)
+		return 0, false, false, fmt.Errorf("d.assignStateKeyNID: %w", err)
 	}
 
-	senderMembershipEventNID, senderMembership, err :=
+	senderMembershipEventNID, senderMembership, isRoomforgotten, err :=
 		d.MembershipTable.SelectMembershipFromRoomAndTarget(
 			ctx, roomNID, requestSenderUserNID,
 		)
 	if err == sql.ErrNoRows {
 		// The user has never been a member of that room
-		return 0, false, nil
+		return 0, false, false, nil
 	} else if err != nil {
 		return
 	}
 
-	return senderMembershipEventNID, senderMembership == tables.MembershipStateJoin, nil
+	return senderMembershipEventNID, senderMembership == tables.MembershipStateJoin, isRoomforgotten, nil
 }
 
 func (d *Database) GetMembershipEventNIDsForRoom(
@@ -992,6 +990,25 @@ func (d *Database) GetKnownRooms(ctx context.Context) ([]string, error) {
 	return d.RoomsTable.SelectRoomIDs(ctx)
 }
 
+// ForgetRoom sets a users room to forgotten
+func (d *Database) ForgetRoom(ctx context.Context, userID, roomID string, forget bool) error {
+	roomNIDs, err := d.RoomsTable.BulkSelectRoomNIDs(ctx, []string{roomID})
+	if err != nil {
+		return err
+	}
+	if len(roomNIDs) > 1 {
+		return fmt.Errorf("expected one room, got %d", len(roomNIDs))
+	}
+	stateKeyNID, err := d.EventStateKeysTable.SelectEventStateKeyNID(ctx, nil, userID)
+	if err != nil {
+		return err
+	}
+
+	return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
+		return d.MembershipTable.UpdateForgetMembership(ctx, nil, roomNIDs[0], stateKeyNID, forget)
+	})
+}
+
 // FIXME TODO: Remove all this - horrible dupe with roomserver/state. Can't use the original impl because of circular loops
 // it should live in this package!
 

roomserver/storage/sqlite3/membership_table.go

@@ -36,13 +36,15 @@ const membershipSchema = `
 		membership_nid INTEGER NOT NULL DEFAULT 1,
 		event_nid INTEGER NOT NULL DEFAULT 0,
 		target_local BOOLEAN NOT NULL DEFAULT false,
+		forgotten BOOLEAN NOT NULL DEFAULT false,
 		UNIQUE (room_nid, target_nid)
 	);
 `
 
 var selectJoinedUsersSetForRoomsSQL = "" +
 	"SELECT target_nid, COUNT(room_nid) FROM roomserver_membership WHERE room_nid IN ($1) AND" +
-	" membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " GROUP BY target_nid"
+	" membership_nid = " + fmt.Sprintf("%d", tables.MembershipStateJoin) + " and forgotten = false" +
+	" GROUP BY target_nid"
 
 // Insert a row in to membership table so that it can be locked by the
 // SELECT FOR UPDATE
@@ -52,37 +54,41 @@ const insertMembershipSQL = "" +
 	" ON CONFLICT DO NOTHING"
 
 const selectMembershipFromRoomAndTargetSQL = "" +
-	"SELECT membership_nid, event_nid FROM roomserver_membership" +
+	"SELECT membership_nid, event_nid, forgotten FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND target_nid = $2"
 
 const selectMembershipsFromRoomAndMembershipSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
-	" WHERE room_nid = $1 AND membership_nid = $2"
+	" WHERE room_nid = $1 AND membership_nid = $2 and forgotten = false"
 
 const selectLocalMembershipsFromRoomAndMembershipSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND membership_nid = $2" +
-	" AND target_local = true"
+	" AND target_local = true and forgotten = false"
 
 const selectMembershipsFromRoomSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
-	" WHERE room_nid = $1"
+	" WHERE room_nid = $1 and forgotten = false"
 
 const selectLocalMembershipsFromRoomSQL = "" +
 	"SELECT event_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1" +
-	" AND target_local = true"
+	" AND target_local = true and forgotten = false"
 
 const selectMembershipForUpdateSQL = "" +
 	"SELECT membership_nid FROM roomserver_membership" +
 	" WHERE room_nid = $1 AND target_nid = $2"
 
 const updateMembershipSQL = "" +
-	"UPDATE roomserver_membership SET sender_nid = $1, membership_nid = $2, event_nid = $3" +
+	"UPDATE roomserver_membership SET sender_nid = $1, membership_nid = $2, event_nid = $3, forgotten = $4" +
-	" WHERE room_nid = $4 AND target_nid = $5"
+	" WHERE room_nid = $5 AND target_nid = $6"
+
+const updateMembershipForgetRoom = "" +
+	"UPDATE roomserver_membership SET forgotten = $1" +
+	" WHERE room_nid = $2 AND target_nid = $3"
 
 const selectRoomsWithMembershipSQL = "" +
-	"SELECT room_nid FROM roomserver_membership WHERE membership_nid = $1 AND target_nid = $2"
+	"SELECT room_nid FROM roomserver_membership WHERE membership_nid = $1 AND target_nid = $2 and forgotten = false"
 
 // selectKnownUsersSQL uses a sub-select statement here to find rooms that the user is
 // joined to. Since this information is used to populate the user directory, we will
@@ -106,6 +112,7 @@ type membershipStatements struct {
 	selectRoomsWithMembershipStmt                   *sql.Stmt
 	updateMembershipStmt                            *sql.Stmt
 	selectKnownUsersStmt                            *sql.Stmt
+	updateMembershipForgetRoomStmt                  *sql.Stmt
 }
 
 func NewSqliteMembershipTable(db *sql.DB) (tables.Membership, error) {
@@ -128,6 +135,7 @@ func NewSqliteMembershipTable(db *sql.DB) (tables.Membership, error) {
 		{&s.updateMembershipStmt, updateMembershipSQL},
 		{&s.selectRoomsWithMembershipStmt, selectRoomsWithMembershipSQL},
 		{&s.selectKnownUsersStmt, selectKnownUsersSQL},
+		{&s.updateMembershipForgetRoomStmt, updateMembershipForgetRoom},
 	}.Prepare(db)
 }
 
@@ -155,10 +163,10 @@ func (s *membershipStatements) SelectMembershipForUpdate(
 func (s *membershipStatements) SelectMembershipFromRoomAndTarget(
 	ctx context.Context,
 	roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
-) (eventNID types.EventNID, membership tables.MembershipState, err error) {
+) (eventNID types.EventNID, membership tables.MembershipState, forgotten bool, err error) {
 	err = s.selectMembershipFromRoomAndTargetStmt.QueryRowContext(
 		ctx, roomNID, targetUserNID,
-	).Scan(&membership, &eventNID)
+	).Scan(&membership, &eventNID, &forgotten)
 	return
 }
 
@@ -216,13 +224,12 @@ func (s *membershipStatements) SelectMembershipsFromRoomAndMembership(
 
 func (s *membershipStatements) UpdateMembership(
 	ctx context.Context, txn *sql.Tx,
-	roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
+	roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, senderUserNID types.EventStateKeyNID, membership tables.MembershipState,
-	senderUserNID types.EventStateKeyNID, membership tables.MembershipState,
+	eventNID types.EventNID, forgotten bool,
-	eventNID types.EventNID,
 ) error {
 	stmt := sqlutil.TxStmt(txn, s.updateMembershipStmt)
 	_, err := stmt.ExecContext(
-		ctx, senderUserNID, membership, eventNID, roomNID, targetUserNID,
+		ctx, senderUserNID, membership, eventNID, forgotten, roomNID, targetUserNID,
 	)
 	return err
 }
@@ -285,3 +292,14 @@ func (s *membershipStatements) SelectKnownUsers(ctx context.Context, userID type
 	}
 	return result, rows.Err()
 }
+
+func (s *membershipStatements) UpdateForgetMembership(
+	ctx context.Context,
+	txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID,
+	forget bool,
+) error {
+	_, err := sqlutil.TxStmt(txn, s.updateMembershipForgetRoomStmt).ExecContext(
+		ctx, forget, roomNID, targetUserNID,
+	)
+	return err
+}

roomserver/storage/tables/interface.go

@@ -123,15 +123,16 @@ const (
 type Membership interface {
 	InsertMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, localTarget bool) error
 	SelectMembershipForUpdate(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID) (MembershipState, error)
-	SelectMembershipFromRoomAndTarget(ctx context.Context, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID) (types.EventNID, MembershipState, error)
+	SelectMembershipFromRoomAndTarget(ctx context.Context, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID) (types.EventNID, MembershipState, bool, error)
 	SelectMembershipsFromRoom(ctx context.Context, roomNID types.RoomNID, localOnly bool) (eventNIDs []types.EventNID, err error)
 	SelectMembershipsFromRoomAndMembership(ctx context.Context, roomNID types.RoomNID, membership MembershipState, localOnly bool) (eventNIDs []types.EventNID, err error)
-	UpdateMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, senderUserNID types.EventStateKeyNID, membership MembershipState, eventNID types.EventNID) error
+	UpdateMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, senderUserNID types.EventStateKeyNID, membership MembershipState, eventNID types.EventNID, forgotten bool) error
 	SelectRoomsWithMembership(ctx context.Context, userID types.EventStateKeyNID, membershipState MembershipState) ([]types.RoomNID, error)
 	// SelectJoinedUsersSetForRooms returns the set of all users in the rooms who are joined to any of these rooms, along with the
 	// counts of how many rooms they are joined.
 	SelectJoinedUsersSetForRooms(ctx context.Context, roomNIDs []types.RoomNID) (map[types.EventStateKeyNID]int, error)
 	SelectKnownUsers(ctx context.Context, userID types.EventStateKeyNID, searchString string, limit int) ([]string, error)
+	UpdateForgetMembership(ctx context.Context, txn *sql.Tx, roomNID types.RoomNID, targetUserNID types.EventStateKeyNID, forget bool) error
 }
 
 type Published interface {