dropchannelgraph: allow removing single channel

cmd/chantools/dropchannelgraph.go

@@ -2,7 +2,6 @@ package main
 
 import (
 	"fmt"
-
 	"github.com/guggero/chantools/lnd"
 	"github.com/spf13/cobra"
 )
@@ -16,6 +15,8 @@ var (
 type dropChannelGraphCommand struct {
 	ChannelDB string
 
+	SingleChannel uint64
+
 	cmd *cobra.Command
 }
 
@@ -27,17 +28,29 @@ func newDropChannelGraphCommand() *cobra.Command {
 		Long: `This command removes all graph data from a channel DB,
 forcing the lnd node to do a full graph sync.
 
+Or if a single channel is specified, that channel is purged from the graph
+without removing any other data.
+
 CAUTION: Running this command will make it impossible to use the channel DB
 with an older version of lnd. Downgrading is not possible and you'll need to
 run lnd v0.13.1-beta or later after using this command!'`,
 		Example: `chantools dropchannelgraph \
-	--channeldb ~/.lnd/data/graph/mainnet/channel.db`,
+	--channeldb ~/.lnd/data/graph/mainnet/channel.db
+
+chantools dropchannelgraph \
+	--channeldb ~/.lnd/data/graph/mainnet/channel.db \
+	--single_channel 726607861215512345`,
 		RunE: cc.Execute,
 	}
 	cc.cmd.Flags().StringVar(
 		&cc.ChannelDB, "channeldb", "", "lnd channel.db file to dump "+
 			"channels from",
 	)
+	cc.cmd.Flags().Uint64Var(
+		&cc.SingleChannel, "single_channel", 0, "the single channel "+
+			"identified by its short channel ID (CID) to remove "+
+			"from the graph",
+	)
 
 	return cc.cmd
 }
@@ -53,20 +66,28 @@ func (c *dropChannelGraphCommand) Execute(_ *cobra.Command, _ []string) error {
 	}
 	defer func() { _ = db.Close() }()
 
-	rwTx, err := db.BeginReadWriteTx()
-	if err != nil {
-		return err
-	}
+	if c.SingleChannel != 0 {
+		log.Infof("Removing single channel %d", c.SingleChannel)
+		return db.ChannelGraph().DeleteChannelEdges(
+			true, c.SingleChannel,
+		)
+	} else {
+		log.Infof("Dropping all graph related buckets")
 
-	if err := rwTx.DeleteTopLevelBucket(nodeBucket); err != nil {
-		return err
-	}
-	if err := rwTx.DeleteTopLevelBucket(edgeBucket); err != nil {
-		return err
-	}
-	if err := rwTx.DeleteTopLevelBucket(graphMetaBucket); err != nil {
-		return err
-	}
+		rwTx, err := db.BeginReadWriteTx()
+		if err != nil {
+			return err
+		}
+		if err := rwTx.DeleteTopLevelBucket(nodeBucket); err != nil {
+			return err
+		}
+		if err := rwTx.DeleteTopLevelBucket(edgeBucket); err != nil {
+			return err
+		}
+		if err := rwTx.DeleteTopLevelBucket(graphMetaBucket); err != nil {
+			return err
+		}
 
-	return rwTx.Commit()
+		return rwTx.Commit()
+	}
 }