Add QueryStateAndAuthChainIDs

roomserver/storage/interface.go

@@ -87,6 +87,8 @@ type Database interface {
 	// Lookup the event IDs for a batch of event numeric IDs.
 	// Returns an error if the retrieval went wrong.
 	EventIDs(ctx context.Context, eventNIDs []types.EventNID) (map[types.EventNID]string, error)
+	// AuthEventNIDs returns the auth event NIDs for the given events.
+	AuthEventNIDs(ctx context.Context, events []types.EventNID) (types.EventNIDs, error)
 	// Look up the latest events in a room in preparation for an update.
 	// The RoomRecentEventsUpdater must have Commit or Rollback called on it if this doesn't return an error.
 	// Returns the latest events in the room and the last eventID sent to the log along with an updater.

roomserver/storage/postgres/events_table.go

@@ -134,6 +134,9 @@ const selectMaxEventDepthSQL = "" +
 const selectRoomNIDsForEventNIDsSQL = "" +
 	"SELECT event_nid, room_nid FROM roomserver_events WHERE event_nid = ANY($1)"
 
+const bulkSelectEventAuthEventNIDsSQL = "" +
+	"SELECT auth_event_nids FROM roomserver_events WHERE event_nid = ANY($1)"
+
 type eventStatements struct {
 	insertEventStmt                        *sql.Stmt
 	selectEventStmt                        *sql.Stmt
@@ -150,6 +153,7 @@ type eventStatements struct {
 	bulkSelectEventNIDStmt                 *sql.Stmt
 	selectMaxEventDepthStmt                *sql.Stmt
 	selectRoomNIDsForEventNIDsStmt         *sql.Stmt
+	bulkSelectEventAuthEventNIDsStmt       *sql.Stmt
 }
 
 func createEventsTable(db *sql.DB) error {
@@ -176,6 +180,7 @@ func prepareEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.bulkSelectEventNIDStmt, bulkSelectEventNIDSQL},
 		{&s.selectMaxEventDepthStmt, selectMaxEventDepthSQL},
 		{&s.selectRoomNIDsForEventNIDsStmt, selectRoomNIDsForEventNIDsSQL},
+		{&s.bulkSelectEventAuthEventNIDsStmt, bulkSelectEventAuthEventNIDsSQL},
 	}.Prepare(db)
 }
 
@@ -502,6 +507,28 @@ func (s *eventStatements) SelectRoomNIDsForEventNIDs(
 	return result, nil
 }
 
+func (s *eventStatements) SelectEventAuthEventNIDs(
+	ctx context.Context, eventNIDs []types.EventNID,
+) (map[types.EventNID][]types.EventNID, error) {
+	rows, err := s.bulkSelectEventAuthEventNIDsStmt.QueryContext(ctx, eventNIDsAsArray(eventNIDs))
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "selectRoomNIDsForEventNIDsStmt: rows.close() failed")
+	result := make(map[types.EventNID][]types.EventNID)
+	for rows.Next() {
+		var eventNID types.EventNID
+		var authEventNIDs pq.Int64Array
+		if err = rows.Scan(&authEventNIDs); err != nil {
+			return nil, err
+		}
+		for _, a := range authEventNIDs {
+			result[eventNID] = append(result[eventNID], types.EventNID(a))
+		}
+	}
+	return result, nil
+}
+
 func eventNIDsAsArray(eventNIDs []types.EventNID) pq.Int64Array {
 	nids := make([]int64, len(eventNIDs))
 	for i := range eventNIDs {

roomserver/storage/shared/storage.go

@@ -292,6 +292,22 @@ func (d *Database) StateEntries(
 	return lists, nil
 }
 
+func (d *Database) AuthEventNIDs(
+	ctx context.Context, events []types.EventNID,
+) (types.EventNIDs, error) {
+	entries, err := d.EventsTable.SelectEventAuthEventNIDs(
+		ctx, events,
+	)
+	if err != nil {
+		return nil, fmt.Errorf("d.EventsTable.SelectEventAuthEventNIDs: %w", err)
+	}
+	var lists types.EventNIDs
+	for _, nids := range entries {
+		lists = append(lists, nids...)
+	}
+	return lists[:util.SortAndUnique(lists)], nil
+}
+
 func (d *Database) SetRoomAlias(ctx context.Context, alias string, roomID string, creatorUserID string) error {
 	return d.Writer.Do(d.DB, nil, func(txn *sql.Tx) error {
 		return d.RoomAliasesTable.InsertRoomAlias(ctx, txn, alias, roomID, creatorUserID)

roomserver/storage/sqlite3/events_table.go

@@ -23,6 +23,7 @@ import (
 	"sort"
 	"strings"
 
+	"github.com/lib/pq"
 	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/storage/shared"
@@ -104,6 +105,9 @@ const selectMaxEventDepthSQL = "" +
 const selectRoomNIDsForEventNIDsSQL = "" +
 	"SELECT event_nid, room_nid FROM roomserver_events WHERE event_nid IN ($1)"
 
+const bulkSelectEventAuthEventNIDsSQL = "" +
+	"SELECT auth_event_nids FROM roomserver_events WHERE event_nid IN ($1)"
+
 type eventStatements struct {
 	db                                     *sql.DB
 	insertEventStmt                        *sql.Stmt
@@ -119,6 +123,7 @@ type eventStatements struct {
 	bulkSelectEventIDStmt                  *sql.Stmt
 	bulkSelectEventNIDStmt                 *sql.Stmt
 	//selectRoomNIDsForEventNIDsStmt           *sql.Stmt
+	//bulkSelectEventAuthEventNIDsStmt         *sql.Stmt
 }
 
 func createEventsTable(db *sql.DB) error {
@@ -145,6 +150,7 @@ func prepareEventsTable(db *sql.DB) (tables.Events, error) {
 		{&s.bulkSelectEventIDStmt, bulkSelectEventIDSQL},
 		{&s.bulkSelectEventNIDStmt, bulkSelectEventNIDSQL},
 		//{&s.selectRoomNIDForEventNIDStmt, selectRoomNIDForEventNIDSQL},
+		//{&s.bulkSelectEventAuthEventNIDsStmt, bulkSelectEventAuthEventNIDsSQL},
 	}.Prepare(db)
 }
 
@@ -571,6 +577,37 @@ func (s *eventStatements) SelectRoomNIDsForEventNIDs(
 	return result, nil
 }
 
+func (s *eventStatements) SelectEventAuthEventNIDs(
+	ctx context.Context, eventNIDs []types.EventNID,
+) (map[types.EventNID][]types.EventNID, error) {
+	sqlStr := strings.Replace(bulkSelectEventAuthEventNIDsSQL, "($1)", sqlutil.QueryVariadic(len(eventNIDs)), 1)
+	sqlPrep, err := s.db.Prepare(sqlStr)
+	if err != nil {
+		return nil, err
+	}
+	iEventNIDs := make([]interface{}, len(eventNIDs))
+	for i, v := range eventNIDs {
+		iEventNIDs[i] = v
+	}
+	rows, err := sqlPrep.QueryContext(ctx, iEventNIDs...)
+	if err != nil {
+		return nil, err
+	}
+	defer internal.CloseAndLogIfError(ctx, rows, "bulkSelectEventAuthEventNIDsStmt: rows.close() failed")
+	result := make(map[types.EventNID][]types.EventNID)
+	for rows.Next() {
+		var eventNID types.EventNID
+		var authEventNIDs pq.Int64Array
+		if err = rows.Scan(&authEventNIDs); err != nil {
+			return nil, err
+		}
+		for _, a := range authEventNIDs {
+			result[eventNID] = append(result[eventNID], types.EventNID(a))
+		}
+	}
+	return result, nil
+}
+
 func eventNIDsAsArray(eventNIDs []types.EventNID) string {
 	b, _ := json.Marshal(eventNIDs)
 	return string(b)

roomserver/storage/tables/interface.go

@@ -61,6 +61,7 @@ type Events interface {
 	BulkSelectEventNID(ctx context.Context, eventIDs []string) (map[string]types.EventNID, error)
 	SelectMaxEventDepth(ctx context.Context, txn *sql.Tx, eventNIDs []types.EventNID) (int64, error)
 	SelectRoomNIDsForEventNIDs(ctx context.Context, eventNIDs []types.EventNID) (roomNIDs map[types.EventNID]types.RoomNID, err error)
+	SelectEventAuthEventNIDs(ctx context.Context, eventNIDs []types.EventNID) (map[types.EventNID][]types.EventNID, error)
 }
 
 type Rooms interface {