Additional testing

state/event_table.go

@@ -320,8 +320,14 @@ type beforeSnapshotUpdate struct {
 //     which update "wins", see e.g.
 //     https://www.postgresql.org/docs/14/sql-update.html#id-1.9.3.182.8).
 func (t *EventTable) UpdateBeforeSnapshotIDs(txn *sqlx.Tx, updates []beforeSnapshotUpdate) error {
-	chunks := sqlutil.Chunkify2[beforeSnapshotUpdate](3, MaxPostgresParameters, updates)
+	return t.updateBeforeSnapshotIDs(txn, updates, MaxPostgresParameters)
+}
 
+// updateBeforeSnapshotIDs is the core of UpdateBeforeSnapshotIDs. It's pulled out so
+// we can test the updates work correctly when there is more than one chunk, without
+// having to generate an absurdly large number of updates.
+func (t *EventTable) updateBeforeSnapshotIDs(txn *sqlx.Tx, updates []beforeSnapshotUpdate, maxParams int) error {
+	chunks := sqlutil.Chunkify2[beforeSnapshotUpdate](3, maxParams, updates)
 	for _, chunk := range chunks {
 		// There's quite a lot of clunkiness in the query below.
 		//
@@ -351,7 +357,7 @@ func (t *EventTable) UpdateBeforeSnapshotIDs(txn *sqlx.Tx, updates []beforeSnaps
 		) AS u(before_state_snapshot_id, event_replaces_nid, event_nid)
 		WHERE e.event_nid = u.event_nid`, chunk)
 		if err != nil {
-			return err
+			return fmt.Errorf("failed to UpdateBeforeSnapshotIDs: %w", err)
 		}
 	}
 	return nil

state/event_table_test.go

@@ -1080,45 +1080,45 @@ func TestEventTableUpdateBeforeSnapshotIDs(t *testing.T) {
 	if err != nil {
 		t.Fatalf("failed to start txn: %s", err)
 	}
-	defer txn.Commit()
+	defer txn.Rollback()
 	const roomID = "!1:localhost"
 
 	// Note: there shouldn't be any other events with these IDs inserted before this
 	// transaction. $A and $B seem to be inserted and commit in TestEventTablePrevBatch.
-	const eventID1 = "$A-UpdateBeforeSnapshotIDs"
-	const eventID2 = "$B-UpdateBeforeSnapshotIDs"
-	const eventID3 = "$C-UpdateBeforeSnapshotIDs"
-	const eventID4 = "$D-UpdateBeforeSnapshotIDs"
+	const eventID0 = "$A-UpdateBeforeSnapshotIDs"
+	const eventID1 = "$B-UpdateBeforeSnapshotIDs"
+	const eventID2 = "$C-UpdateBeforeSnapshotIDs"
+	const eventID3 = "$D-UpdateBeforeSnapshotIDs"
 
 	knownEvents := []Event{
 		{
 			Type:     "m.room.create",
 			StateKey: "",
 			IsState:  true,
-			ID:       eventID1,
+			ID:       eventID0,
 			RoomID:   roomID,
 		},
 		{
 			Type:     "m.room.name",
 			StateKey: "",
 			IsState:  true,
-			ID:       eventID2,
+			ID:       eventID1,
 			RoomID:   roomID,
 		},
 		// Update the room name
 		{
 			Type:     "m.room.name",
 			StateKey: "",
 			IsState:  true,
-			ID:       eventID3,
+			ID:       eventID2,
 			RoomID:   roomID,
 		},
 		// Set topic
 		{
 			Type:     "m.room.topic",
 			StateKey: "",
 			IsState:  true,
-			ID:       eventID4,
+			ID:       eventID3,
 			RoomID:   roomID,
 		},
 	}
@@ -1130,39 +1130,162 @@ func TestEventTableUpdateBeforeSnapshotIDs(t *testing.T) {
 		t.Fatalf("failed to insert event: %s", err)
 	}
 
-	updates := []beforeSnapshotUpdate{
-		// m.room.create
-		{
-			NID:                   nidMap[eventID1],
-			BeforeStateSnapshotID: 0,
-			ReplacesNID:           0,
-		},
-		// m.room.name
-		{
-			NID:                   nidMap[eventID2],
-			BeforeStateSnapshotID: 1,
-			ReplacesNID:           0,
-		},
-		// m.room.name update
-		{
-			NID:                   nidMap[eventID3],
-			BeforeStateSnapshotID: 2,
-			ReplacesNID:           nidMap[eventID2],
-		},
-		// topic set
-		{
-			NID:                   nidMap[eventID4],
-			BeforeStateSnapshotID: 3,
-			ReplacesNID:           0,
-		},
+	var expectations []beforeSnapshotUpdate
+	checkUpdatesApplied := func() {
+		// Reselect the events and check the updates have applied
+		updatedEvents, err := table.SelectByNIDs(txn, true, []int64{nidMap[eventID0], nidMap[eventID1], nidMap[eventID2], nidMap[eventID3]})
+		if err != nil {
+			t.Fatal(err)
+		}
+		for i, expectation := range expectations {
+			event := updatedEvents[i]
+			if event.NID != expectation.NID {
+				t.Errorf("NID mismatch (%s): got %d, expected %d", event.ID, event.NID, expectation.NID)
+			}
+			if event.BeforeStateSnapshotID != expectation.BeforeStateSnapshotID {
+				t.Errorf("BeforeStateSnapshotID mismatch (%s): got %d, expected %d", event.ID, event.BeforeStateSnapshotID, expectation.BeforeStateSnapshotID)
+			}
+			if event.ReplacesNID != expectation.ReplacesNID {
+				t.Errorf("ReplacesNID mismatch (%s): got %d, expected %d", event.ID, event.ReplacesNID, expectation.ReplacesNID)
+			}
+		}
+	}
+
+	t.Run("Can update all events", func(t *testing.T) {
+		// Give each event a state snapshot and a replaces NID.
+		updates := []beforeSnapshotUpdate{
+			// m.room.create
+			{nidMap[eventID0], 0, 0},
+			// m.room.name
+			{nidMap[eventID1], 1, 0},
+			// m.room.name update
+			{nidMap[eventID2], 2, nidMap[eventID1]},
+			// topic set
+			{nidMap[eventID3], 3, 0},
+		}
+		err = table.UpdateBeforeSnapshotIDs(txn, updates)
+		if err != nil {
+			t.Fatal(err)
+		}
+		expectations = updates
+		checkUpdatesApplied()
+	})
+
+	t.Run("Can update a single event", func(t *testing.T) {
+		updates := []beforeSnapshotUpdate{
+			// Prove that we can update a single event without altering the others.
+			// We'd never do this in the application proper; we're only doing this to
+			// demonstrate the correctness of the batch SQL update.
+			{nidMap[eventID3], 5, 123},
+		}
+		err = table.UpdateBeforeSnapshotIDs(txn, updates)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// This update should apply to event 3; the others remain unaffected.
+		expectations[3] = updates[0]
+		checkUpdatesApplied()
+	})
+
+	t.Run("Can update some events", func(t *testing.T) {
+		updates := []beforeSnapshotUpdate{
+			{nidMap[eventID2], 6, 456},
+			{nidMap[eventID3], 7, 789},
+		}
+		err = table.UpdateBeforeSnapshotIDs(txn, updates)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// This update should apply to event 3; the others remain unaffected.
+		expectations[2] = updates[0]
+		expectations[3] = updates[1]
+		checkUpdatesApplied()
+	})
+
+	t.Run("Can update some events where one is unknown", func(t *testing.T) {
+		// The NIDs are generated by the DB using a postgres sequence starting from 1.
+		// So we should never have a negative NID.
+		const unknownEventNID = -999999
+		updates := []beforeSnapshotUpdate{
+			{nidMap[eventID1], 8, 1234},
+			{unknownEventNID, 9, 1234},
+		}
+		err = table.UpdateBeforeSnapshotIDs(txn, updates)
+		if err != nil {
+			t.Fatal(err)
+		}
+		// This update should apply to event 3; the others remain unaffected.
+		expectations[1] = updates[0]
+		checkUpdatesApplied()
+
+		// Additionally check that there is no event with the unknown NID.
+		selected, err := table.SelectByNIDs(txn, false, []int64{unknownEventNID})
+		if err != nil {
+			t.Fatal(err)
+		}
+		if len(selected) != 0 {
+			t.Fatalf("Expected there to be no event with NID %d, but got %d such events", unknownEventNID, len(selected))
+		}
+	})
+
+}
+
+func TestEventTableUpdateBeforeSnapshotIDsMultipleChunks(t *testing.T) {
+	db, close := connectToDB(t)
+	defer close()
+	txn, err := db.Beginx()
+	if err != nil {
+		t.Fatalf("failed to start txn: %s", err)
+	}
+	defer txn.Rollback()
+
+	// We're going to try and update 50 rows. Each update takes 3 params for a total of
+	// 150 params. We'll set the maxParams per chunk to 40, meaning that we should use
+	// at least 4 chunks.
+	const numUpdates = 50
+	const maxParams = 40
+
+	eventIDs := make([]string, 0, numUpdates)
+	events := make([]Event, 0, numUpdates)
+	for i := 0; i < numUpdates; i++ {
+		ev := Event{
+			Type:     "com.example.foo",
+			StateKey: "",
+			IsState:  true,
+			ID:       fmt.Sprintf("$%d-UpdateBeforeSnapshotID", i),
+			RoomID:   "!room2:unimportant",
+		}
+		events = append(events, ev)
+		eventIDs = append(eventIDs, ev.ID)
+	}
+
+	// Insert the events
+	table := NewEventTable(db)
+	nidMap, err := table.Insert(txn, events, false)
+	if err != nil {
+		t.Fatalf("failed to insert event: %s", err)
+	}
+	nids := make([]int64, 0, numUpdates)
+	for _, eventID := range eventIDs {
+		nids = append(nids, nidMap[eventID])
+	}
+
+	// Update 'em all
+	updates := make([]beforeSnapshotUpdate, 0, numUpdates)
+	for i := 0; i < numUpdates; i++ {
+		updates = append(updates, beforeSnapshotUpdate{
+			NID:                   nidMap[eventIDs[i]],
+			BeforeStateSnapshotID: int64(i),
+			ReplacesNID:           int64(numUpdates + i),
+		})
 	}
-	err = table.UpdateBeforeSnapshotIDs(txn, updates)
+	err = table.updateBeforeSnapshotIDs(txn, updates, maxParams)
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	// Reselect the events and check the updates have applied
-	updatedEvents, err := table.SelectByNIDs(txn, true, []int64{nidMap[eventID1], nidMap[eventID2], nidMap[eventID3], nidMap[eventID4]})
+	// Reselect and check that the update applied correctly.
+	updatedEvents, err := table.SelectByNIDs(txn, true, nids)
 	if err != nil {
 		t.Fatal(err)
 	}