graph/db+sqldb: queries and logic for batch fetching node data

In this commit, we add the queries that will be needed to batch-fetch
the data of a set of nodes. The logic for using these new queries is
also added but not used yet.

graph/db/sql_store.go

@@ -69,15 +69,18 @@ type SQLQueries interface {
 	DeleteNode(ctx context.Context, id int64) error
 
 	GetExtraNodeTypes(ctx context.Context, nodeID int64) ([]sqlc.GraphNodeExtraType, error)
+	GetNodeExtraTypesBatch(ctx context.Context, ids []int64) ([]sqlc.GraphNodeExtraType, error)
 	UpsertNodeExtraType(ctx context.Context, arg sqlc.UpsertNodeExtraTypeParams) error
 	DeleteExtraNodeType(ctx context.Context, arg sqlc.DeleteExtraNodeTypeParams) error
 
 	InsertNodeAddress(ctx context.Context, arg sqlc.InsertNodeAddressParams) error
 	GetNodeAddresses(ctx context.Context, nodeID int64) ([]sqlc.GetNodeAddressesRow, error)
+	GetNodeAddressesBatch(ctx context.Context, ids []int64) ([]sqlc.GraphNodeAddress, error)
 	DeleteNodeAddresses(ctx context.Context, nodeID int64) error
 
 	InsertNodeFeature(ctx context.Context, arg sqlc.InsertNodeFeatureParams) error
 	GetNodeFeatures(ctx context.Context, nodeID int64) ([]sqlc.GraphNodeFeature, error)
+	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
 
@@ -4672,3 +4675,155 @@ func channelIDToBytes(channelID uint64) []byte {
 
 	return chanIDB[:]
 }
+
+// batchNodeData holds all the related data for a batch of nodes.
+type batchNodeData struct {
+	// features is a map from a DB node ID to the feature bits for that
+	// node.
+	features map[int64][]int
+
+	// addresses is a map from a DB node ID to the node's addresses.
+	addresses map[int64][]nodeAddress
+
+	// extraFields is a map from a DB node ID to the extra signed fields
+	// for that node.
+	extraFields map[int64]map[uint64][]byte
+}
+
+// nodeAddress holds the address type, position and address string for a
+// node. This is used to batch the fetching of node addresses.
+type nodeAddress struct {
+	addrType dbAddressType
+	position int32
+	address  string
+}
+
+// batchLoadNodeData loads all related data for a batch of node IDs using the
+// provided SQLQueries interface. It returns a batchNodeData instance containing
+// the node features, addresses and extra signed fields.
+func batchLoadNodeData(ctx context.Context, cfg *sqldb.PagedQueryConfig,
+	db SQLQueries, nodeIDs []int64) (*batchNodeData, error) {
+
+	// Batch load the node features.
+	features, err := batchLoadNodeFeaturesHelper(ctx, cfg, db, nodeIDs)
+	if err != nil {
+		return nil, fmt.Errorf("unable to batch load node "+
+			"features: %w", err)
+	}
+
+	// Batch load the node addresses.
+	addrs, err := batchLoadNodeAddressesHelper(ctx, cfg, db, nodeIDs)
+	if err != nil {
+		return nil, fmt.Errorf("unable to batch load node "+
+			"addresses: %w", err)
+	}
+
+	// Batch load the node extra signed fields.
+	extraTypes, err := batchLoadNodeExtraTypesHelper(ctx, cfg, db, nodeIDs)
+	if err != nil {
+		return nil, fmt.Errorf("unable to batch load node extra "+
+			"signed fields: %w", err)
+	}
+
+	return &batchNodeData{
+		features:    features,
+		addresses:   addrs,
+		extraFields: extraTypes,
+	}, nil
+}
+
+// batchLoadNodeFeaturesHelper loads node features for a batch of node IDs
+// using ExecutePagedQuery wrapper around the GetNodeFeaturesBatch query.
+func batchLoadNodeFeaturesHelper(ctx context.Context,
+	cfg *sqldb.PagedQueryConfig, db SQLQueries,
+	nodeIDs []int64) (map[int64][]int, error) {
+
+	features := make(map[int64][]int)
+
+	return features, sqldb.ExecutePagedQuery(
+		ctx, cfg, nodeIDs,
+		func(id int64) int64 {
+			return id
+		},
+		func(ctx context.Context, ids []int64) ([]sqlc.GraphNodeFeature,
+			error) {
+
+			return db.GetNodeFeaturesBatch(ctx, ids)
+		},
+		func(ctx context.Context, feature sqlc.GraphNodeFeature) error {
+			features[feature.NodeID] = append(
+				features[feature.NodeID],
+				int(feature.FeatureBit),
+			)
+
+			return nil
+		},
+	)
+}
+
+// batchLoadNodeAddressesHelper loads node addresses using ExecutePagedQuery
+// wrapper around the GetNodeAddressesBatch query. It returns a map from
+// node ID to a slice of nodeAddress structs.
+func batchLoadNodeAddressesHelper(ctx context.Context,
+	cfg *sqldb.PagedQueryConfig, db SQLQueries,
+	nodeIDs []int64) (map[int64][]nodeAddress, error) {
+
+	addrs := make(map[int64][]nodeAddress)
+
+	return addrs, sqldb.ExecutePagedQuery(
+		ctx, cfg, nodeIDs,
+		func(id int64) int64 {
+			return id
+		},
+		func(ctx context.Context, ids []int64) ([]sqlc.GraphNodeAddress,
+			error) {
+
+			return db.GetNodeAddressesBatch(ctx, ids)
+		},
+		func(ctx context.Context, addr sqlc.GraphNodeAddress) error {
+			addrs[addr.NodeID] = append(
+				addrs[addr.NodeID], nodeAddress{
+					addrType: dbAddressType(addr.Type),
+					position: addr.Position,
+					address:  addr.Address,
+				},
+			)
+
+			return nil
+		},
+	)
+}
+
+// batchLoadNodeExtraTypesHelper loads node extra type bytes for a batch of
+// node IDs using ExecutePagedQuery wrapper around the GetNodeExtraTypesBatch
+// query.
+func batchLoadNodeExtraTypesHelper(ctx context.Context,
+	cfg *sqldb.PagedQueryConfig, db SQLQueries,
+	nodeIDs []int64) (map[int64]map[uint64][]byte, error) {
+
+	extraFields := make(map[int64]map[uint64][]byte)
+
+	callback := func(ctx context.Context,
+		field sqlc.GraphNodeExtraType) error {
+
+		if extraFields[field.NodeID] == nil {
+			extraFields[field.NodeID] = make(map[uint64][]byte)
+		}
+		extraFields[field.NodeID][uint64(field.Type)] = field.Value
+
+		return nil
+	}
+
+	return extraFields, sqldb.ExecutePagedQuery(
+		ctx, cfg, nodeIDs,
+		func(id int64) int64 {
+			return id
+		},
+		func(ctx context.Context, ids []int64) (
+			[]sqlc.GraphNodeExtraType, error) {
+
+			return db.GetNodeExtraTypesBatch(ctx, ids)
+		},
+		callback,
+	)
+}

sqldb/sqlc/graph.sql.go

@@ -1473,6 +1473,51 @@ func (q *Queries) GetNodeAddresses(ctx context.Context, nodeID int64) ([]GetNode
 	return items, nil
 }
 
+const getNodeAddressesBatch = `-- name: GetNodeAddressesBatch :many
+SELECT node_id, type, position, address
+FROM graph_node_addresses
+WHERE node_id IN (/*SLICE:ids*/?)
+ORDER BY node_id, type, position
+`
+
+func (q *Queries) GetNodeAddressesBatch(ctx context.Context, ids []int64) ([]GraphNodeAddress, error) {
+	query := getNodeAddressesBatch
+	var queryParams []interface{}
+	if len(ids) > 0 {
+		for _, v := range ids {
+			queryParams = append(queryParams, v)
+		}
+		query = strings.Replace(query, "/*SLICE:ids*/?", makeQueryParams(len(queryParams), len(ids)), 1)
+	} else {
+		query = strings.Replace(query, "/*SLICE:ids*/?", "NULL", 1)
+	}
+	rows, err := q.db.QueryContext(ctx, query, queryParams...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []GraphNodeAddress
+	for rows.Next() {
+		var i GraphNodeAddress
+		if err := rows.Scan(
+			&i.NodeID,
+			&i.Type,
+			&i.Position,
+			&i.Address,
+		); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const getNodeByPubKey = `-- name: GetNodeByPubKey :one
 SELECT id, version, pub_key, alias, last_update, color, signature
 FROM graph_nodes
@@ -1500,6 +1545,46 @@ func (q *Queries) GetNodeByPubKey(ctx context.Context, arg GetNodeByPubKeyParams
 	return i, err
 }
 
+const getNodeExtraTypesBatch = `-- name: GetNodeExtraTypesBatch :many
+SELECT node_id, type, value
+FROM graph_node_extra_types
+WHERE node_id IN (/*SLICE:ids*/?)
+ORDER BY node_id, type
+`
+
+func (q *Queries) GetNodeExtraTypesBatch(ctx context.Context, ids []int64) ([]GraphNodeExtraType, error) {
+	query := getNodeExtraTypesBatch
+	var queryParams []interface{}
+	if len(ids) > 0 {
+		for _, v := range ids {
+			queryParams = append(queryParams, v)
+		}
+		query = strings.Replace(query, "/*SLICE:ids*/?", makeQueryParams(len(queryParams), len(ids)), 1)
+	} else {
+		query = strings.Replace(query, "/*SLICE:ids*/?", "NULL", 1)
+	}
+	rows, err := q.db.QueryContext(ctx, query, queryParams...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []GraphNodeExtraType
+	for rows.Next() {
+		var i GraphNodeExtraType
+		if err := rows.Scan(&i.NodeID, &i.Type, &i.Value); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const getNodeFeatures = `-- name: GetNodeFeatures :many
 SELECT node_id, feature_bit
 FROM graph_node_features
@@ -1529,6 +1614,46 @@ func (q *Queries) GetNodeFeatures(ctx context.Context, nodeID int64) ([]GraphNod
 	return items, nil
 }
 
+const getNodeFeaturesBatch = `-- name: GetNodeFeaturesBatch :many
+SELECT node_id, feature_bit
+FROM graph_node_features
+WHERE node_id IN (/*SLICE:ids*/?)
+ORDER BY node_id, feature_bit
+`
+
+func (q *Queries) GetNodeFeaturesBatch(ctx context.Context, ids []int64) ([]GraphNodeFeature, error) {
+	query := getNodeFeaturesBatch
+	var queryParams []interface{}
+	if len(ids) > 0 {
+		for _, v := range ids {
+			queryParams = append(queryParams, v)
+		}
+		query = strings.Replace(query, "/*SLICE:ids*/?", makeQueryParams(len(queryParams), len(ids)), 1)
+	} else {
+		query = strings.Replace(query, "/*SLICE:ids*/?", "NULL", 1)
+	}
+	rows, err := q.db.QueryContext(ctx, query, queryParams...)
+	if err != nil {
+		return nil, err
+	}
+	defer rows.Close()
+	var items []GraphNodeFeature
+	for rows.Next() {
+		var i GraphNodeFeature
+		if err := rows.Scan(&i.NodeID, &i.FeatureBit); err != nil {
+			return nil, err
+		}
+		items = append(items, i)
+	}
+	if err := rows.Close(); err != nil {
+		return nil, err
+	}
+	if err := rows.Err(); err != nil {
+		return nil, err
+	}
+	return items, nil
+}
+
 const getNodeFeaturesByPubKey = `-- name: GetNodeFeaturesByPubKey :many
 SELECT f.feature_bit
 FROM graph_nodes n

sqldb/sqlc/querier.go

@@ -60,8 +60,11 @@ type Querier interface {
 	GetKVInvoicePaymentHashByAddIndex(ctx context.Context, addIndex int64) ([]byte, error)
 	GetMigration(ctx context.Context, version int32) (time.Time, error)
 	GetNodeAddresses(ctx context.Context, nodeID int64) ([]GetNodeAddressesRow, error)
+	GetNodeAddressesBatch(ctx context.Context, ids []int64) ([]GraphNodeAddress, error)
 	GetNodeByPubKey(ctx context.Context, arg GetNodeByPubKeyParams) (GraphNode, error)
+	GetNodeExtraTypesBatch(ctx context.Context, ids []int64) ([]GraphNodeExtraType, error)
 	GetNodeFeatures(ctx context.Context, nodeID int64) ([]GraphNodeFeature, error)
+	GetNodeFeaturesBatch(ctx context.Context, ids []int64) ([]GraphNodeFeature, error)
 	GetNodeFeaturesByPubKey(ctx context.Context, arg GetNodeFeaturesByPubKeyParams) ([]int32, error)
 	GetNodeIDByPubKey(ctx context.Context, arg GetNodeIDByPubKeyParams) (int64, error)
 	GetNodesByLastUpdateRange(ctx context.Context, arg GetNodesByLastUpdateRangeParams) ([]GraphNode, error)

sqldb/sqlc/queries/graph.sql

@@ -114,6 +114,12 @@ FROM graph_nodes n
 WHERE n.pub_key = $1
   AND n.version = $2;
 
+-- name: GetNodeFeaturesBatch :many
+SELECT node_id, feature_bit
+FROM graph_node_features
+WHERE node_id IN (sqlc.slice('ids')/*SLICE:ids*/)
+ORDER BY node_id, feature_bit;
+
 -- name: DeleteNodeFeature :exec
 DELETE FROM graph_node_features
 WHERE node_id = $1
@@ -140,6 +146,12 @@ FROM graph_node_addresses
 WHERE node_id = $1
 ORDER BY type ASC, position ASC;
 
+-- name: GetNodeAddressesBatch :many
+SELECT node_id, type, position, address
+FROM graph_node_addresses
+WHERE node_id IN (sqlc.slice('ids')/*SLICE:ids*/)
+ORDER BY node_id, type, position;
+
 -- name: GetNodesByLastUpdateRange :many
 SELECT *
 FROM graph_nodes
@@ -170,6 +182,12 @@ SELECT *
 FROM graph_node_extra_types
 WHERE node_id = $1;
 
+-- name: GetNodeExtraTypesBatch :many
+SELECT node_id, type, value
+FROM graph_node_extra_types
+WHERE node_id IN (sqlc.slice('ids')/*SLICE:ids*/)
+ORDER BY node_id, type;
+
 -- name: DeleteExtraNodeType :exec
 DELETE FROM graph_node_extra_types
 WHERE node_id = $1