graph/db: version ForEachSourceNodeChannel

diff --git a/graph/db/graph.go b/graph/db/graph.go
index baee3ec..baac06e 100644
--- a/graph/db/graph.go
+++ b/graph/db/graph.go
@@ -587,14 +587,6 @@ func (c *ChannelGraph) UpdateEdgePolicy(ctx context.Context,
 	return nil
 }
 
-// ForEachSourceNodeChannel iterates through all channels of the source node.
-func (c *ChannelGraph) ForEachSourceNodeChannel(ctx context.Context,
-	cb func(chanPoint wire.OutPoint, havePolicy bool,
-		otherNode *models.Node) error, reset func()) error {
-
-	return c.db.ForEachSourceNodeChannel(ctx, cb, reset)
-}
-
 // ForEachNodeChannel iterates through all channels of the given node.
 func (c *ChannelGraph) ForEachNodeChannel(ctx context.Context,
 	v lnwire.GossipVersion, nodePub route.Vertex,
@@ -906,6 +898,14 @@ func (c *VersionedGraph) HasChannelEdge(chanID uint64) (bool, bool, error) {
 	return c.db.HasChannelEdge(c.v, chanID)
 }
 
+// ForEachSourceNodeChannel iterates through all channels of the source node.
+func (c *VersionedGraph) ForEachSourceNodeChannel(ctx context.Context,
+	cb func(chanPoint wire.OutPoint, havePolicy bool,
+		otherNode *models.Node) error, reset func()) error {
+
+	return c.db.ForEachSourceNodeChannel(ctx, c.v, cb, reset)
+}
+
 // ForEachNodeChannel iterates through all channels of the given node.
 func (c *VersionedGraph) ForEachNodeChannel(ctx context.Context,
 	nodePub route.Vertex, cb func(*models.ChannelEdgeInfo,
diff --git a/graph/db/graph_test.go b/graph/db/graph_test.go
index 2edfc06..5e71e29 100644
--- a/graph/db/graph_test.go
+++ b/graph/db/graph_test.go
@@ -154,6 +154,14 @@ var versionedTests = []versionedTest{
 		name: "incomplete channel policies",
 		test: testIncompleteChannelPolicies,
 	},
+	{
+		name: "add channel edge shell nodes",
+		test: testAddChannelEdgeShellNodes,
+	},
+	{
+		name: "for each source node channel",
+		test: testForEachSourceNodeChannel,
+	},
 	{
 		name: "partial node",
 		test: testPartialNode,
@@ -1605,16 +1613,16 @@ func testAddEdgeProof(t *testing.T, v lnwire.GossipVersion) {
 	require.NotNil(t, dbEdge2.AuthProof)
 }
 
-// TestForEachSourceNodeChannel tests that the ForEachSourceNodeChannel
+// testForEachSourceNodeChannel tests that the ForEachSourceNodeChannel
 // correctly iterates through the channels of the set source node.
-func TestForEachSourceNodeChannel(t *testing.T) {
+func testForEachSourceNodeChannel(t *testing.T, v lnwire.GossipVersion) {
 	t.Parallel()
 	ctx := t.Context()
 
-	graph := MakeTestGraph(t)
+	graph := NewVersionedGraph(MakeTestGraph(t), v)
 
 	// Create a source node (A) and set it as such in the DB.
-	nodeA := createTestVertex(t, lnwire.GossipVersion1)
+	nodeA := createTestVertex(t, v)
 	require.NoError(t, graph.SetSourceNode(ctx, nodeA))
 
 	// Now, create a few more nodes (B, C, D) along with some channels
@@ -1630,32 +1638,40 @@ func TestForEachSourceNodeChannel(t *testing.T) {
 	// outgoing policy but for the A-C channel, we will set only an incoming
 	// policy.
 
-	nodeB := createTestVertex(t, lnwire.GossipVersion1)
-	nodeC := createTestVertex(t, lnwire.GossipVersion1)
-	nodeD := createTestVertex(t, lnwire.GossipVersion1)
+	nodeB := createTestVertex(t, v)
+	nodeC := createTestVertex(t, v)
+	nodeD := createTestVertex(t, v)
 
-	abEdge, abPolicy1, abPolicy2 := createChannelEdge(nodeA, nodeB)
+	abEdge, _ := createEdge(v, 100, 0, 0, 0, nodeA, nodeB)
 	require.NoError(t, graph.AddChannelEdge(ctx, abEdge))
-	acEdge, acPolicy1, acPolicy2 := createChannelEdge(nodeA, nodeC)
+	acEdge, _ := createEdge(v, 200, 0, 0, 1, nodeA, nodeC)
 	require.NoError(t, graph.AddChannelEdge(ctx, acEdge))
-	bdEdge, _, _ := createChannelEdge(nodeB, nodeD)
+	bdEdge, _ := createEdge(v, 300, 0, 0, 2, nodeB, nodeD)
 	require.NoError(t, graph.AddChannelEdge(ctx, bdEdge))
 
-	// Figure out which of the policies returned above are node A's so that
-	// we know which to persist.
-	//
-	// First, set the outgoing policy for the A-B channel.
-	abPolicyAOutgoing := abPolicy1
-	if !bytes.Equal(abPolicy1.ToNode[:], nodeB.PubKeyBytes[:]) {
-		abPolicyAOutgoing = abPolicy2
+	newPolicy := func(edge *models.ChannelEdgeInfo, fromNode,
+		toNode route.Vertex) *models.ChannelEdgePolicy {
+
+		isNode1 := bytes.Equal(fromNode[:], edge.NodeKey1Bytes[:])
+		policy := newEdgePolicy(
+			v, edge.ChannelID, nextUpdateTime().Unix(), isNode1,
+		)
+		policy.ToNode = toNode
+		policy.SigBytes = testSig.Serialize()
+
+		return policy
 	}
+
+	// First, set the outgoing policy for the A-B channel.
+	abPolicyAOutgoing := newPolicy(
+		abEdge, nodeA.PubKeyBytes, nodeB.PubKeyBytes,
+	)
 	require.NoError(t, graph.UpdateEdgePolicy(ctx, abPolicyAOutgoing))
 
 	// Now, set the incoming policy for the A-C channel.
-	acPolicyAIncoming := acPolicy1
-	if !bytes.Equal(acPolicy1.ToNode[:], nodeA.PubKeyBytes[:]) {
-		acPolicyAIncoming = acPolicy2
-	}
+	acPolicyAIncoming := newPolicy(
+		acEdge, nodeC.PubKeyBytes, nodeA.PubKeyBytes,
+	)
 	require.NoError(t, graph.UpdateEdgePolicy(ctx, acPolicyAIncoming))
 
 	type sourceNodeChan struct {
@@ -1677,21 +1693,24 @@ func TestForEachSourceNodeChannel(t *testing.T) {
 
 	// Now, we'll use the ForEachSourceNodeChannel and assert that it
 	// returns the expected data in the call-back.
-	err := graph.ForEachSourceNodeChannel(ctx, func(chanPoint wire.OutPoint,
-		havePolicy bool, otherNode *models.Node) error {
+	err := graph.ForEachSourceNodeChannel(
+		ctx, func(chanPoint wire.OutPoint, havePolicy bool,
+			otherNode *models.Node) error {
 
-		require.Contains(t, expectedSrcChans, chanPoint)
-		expected := expectedSrcChans[chanPoint]
+			require.Contains(t, expectedSrcChans, chanPoint)
+			expected := expectedSrcChans[chanPoint]
 
-		require.Equal(
-			t, expected.otherNode[:], otherNode.PubKeyBytes[:],
-		)
-		require.Equal(t, expected.havePolicy, havePolicy)
+			require.Equal(
+				t, expected.otherNode[:],
+				otherNode.PubKeyBytes[:],
+			)
+			require.Equal(t, expected.havePolicy, havePolicy)
 
-		delete(expectedSrcChans, chanPoint)
+			delete(expectedSrcChans, chanPoint)
 
-		return nil
-	}, func() {})
+			return nil
+		}, func() {},
+	)
 	require.NoError(t, err)
 	require.Empty(t, expectedSrcChans)
 }
@@ -4106,25 +4125,25 @@ func TestPruneGraphNodes(t *testing.T) {
 	require.NotNil(t, err)
 }
 
-// TestAddChannelEdgeShellNodes tests that when we attempt to add a ChannelEdge
+// testAddChannelEdgeShellNodes tests that when we attempt to add a ChannelEdge
 // to the graph, one or both of the nodes the edge involves aren't found in the
 // database, then shell edges are created for each node if needed.
-func TestAddChannelEdgeShellNodes(t *testing.T) {
+func testAddChannelEdgeShellNodes(t *testing.T, v lnwire.GossipVersion) {
 	t.Parallel()
 	ctx := t.Context()
 
-	graph := NewVersionedGraph(MakeTestGraph(t), lnwire.GossipVersion1)
+	graph := NewVersionedGraph(MakeTestGraph(t), v)
 
 	// To start, we'll create two nodes, and only add one of them to the
 	// channel graph.
-	node1 := createTestVertex(t, lnwire.GossipVersion1)
+	node1 := createTestVertex(t, v)
 	require.NoError(t, graph.SetSourceNode(ctx, node1))
-	node2 := createTestVertex(t, lnwire.GossipVersion1)
+	node2 := createTestVertex(t, v)
 
 	// We'll now create an edge between the two nodes, as a result, node2
 	// should be inserted into the database as a shell node.
 	edgeInfo, _ := createEdge(
-		lnwire.GossipVersion1, 100, 0, 0, 0, node1, node2,
+		v, 100, 0, 0, 0, node1, node2,
 	)
 	require.NoError(t, graph.AddChannelEdge(ctx, edgeInfo))
 
diff --git a/graph/db/interfaces.go b/graph/db/interfaces.go
index 0936386..9c904b5 100644
--- a/graph/db/interfaces.go
+++ b/graph/db/interfaces.go
@@ -59,7 +59,7 @@ type Store interface { //nolint:interfacebloat
 	// node, executing the passed callback on each. The call-back is
 	// provided with the channel's outpoint, whether we have a policy for
 	// the channel and the channel peer's node information.
-	ForEachSourceNodeChannel(ctx context.Context,
+	ForEachSourceNodeChannel(ctx context.Context, v lnwire.GossipVersion,
 		cb func(chanPoint wire.OutPoint, havePolicy bool,
 			otherNode *models.Node) error,
 		reset func()) error
diff --git a/graph/db/kv_store.go b/graph/db/kv_store.go
index dcb84d3..7c81d5f 100644
--- a/graph/db/kv_store.go
+++ b/graph/db/kv_store.go
@@ -3722,8 +3722,13 @@ func (c *KVStore) ForEachNodeChannel(_ context.Context,
 // channel's outpoint, whether we have a policy for the channel and the channel
 // peer's node information.
 func (c *KVStore) ForEachSourceNodeChannel(_ context.Context,
-	cb func(chanPoint wire.OutPoint, havePolicy bool,
-		otherNode *models.Node) error, reset func()) error {
+	v lnwire.GossipVersion, cb func(chanPoint wire.OutPoint,
+		havePolicy bool, otherNode *models.Node) error,
+	reset func()) error {
+
+	if v != lnwire.GossipVersion1 {
+		return ErrVersionNotSupportedForKVDB
+	}
 
 	return kvdb.View(c.db, func(tx kvdb.RTx) error {
 		nodes := tx.ReadBucket(nodeBucket)
diff --git a/graph/db/sql_store.go b/graph/db/sql_store.go
index ce95c10..a6b05da 100644
--- a/graph/db/sql_store.go
+++ b/graph/db/sql_store.go
@@ -902,20 +902,19 @@ func (s *SQLStore) updateEdgeCache(e *models.ChannelEdgePolicy,
 //
 // NOTE: part of the Store interface.
 func (s *SQLStore) ForEachSourceNodeChannel(ctx context.Context,
-	cb func(chanPoint wire.OutPoint, havePolicy bool,
-		otherNode *models.Node) error, reset func()) error {
+	v lnwire.GossipVersion, cb func(chanPoint wire.OutPoint,
+		havePolicy bool, otherNode *models.Node) error,
+	reset func()) error {
 
 	return s.db.ExecTx(ctx, sqldb.ReadTxOpt(), func(db SQLQueries) error {
-		nodeID, nodePub, err := s.getSourceNode(
-			ctx, db, lnwire.GossipVersion1,
-		)
+		nodeID, nodePub, err := s.getSourceNode(ctx, db, v)
 		if err != nil {
 			return fmt.Errorf("unable to fetch source node: %w",
 				err)
 		}
 
 		return forEachNodeChannel(
-			ctx, db, s.cfg, lnwire.GossipVersion1, nodeID,
+			ctx, db, s.cfg, v, nodeID,
 			func(info *models.ChannelEdgeInfo,
 				outPolicy *models.ChannelEdgePolicy,
 				_ *models.ChannelEdgePolicy) error {
@@ -937,8 +936,8 @@ func (s *SQLStore) ForEachSourceNodeChannel(ctx context.Context,
 				}
 
 				_, otherNode, err := getNodeByPubKey(
-					ctx, s.cfg.QueryCfg, db,
-					lnwire.GossipVersion1, otherNodePub,
+					ctx, s.cfg.QueryCfg, db, v,
+					otherNodePub,
 				)
 				if err != nil {
 					return fmt.Errorf("unable to fetch "+
diff --git a/server.go b/server.go
index ad97261..9092463 100644
--- a/server.go
+++ b/server.go
@@ -1128,7 +1128,7 @@ func newServer(ctx context.Context, cfg *Config, listenAddrs []net.Addr,
 		DefaultRoutingPolicy: cc.RoutingPolicy,
 		ForAllOutgoingChannels: func(ctx context.Context,
 			cb func(*models.ChannelEdgeInfo,
-				*models.ChannelEdgePolicy) error,
+			*models.ChannelEdgePolicy) error,
 			reset func()) error {
 
 			return s.v1Graph.ForEachNodeChannel(
@@ -3576,7 +3576,10 @@ func (s *server) establishPersistentConnections(ctx context.Context) error {
 		graphAddrs[pubStr] = n
 		return nil
 	}
-	err = s.graphDB.ForEachSourceNodeChannel(
+
+	// TODO(elle): for now, we only fetch our V1 channels. This should be
+	//  updated to fetch channels across all versions.
+	err = s.v1Graph.ForEachSourceNodeChannel(
 		ctx, forEachSrcNodeChan, func() {
 			clear(graphAddrs)
 		},