diff --git a/graph/builder.go b/graph/builder.go
index a471e9525..a47cc18d1 100644
--- a/graph/builder.go
+++ b/graph/builder.go
@@ -366,7 +366,9 @@ func (b *Builder) syncGraphWithChain() error {
 			"(hash=%v)", pruneHeight, pruneHash)
 		// Prune the graph for every channel that was opened at height
 		// >= pruneHeight.
-		_, err := b.cfg.Graph.DisconnectBlockAtHeight(pruneHeight)
+		_, err := b.cfg.Graph.DisconnectBlockAtHeight(
+			context.TODO(), pruneHeight,
+		)
 		if err != nil {
 			return err
 		}
@@ -687,7 +689,7 @@ func (b *Builder) networkHandler() {
 			// Update the channel graph to reflect that this block
 			// was disconnected.
 			_, err := b.cfg.Graph.DisconnectBlockAtHeight(
-				blockHeight,
+				context.TODO(), blockHeight,
 			)
 			if err != nil {
 				log.Errorf("unable to prune graph with stale "+
diff --git a/graph/db/graph.go b/graph/db/graph.go
index fd32348e9..c42f54301 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -419,10 +419,10 @@ func (c *ChannelGraph) DeleteChannelEdges(ctx context.Context,
 // set to the last prune height valid for the remaining chain.
 // Channels that were removed from the graph resulting from the
 // disconnected block are returned.
-func (c *ChannelGraph) DisconnectBlockAtHeight(height uint32) (
-	[]*models.ChannelEdgeInfo, error) {
+func (c *ChannelGraph) DisconnectBlockAtHeight(ctx context.Context,
+	height uint32) ([]*models.ChannelEdgeInfo, error) {
 
-	edges, err := c.db.DisconnectBlockAtHeight(height)
+	edges, err := c.db.DisconnectBlockAtHeight(ctx, height)
 	if err != nil {
 		return nil, err
 	}
diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go
index 4ad0ae336..d9c26a909 100644
--- a/graph/db/graph_test.go
+++ b/graph/db/graph_test.go
@@ -905,7 +905,7 @@ func TestDisconnectBlockAtHeight(t *testing.T) {
 
 	// Call DisconnectBlockAtHeight, which should prune every channel
 	// that has a funding height of 'height' or greater.
-	removed, err := graph.DisconnectBlockAtHeight(uint32(height))
+	removed, err := graph.DisconnectBlockAtHeight(ctx, height)
 	require.NoError(t, err)
 	assertNoEdge(t, graph, edgeInfo.ChannelID)
 	assertNoEdge(t, graph, edgeInfo2.ChannelID)
@@ -3437,7 +3437,7 @@ func TestStressTestChannelGraphAPI(t *testing.T) {
 			name: "DisconnectBlockAtHeight",
 			fn: func() error {
 				_, err := graph.DisconnectBlockAtHeight(
-					newBlockHeight(),
+					ctx, newBlockHeight(),
 				)
 
 				return err
diff --git a/graph/db/interfaces.go b/graph/db/interfaces.go
index 27dcb042c..615b2cb0d 100644
--- a/graph/db/interfaces.go
+++ b/graph/db/interfaces.go
@@ -412,6 +412,6 @@ type Store interface { //nolint:interfacebloat
 	// set to the last prune height valid for the remaining chain.
 	// Channels that were removed from the graph resulting from the
 	// disconnected block are returned.
-	DisconnectBlockAtHeight(height uint32) ([]*models.ChannelEdgeInfo,
-		error)
+	DisconnectBlockAtHeight(ctx context.Context,
+		height uint32) ([]*models.ChannelEdgeInfo, error)
 }
diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go
index a5a324d28..7a57db673 100644
--- a/graph/db/kv_store.go
+++ b/graph/db/kv_store.go
@@ -1760,8 +1760,8 @@ func (c *KVStore) pruneGraphNodes(nodes kvdb.RwBucket,
 // set to the last prune height valid for the remaining chain.
 // Channels that were removed from the graph resulting from the
 // disconnected block are returned.
-func (c *KVStore) DisconnectBlockAtHeight(height uint32) (
-	[]*models.ChannelEdgeInfo, error) {
+func (c *KVStore) DisconnectBlockAtHeight(_ context.Context,
+	height uint32) ([]*models.ChannelEdgeInfo, error) {
 
 	// Every channel having a ShortChannelID starting at 'height'
 	// will no longer be confirmed.
diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go
index 406b16d39..f1aad755d 100644
--- a/graph/db/sql_store.go
+++ b/graph/db/sql_store.go
@@ -3112,10 +3112,8 @@ func (s *SQLStore) pruneGraphNodes(ctx context.Context,
 // disconnected block are returned.
 //
 // NOTE: part of the Store interface.
-func (s *SQLStore) DisconnectBlockAtHeight(height uint32) (
-	[]*models.ChannelEdgeInfo, error) {
-
-	ctx := context.TODO()
+func (s *SQLStore) DisconnectBlockAtHeight(ctx context.Context,
+	height uint32) ([]*models.ChannelEdgeInfo, error) {
 
 	var (
 		// Every channel having a ShortChannelID starting at 'height'