1 files changed, 40 insertions, 8 deletions
diff --git a/syncapi/storage/sqlite3/current_room_state_table.go b/syncapi/storage/sqlite3/current_room_state_table.go
index c91ca692..587f9d24 100644
--- a/syncapi/storage/sqlite3/current_room_state_table.go
+++ b/syncapi/storage/sqlite3/current_room_state_table.go
@@ -66,6 +66,9 @@ const DeleteRoomStateForRoomSQL = "" +
 const selectRoomIDsWithMembershipSQL = "" +
 	"SELECT DISTINCT room_id FROM syncapi_current_room_state WHERE type = 'm.room.member' AND state_key = $1 AND membership = $2"
 
+const selectRoomIDsWithAnyMembershipSQL = "" +
+	"SELECT DISTINCT room_id, membership FROM syncapi_current_room_state WHERE type = 'm.room.member' AND state_key = $1"
+
 const selectCurrentStateSQL = "" +
 	"SELECT event_id, headered_event_json FROM syncapi_current_room_state WHERE room_id = $1"
 
@@ -86,14 +89,15 @@ const selectEventsWithEventIDsSQL = "" +
 	" FROM syncapi_current_room_state WHERE event_id IN ($1)"
 
 type currentRoomStateStatements struct {
-	db                              *sql.DB
-	streamIDStatements              *streamIDStatements
-	upsertRoomStateStmt             *sql.Stmt
-	deleteRoomStateByEventIDStmt    *sql.Stmt
-	DeleteRoomStateForRoomStmt      *sql.Stmt
-	selectRoomIDsWithMembershipStmt *sql.Stmt
-	selectJoinedUsersStmt           *sql.Stmt
-	selectStateEventStmt            *sql.Stmt
+	db                                 *sql.DB
+	streamIDStatements                 *streamIDStatements
+	upsertRoomStateStmt                *sql.Stmt
+	deleteRoomStateByEventIDStmt       *sql.Stmt
+	DeleteRoomStateForRoomStmt         *sql.Stmt
+	selectRoomIDsWithMembershipStmt    *sql.Stmt
+	selectRoomIDsWithAnyMembershipStmt *sql.Stmt
+	selectJoinedUsersStmt              *sql.Stmt
+	selectStateEventStmt               *sql.Stmt
 }
 
 func NewSqliteCurrentRoomStateTable(db *sql.DB, streamID *streamIDStatements) (tables.CurrentRoomState, error) {
@@ -117,6 +121,9 @@ func NewSqliteCurrentRoomStateTable(db *sql.DB, streamID *streamIDStatements) (t
 	if s.selectRoomIDsWithMembershipStmt, err = db.Prepare(selectRoomIDsWithMembershipSQL); err != nil {
 		return nil, err
 	}
+	if s.selectRoomIDsWithAnyMembershipStmt, err = db.Prepare(selectRoomIDsWithAnyMembershipSQL); err != nil {
+		return nil, err
+	}
 	if s.selectJoinedUsersStmt, err = db.Prepare(selectJoinedUsersSQL); err != nil {
 		return nil, err
 	}
@@ -175,6 +182,31 @@ func (s *currentRoomStateStatements) SelectRoomIDsWithMembership(
 	return result, nil
 }
 
+// SelectRoomIDsWithAnyMembership returns a map of all memberships for the given user.
+func (s *currentRoomStateStatements) SelectRoomIDsWithAnyMembership(
+	ctx context.Context,
+	txn *sql.Tx,
+	userID string,
+) (map[string]string, error) {
+	stmt := sqlutil.TxStmt(txn, s.selectRoomIDsWithAnyMembershipStmt)
+	rows, err := stmt.QueryContext(ctx, userID)
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomIDsWithAnyMembership: rows.close() failed")
+
+	result := map[string]string{}
+	for rows.Next() {
+		var roomID string
+		var membership string
+		if err := rows.Scan(&roomID, &membership); err != nil {
+			return nil, err
+		}
+		result[roomID] = membership
+	}
+	return result, rows.Err()
+}
+
 // CurrentState returns all the current state events for the given room.
 func (s *currentRoomStateStatements) SelectCurrentState(
 	ctx context.Context, txn *sql.Tx, roomID string,