graph/db: version disabled channel IDs

diff --git a/graph/builder.go b/graph/builder.go
index 9d3baeb..69e7555 100644
--- a/graph/builder.go
+++ b/graph/builder.go
@@ -573,7 +573,9 @@ func (b *Builder) pruneZombieChans() error {
 	// both edges. If they're both disabled, then we can interpret this as
 	// the channel being closed and can prune it from our graph.
 	if b.cfg.AssumeChannelValid {
-		disabledChanIDs, err := b.cfg.Graph.DisabledChannelIDs()
+		disabledChanIDs, err := b.cfg.Graph.DisabledChannelIDs(
+			lnwire.GossipVersion1,
+		)
 		if err != nil {
 			return fmt.Errorf("unable to get disabled channels "+
 				"ids chans: %v", err)
diff --git a/graph/db/graph.go b/graph/db/graph.go
index 4c3ffd9..4579f16 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -660,8 +660,10 @@ func (c *ChannelGraph) ForEachChannel(ctx context.Context,
 }
 
 // DisabledChannelIDs returns the channel ids of disabled channels.
-func (c *ChannelGraph) DisabledChannelIDs() ([]uint64, error) {
-	return c.db.DisabledChannelIDs()
+func (c *ChannelGraph) DisabledChannelIDs(v lnwire.GossipVersion) (
+	[]uint64, error) {
+
+	return c.db.DisabledChannelIDs(v)
 }
 
 // HasV1ChannelEdge returns true if the database knows of a channel edge.
@@ -943,6 +945,11 @@ func (c *VersionedGraph) ForEachChannelCacheable(
 	return c.db.ForEachChannelCacheable(c.v, cb, reset)
 }
 
+// DisabledChannelIDs returns the channel ids of disabled channels.
+func (c *VersionedGraph) DisabledChannelIDs() ([]uint64, error) {
+	return c.db.DisabledChannelIDs(c.v)
+}
+
 // ChannelID attempts to lookup the 8-byte compact channel ID.
 func (c *VersionedGraph) ChannelID(chanPoint *wire.OutPoint) (uint64, error) {
 	return c.db.ChannelID(c.v, chanPoint)
diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go
index 47c8f51..431d818 100644
--- a/graph/db/graph_test.go
+++ b/graph/db/graph_test.go
@@ -186,6 +186,10 @@ var versionedTests = []versionedTest{
 		name: "batched update edge policy",
 		test: testBatchedUpdateEdgePolicy,
 	},
+	{
+		name: "disabled channel ids",
+		test: testDisabledChannelIDs,
+	},
 }
 
 // TestVersionedDBs runs various tests against both v1 and v2 versioned
@@ -4577,29 +4581,32 @@ func BenchmarkIsPublicNode(b *testing.B) {
 // TestDisabledChannelIDs ensures that the disabled channels within the
 // disabledEdgePolicyBucket are managed properly and the list returned from
 // DisabledChannelIDs is correct.
-func TestDisabledChannelIDs(t *testing.T) {
+func testDisabledChannelIDs(t *testing.T, v lnwire.GossipVersion) {
 	t.Parallel()
 	ctx := t.Context()
 
-	graph := MakeTestGraph(t)
+	graph := NewVersionedGraph(MakeTestGraph(t), v)
 
 	// Create first node and add it to the graph.
-	node1 := createTestVertex(t, lnwire.GossipVersion1)
+	node1 := createTestVertex(t, v)
 	if err := graph.AddNode(ctx, node1); err != nil {
 		t.Fatalf("unable to add node: %v", err)
 	}
 
 	// Create second node and add it to the graph.
-	node2 := createTestVertex(t, lnwire.GossipVersion1)
+	node2 := createTestVertex(t, v)
 	if err := graph.AddNode(ctx, node2); err != nil {
 		t.Fatalf("unable to add node: %v", err)
 	}
 
 	// Adding a new channel edge to the graph.
-	edgeInfo, edge1, edge2 := createChannelEdge(
-		node1, node2, lnwire.GossipVersion1,
-	)
-	node2.LastUpdate = nextUpdateTime()
+	edgeInfo, edge1, edge2 := createChannelEdge(node1, node2, v)
+	switch v {
+	case lnwire.GossipVersion1:
+		node2.LastUpdate = nextUpdateTime()
+	case lnwire.GossipVersion2:
+		node2.LastBlockHeight = nextBlockHeight()
+	}
 	if err := graph.AddNode(ctx, node2); err != nil {
 		t.Fatalf("unable to add node: %v", err)
 	}
@@ -4618,7 +4625,12 @@ func TestDisabledChannelIDs(t *testing.T) {
 
 	// Add one disabled policy and ensure the channel is still not in the
 	// disabled list.
-	edge1.ChannelFlags |= lnwire.ChanUpdateDisabled
+	switch v {
+	case lnwire.GossipVersion1:
+		edge1.ChannelFlags |= lnwire.ChanUpdateDisabled
+	case lnwire.GossipVersion2:
+		edge1.DisableFlags |= lnwire.ChanUpdateDisableIncoming
+	}
 	if err := graph.UpdateEdgePolicy(ctx, edge1); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
@@ -4631,7 +4643,12 @@ func TestDisabledChannelIDs(t *testing.T) {
 
 	// Add second disabled policy and ensure the channel is now in the
 	// disabled list.
-	edge2.ChannelFlags |= lnwire.ChanUpdateDisabled
+	switch v {
+	case lnwire.GossipVersion1:
+		edge2.ChannelFlags |= lnwire.ChanUpdateDisabled
+	case lnwire.GossipVersion2:
+		edge2.DisableFlags |= lnwire.ChanUpdateDisableIncoming
+	}
 	if err := graph.UpdateEdgePolicy(ctx, edge2); err != nil {
 		t.Fatalf("unable to update edge: %v", err)
 	}
diff --git a/graph/db/interfaces.go b/graph/db/interfaces.go
index 95d8bb3..9613ff6 100644
--- a/graph/db/interfaces.go
+++ b/graph/db/interfaces.go
@@ -190,7 +190,7 @@ type Store interface { //nolint:interfacebloat
 	// DisabledChannelIDs returns the channel ids of disabled channels.
 	// A channel is disabled when two of the associated ChanelEdgePolicies
 	// have their disabled bit on.
-	DisabledChannelIDs() ([]uint64, error)
+	DisabledChannelIDs(v lnwire.GossipVersion) ([]uint64, error)
 
 	// AddChannelEdge adds a new (undirected, blank) edge to the graph
 	// database. An undirected edge from the two target nodes are created.
diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go
index bdb1a78..d61713a 100644
--- a/graph/db/kv_store.go
+++ b/graph/db/kv_store.go
@@ -772,7 +772,13 @@ func (c *KVStore) ForEachNodeCached(ctx context.Context, withAddrs bool,
 // DisabledChannelIDs returns the channel ids of disabled channels.
 // A channel is disabled when two of the associated ChanelEdgePolicies
 // have their disabled bit on.
-func (c *KVStore) DisabledChannelIDs() ([]uint64, error) {
+func (c *KVStore) DisabledChannelIDs(
+	v lnwire.GossipVersion) ([]uint64, error) {
+
+	if v != lnwire.GossipVersion1 {
+		return nil, ErrVersionNotSupportedForKVDB
+	}
+
 	var disabledChanIDs []uint64
 	var chanEdgeFound map[uint64]struct{}
 
diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go
index 65013bc..4c64e2e 100644
--- a/graph/db/sql_store.go
+++ b/graph/db/sql_store.go
@@ -76,6 +76,8 @@ type SQLQueries interface {
 	GetNodeFeaturesBatch(ctx context.Context, ids []int64) ([]sqlc.GraphNodeFeature, error)
 	GetNodeFeaturesByPubKey(ctx context.Context, arg sqlc.GetNodeFeaturesByPubKeyParams) ([]int32, error)
 	DeleteNodeFeature(ctx context.Context, arg sqlc.DeleteNodeFeatureParams) error
+	GetV1DisabledSCIDs(ctx context.Context) ([][]byte, error)
+	GetV2DisabledSCIDs(ctx context.Context) ([][]byte, error)
 
 	/*
 		Source node queries.
@@ -119,7 +121,6 @@ type SQLQueries interface {
 	*/
 	UpsertEdgePolicy(ctx context.Context, arg sqlc.UpsertEdgePolicyParams) (int64, error)
 	GetChannelPolicyByChannelAndNode(ctx context.Context, arg sqlc.GetChannelPolicyByChannelAndNodeParams) (sqlc.GraphChannelPolicy, error)
-	GetV1DisabledSCIDs(ctx context.Context) ([][]byte, error)
 
 	UpsertChanPolicyExtraType(ctx context.Context, arg sqlc.UpsertChanPolicyExtraTypeParams) error
 	GetChannelPolicyExtraTypesBatch(ctx context.Context, policyIds []int64) ([]sqlc.GetChannelPolicyExtraTypesBatchRow, error)
@@ -459,13 +460,26 @@ func (s *SQLStore) FetchNodeFeatures(v lnwire.GossipVersion,
 // have their disabled bit on.
 //
 // NOTE: part of the Store interface.
-func (s *SQLStore) DisabledChannelIDs() ([]uint64, error) {
+func (s *SQLStore) DisabledChannelIDs(
+	v lnwire.GossipVersion) ([]uint64, error) {
+
 	var (
 		ctx     = context.TODO()
 		chanIDs []uint64
 	)
 	err := s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		dbChanIDs, err := db.GetV1DisabledSCIDs(ctx)
+		var (
+			dbChanIDs [][]byte
+			err       error
+		)
+		switch v {
+		case gossipV1:
+			dbChanIDs, err = db.GetV1DisabledSCIDs(ctx)
+		case gossipV2:
+			dbChanIDs, err = db.GetV2DisabledSCIDs(ctx)
+		default:
+			return fmt.Errorf("unsupported gossip version: %d", v)
+		}
 		if err != nil {
 			return fmt.Errorf("unable to fetch disabled "+
 				"channels: %w", err)
diff --git a/sqldb/sqlc/graph.sql.go b/sqldb/sqlc/graph.sql.go
index aa2a046..b0d1f78 100644
--- a/sqldb/sqlc/graph.sql.go
+++ b/sqldb/sqlc/graph.sql.go
@@ -2410,6 +2410,41 @@ func (q *Queries) GetV1DisabledSCIDs(ctx context.Context) ([][]byte, error) {
 	return items, nil
 }
 
+const getV2DisabledSCIDs = `-- name: GetV2DisabledSCIDs :many
+SELECT c.scid
+FROM graph_channels c
+    JOIN graph_channel_policies cp ON cp.channel_id = c.id
+WHERE COALESCE(cp.disable_flags, 0) != 0
+AND c.version = 2
+GROUP BY c.scid
+HAVING COUNT(*) > 1
+`
+
+// NOTE: this is V2 specific since V2 uses a disable flag
+// bit vector instead of a single boolean.
+func (q *Queries) GetV2DisabledSCIDs(ctx context.Context) ([][]byte, error) {
+	rows, err := q.db.QueryContext(ctx, getV2DisabledSCIDs)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items [][]byte
+	for rows.Next() {
+		var scid []byte
+		if err := rows.Scan(&scid); err != nil {
+			return nil, err
+		}
+		items = append(items, scid)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const getZombieChannel = `-- name: GetZombieChannel :one
 SELECT scid, version, node_key_1, node_key_2
 FROM graph_zombie_channels
diff --git a/sqldb/sqlc/querier.go b/sqldb/sqlc/querier.go
index d26f845..8948955 100644
--- a/sqldb/sqlc/querier.go
+++ b/sqldb/sqlc/querier.go
@@ -83,6 +83,9 @@ type Querier interface {
 	// structure will have a more complex disabled bit vector
 	// and so the query for V2 may differ.
 	GetV1DisabledSCIDs(ctx context.Context) ([][]byte, error)
+	// NOTE: this is V2 specific since V2 uses a disable flag
+	// bit vector instead of a single boolean.
+	GetV2DisabledSCIDs(ctx context.Context) ([][]byte, error)
 	GetZombieChannel(ctx context.Context, arg GetZombieChannelParams) (GraphZombieChannel, error)
 	GetZombieChannelsSCIDs(ctx context.Context, arg GetZombieChannelsSCIDsParams) ([]GraphZombieChannel, error)
 	HighestSCID(ctx context.Context, version int16) ([]byte, error)
diff --git a/sqldb/sqlc/queries/graph.sql b/sqldb/sqlc/queries/graph.sql
index a4a42a1..2cc2b22 100644
--- a/sqldb/sqlc/queries/graph.sql
+++ b/sqldb/sqlc/queries/graph.sql
@@ -1066,6 +1066,17 @@ AND c.version = 1
 GROUP BY c.scid
 HAVING COUNT(*) > 1;
 
+-- name: GetV2DisabledSCIDs :many
+SELECT c.scid
+FROM graph_channels c
+    JOIN graph_channel_policies cp ON cp.channel_id = c.id
+-- NOTE: this is V2 specific since V2 uses a disable flag
+-- bit vector instead of a single boolean.
+WHERE COALESCE(cp.disable_flags, 0) != 0
+AND c.version = 2
+GROUP BY c.scid
+HAVING COUNT(*) > 1;
+
 -- name: DeleteChannelPolicyExtraTypes :exec
 DELETE FROM graph_channel_policy_extra_types
 WHERE channel_policy_id = $1;