From 69bcf47dca2e1bd7f5476b91c96ebf8528d4d3d4 Mon Sep 17 00:00:00 2001
From: Elle Mouton <elle.mouton@gmail.com>
Date: Tue, 5 Aug 2025 14:07:16 +0200
Subject: [PATCH] graph/db: use batch loading for PruneGraph

---
 graph/db/sql_store.go   | 51 +++++++++++++++++++++++------------------
 sqldb/sqlc/db_custom.go | 21 +++++++++++++++++
 2 files changed, 50 insertions(+), 22 deletions(-)

diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go
index cd960b166..f1b34b60e 100644
--- a/graph/db/sql_store.go
+++ b/graph/db/sql_store.go
@@ -2325,31 +2325,12 @@ func (s *SQLStore) PruneGraph(spentOutputs []*wire.OutPoint,
 		prunedNodes []route.Vertex
 	)
 	err := s.db.ExecTx(ctx, sqldb.WriteTxOpt(), func(db SQLQueries) error {
-		var chansToDelete []int64
-
-		// Define the callback function for processing each channel.
+		// First, collect all channel rows that need to be pruned.
+		var channelRows []sqlc.GetChannelsByOutpointsRow
 		channelCallback := func(ctx context.Context,
 			row sqlc.GetChannelsByOutpointsRow) error {
 
-			node1, node2, err := buildNodeVertices(
-				row.Node1Pubkey, row.Node2Pubkey,
-			)
-			if err != nil {
-				return err
-			}
-
-			info, err := getAndBuildEdgeInfo(
-				ctx, db, s.cfg.ChainHash, row.GraphChannel,
-				node1, node2,
-			)
-			if err != nil {
-				return err
-			}
-
-			closedChans = append(closedChans, info)
-			chansToDelete = append(
-				chansToDelete, row.GraphChannel.ID,
-			)
+			channelRows = append(channelRows, row)
 
 			return nil
 		}
@@ -2362,6 +2343,32 @@ func (s *SQLStore) PruneGraph(spentOutputs []*wire.OutPoint,
 				"outpoints: %w", err)
 		}
 
+		if len(channelRows) == 0 {
+			// There are no channels to prune. So we can exit early
+			// after updating the prune log.
+			err = db.UpsertPruneLogEntry(
+				ctx, sqlc.UpsertPruneLogEntryParams{
+					BlockHash:   blockHash[:],
+					BlockHeight: int64(blockHeight),
+				},
+			)
+			if err != nil {
+				return fmt.Errorf("unable to insert prune log "+
+					"entry: %w", err)
+			}
+
+			return nil
+		}
+
+		// Batch build all channel edges for pruning.
+		var chansToDelete []int64
+		closedChans, chansToDelete, err = batchBuildChannelInfo(
+			ctx, s.cfg, db, channelRows,
+		)
+		if err != nil {
+			return err
+		}
+
 		err = s.deleteChannels(ctx, db, chansToDelete)
 		if err != nil {
 			return fmt.Errorf("unable to delete channels: %w", err)
diff --git a/sqldb/sqlc/db_custom.go b/sqldb/sqlc/db_custom.go
index 823002817..2b378d4a8 100644
--- a/sqldb/sqlc/db_custom.go
+++ b/sqldb/sqlc/db_custom.go
@@ -105,3 +105,24 @@ func (r GetChannelsBySCIDWithPoliciesRow) Node1Pub() []byte {
 func (r GetChannelsBySCIDWithPoliciesRow) Node2Pub() []byte {
 	return r.GraphNode_2.PubKey
 }
+
+// Channel returns the GraphChannel associated with this interface.
+//
+// NOTE: This method is part of the ChannelAndNodeIDs interface.
+func (r GetChannelsByOutpointsRow) Channel() GraphChannel {
+	return r.GraphChannel
+}
+
+// Node1Pub returns the public key of the first node as a byte slice.
+//
+// NOTE: This method is part of the ChannelAndNodeIDs interface.
+func (r GetChannelsByOutpointsRow) Node1Pub() []byte {
+	return r.Node1Pubkey
+}
+
+// Node2Pub returns the public key of the second node as a byte slice.
+//
+// NOTE: This method is part of the ChannelAndNodeIDs interface.
+func (r GetChannelsByOutpointsRow) Node2Pub() []byte {
+	return r.Node2Pubkey
+}