Try to refactor PostgreSQL migration

roomserver/storage/postgres/deltas/20210416150927_state_blocks_refactor.go

@@ -15,12 +15,10 @@
 package deltas
 
 import (
-	"context"
 	"database/sql"
 	"fmt"
 
 	"github.com/lib/pq"
-	"github.com/matrix-org/dendrite/internal"
 	"github.com/matrix-org/dendrite/internal/sqlutil"
 	"github.com/matrix-org/dendrite/roomserver/types"
 	"github.com/matrix-org/util"
@@ -63,38 +61,68 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 		return fmt.Errorf("tx.Exec: %w", err)
 	}
 	logrus.Warn("New tables created...")
-	snapshotrows, err := tx.Query(`SELECT state_snapshot_nid, room_nid, state_block_nids FROM _roomserver_state_snapshots;`)
+
+	var snapshotcount int
+	err = tx.QueryRow(`
+		SELECT COUNT(DISTINCT state_snapshot_nid) FROM roomserver_state_snapshots;
+	`).Scan(&snapshotcount)
+	if err != nil {
+		return fmt.Errorf("tx.QueryRow.Scan (count snapshots): %w", err)
+	}
+	logrus.Warnf("Will convert %d snapshots...", snapshotcount)
+
+	batchsize := 100
+	batchoffset := 0
+
+	var lastsnapshot types.StateSnapshotNID
+	var newblocks types.StateBlockNIDs
+	var snapshots *sql.Rows
+
+	for ; batchoffset < snapshotcount; batchoffset += batchsize {
+		snapshots, err = tx.Query(`
+			SELECT
+				state_snapshot_nid,
+				room_nid,
+				state_block_nid,
+				ARRAY_AGG(event_nid) AS event_nids
+			FROM (
+				SELECT
+					_roomserver_state_snapshots.state_snapshot_nid,
+					_roomserver_state_snapshots.room_nid,
+					_roomserver_state_block.state_block_nid,
+					_roomserver_state_block.event_nid
+				FROM
+					_roomserver_state_snapshots
+					JOIN _roomserver_state_block ON _roomserver_state_block.state_block_nid = ANY (_roomserver_state_snapshots.state_block_nids)
+				WHERE
+					_roomserver_state_snapshots.state_snapshot_nid = ANY ( SELECT DISTINCT
+							_roomserver_state_snapshots.state_snapshot_nid
+						FROM
+							_roomserver_state_snapshots
+						LIMIT $1 OFFSET $2)) AS _roomserver_state_block
+			GROUP BY
+				state_snapshot_nid,
+				room_nid,
+				state_block_nid;
+		`, batchoffset, batchsize)
 		if err != nil {
 			return fmt.Errorf("tx.Query: %w", err)
 		}
-	defer internal.CloseAndLogIfError(context.TODO(), snapshotrows, "rows.close() failed")
+
-	for snapshotrows.Next() {
+		for snapshots.Next() {
+			logrus.Warnf("Performing %d to %d...", batchoffset, batchoffset+batchsize)
+
 			var snapshot types.StateSnapshotNID
 			var room types.RoomNID
 			var blocksarray pq.Int64Array
-		var blocks []types.StateBlockNID
+			var eventsarray pq.Int64Array
-		if err = snapshotrows.Scan(&snapshot, &room, &blocksarray); err != nil {
+			if err = snapshots.Scan(&snapshot, &room, &blocksarray, &eventsarray); err != nil {
 				return fmt.Errorf("rows.Scan: %w", err)
 			}
-		for _, b := range blocksarray {
-			blocks = append(blocks, types.StateBlockNID(b))
-		}
 
-		var newblocks []types.StateBlockNID
+			var events types.EventNIDs
-		for _, block := range blocks {
+			for _, e := range eventsarray {
-			if err = func() error {
+				events = append(events, types.EventNID(e))
-				blockrows, berr := tx.Query(`SELECT event_nid FROM _roomserver_state_block WHERE state_block_nid = $1`, int64(block))
-				if berr != nil {
-					return fmt.Errorf("tx.Query (event nids from old block): %w", berr)
-				}
-				defer internal.CloseAndLogIfError(context.TODO(), blockrows, "rows.close() failed")
-				events := types.EventNIDs{}
-				for blockrows.Next() {
-					var event types.EventNID
-					if err = blockrows.Scan(&event); err != nil {
-						return fmt.Errorf("rows.Scan (event nids from old block): %w", err)
-					}
-					events = append(events, event)
 			}
 			events = events[:util.SortAndUnique(events)]
 
@@ -109,11 +137,8 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 				return fmt.Errorf("tx.QueryRow.Scan (insert new block): %w", err)
 			}
 			newblocks = append(newblocks, blocknid)
-				return nil
-			}(); err != nil {
-				return err
-			}
 
+			if snapshot != lastsnapshot {
 				var newsnapshot types.StateSnapshotNID
 				err = tx.QueryRow(`
 				INSERT INTO roomserver_state_snapshots (room_nid, state_block_nids)
@@ -135,19 +160,24 @@ func UpStateBlocksRefactor(tx *sql.Tx) error {
 				}
 
 				fmt.Println("Rewrote snapshot", snapshot, "to", newsnapshot)
+				newblocks = newblocks[:0]
+				lastsnapshot = snapshot
+			}
+		}
+
+		if err = snapshots.Close(); err != nil {
+			return fmt.Errorf("snapshots.Close: %w", err)
 		}
 	}
 
-	/*
 	if _, err = tx.Exec(`DROP TABLE _roomserver_state_snapshots;`); err != nil {
 		return fmt.Errorf("tx.Exec (delete old snapshot table): %w", err)
 	}
 	if _, err = tx.Exec(`DROP TABLE _roomserver_state_block;`); err != nil {
 		return fmt.Errorf("tx.Exec (delete old block table): %w", err)
 	}
-	*/
 
-	return nil
+	return fmt.Errorf("stopping here to revert changes")
 }
 
 func DownStateBlocksRefactor(tx *sql.Tx) error {