// AppendEntry adds the given Entry to the Source's facts or edges. It is
// assumed that src.Ticket == kytheuri.ToString(e.Source).
func AppendEntry(src *ipb.Source, e *spb.Entry) {
if graphstore.IsEdge(e) {
kind, ordinal, _ := schema.ParseOrdinal(e.EdgeKind)
group, ok := src.EdgeGroups[kind]
if !ok {
group = &ipb.Source_EdgeGroup{}
src.EdgeGroups[kind] = group
}
ticket := kytheuri.ToString(e.Target)
ord := int32(ordinal)
for _, edge := range group.Edges {
if edge.Ticket == ticket && edge.Ordinal == ord {
// Don't add duplicate edge
return
}
}
group.Edges = append(group.Edges, &ipb.Source_Edge{
Ticket: ticket,
Ordinal: ord,
})
} else {
src.Facts[e.FactName] = e.FactValue
}
}
// getEdges returns edgeTargets with the given node as their source. Only edge
// entries that return true when applied to pred are returned.
func getEdges(ctx context.Context, gs graphstore.Service, node *spb.VName, pred func(*spb.Entry) bool) ([]*edgeTarget, error) {
var targets []*edgeTarget
if err := gs.Read(ctx, &spb.ReadRequest{
Source: node,
EdgeKind: "*",
}, func(entry *spb.Entry) error {
if graphstore.IsEdge(entry) && pred(entry) {
edgeKind, ordinal, _ := schema.ParseOrdinal(entry.EdgeKind)
targets = append(targets, &edgeTarget{edgeKind, entry.Target, int32(ordinal)})
}
return nil
}); err != nil {
return nil, fmt.Errorf("read error: %v", err)
}
return targets, nil
}
// Edges implements part of the Service interface.
func (g *GraphStoreService) Edges(ctx context.Context, req *xpb.EdgesRequest) (*xpb.EdgesReply, error) {
if len(req.Ticket) == 0 {
return nil, errors.New("no tickets specified")
} else if req.PageToken != "" {
return nil, errors.New("UNIMPLEMENTED: page_token")
}
patterns := xrefs.ConvertFilters(req.Filter)
allowedKinds := stringset.New(req.Kind...)
targetSet := stringset.New()
reply := &xpb.EdgesReply{
EdgeSets: make(map[string]*xpb.EdgeSet),
Nodes: make(map[string]*xpb.NodeInfo),
}
for _, ticket := range req.Ticket {
vname, err := kytheuri.ToVName(ticket)
if err != nil {
return nil, fmt.Errorf("invalid ticket %q: %v", ticket, err)
}
var (
// EdgeKind -> TargetTicket -> OrdinalSet
filteredEdges = make(map[string]map[string]map[int32]struct{})
filteredFacts = make(map[string][]byte)
)
if err := g.gs.Read(ctx, &spb.ReadRequest{
Source: vname,
EdgeKind: "*",
}, func(entry *spb.Entry) error {
edgeKind := entry.EdgeKind
if edgeKind == "" {
// node fact
if len(patterns) > 0 && xrefs.MatchesAny(entry.FactName, patterns) {
filteredFacts[entry.FactName] = entry.FactValue
}
} else {
// edge
edgeKind, ordinal, _ := schema.ParseOrdinal(edgeKind)
if len(req.Kind) == 0 || allowedKinds.Contains(edgeKind) {
targets, ok := filteredEdges[edgeKind]
if !ok {
targets = make(map[string]map[int32]struct{})
filteredEdges[edgeKind] = targets
}
ticket := kytheuri.ToString(entry.Target)
ordSet, ok := targets[ticket]
if !ok {
ordSet = make(map[int32]struct{})
targets[ticket] = ordSet
}
ordSet[int32(ordinal)] = struct{}{}
}
}
return nil
}); err != nil {
return nil, fmt.Errorf("failed to retrieve entries for ticket %q", ticket)
}
// Only add a EdgeSet if there are targets for the requested edge kinds.
if len(filteredEdges) > 0 {
groups := make(map[string]*xpb.EdgeSet_Group)
for edgeKind, targets := range filteredEdges {
g := &xpb.EdgeSet_Group{}
for target, ordinals := range targets {
for ordinal := range ordinals {
g.Edge = append(g.Edge, &xpb.EdgeSet_Group_Edge{
TargetTicket: target,
Ordinal: ordinal,
})
}
targetSet.Add(target)
}
groups[edgeKind] = g
}
reply.EdgeSets[ticket] = &xpb.EdgeSet{
Groups: groups,
}
// In addition, only add a NodeInfo if the filters have resulting facts.
if len(filteredFacts) > 0 {
reply.Nodes[ticket] = &xpb.NodeInfo{
Facts: filteredFacts,
}
}
}
}
// Only request Nodes when there are fact filters given.
if len(req.Filter) > 0 {
// Eliminate redundant work by removing already requested nodes from targetSet
for ticket := range reply.Nodes {
targetSet.Remove(ticket)
}
// Batch request all leftover target nodes
nodesReply, err := g.Nodes(ctx, &xpb.NodesRequest{
Ticket: targetSet.Slice(),
Filter: req.Filter,
})
if err != nil {
return nil, fmt.Errorf("failure getting target nodes: %v", err)
}
for ticket, node := range nodesReply.Nodes {
reply.Nodes[ticket] = node
}
}
return reply, nil
}
func (d *DB) copyEntries(entries <-chan *spb.Entry) error {
// Start a transaction for a COPY statement per table
nodesTx, err := d.Begin()
if err != nil {
return err
}
edgesTx, err := d.Begin()
if err != nil {
return err
}
// Create each table in their corresponding transactions to speed up COPY
if _, err := nodesTx.Exec(createNodesTable); err != nil {
return fmt.Errorf("error truncating Nodes table: %v", err)
} else if _, err := edgesTx.Exec(createEdgeTable); err != nil {
return fmt.Errorf("error truncating Edges table: %v", err)
}
copyNode, err := nodesTx.Prepare(pq.CopyIn(
"nodes",
"ticket",
"node_kind",
"subkind",
"text",
"text_encoding",
"start_offset",
"end_offset",
"snippet_start",
"snippet_end",
"other_facts_num",
"other_facts",
))
if err != nil {
return fmt.Errorf("error preparing Nodes copy: %v", err)
}
copyEdge, err := edgesTx.Prepare(pq.CopyIn(
"edges",
"source",
"kind",
"target",
"ordinal",
))
if err != nil {
return fmt.Errorf("error preparing Edges copy: %v", err)
}
var node srvpb.Node
var nodeKind string
var subkind, textEncoding *string
var text *[]byte
var startOffset, endOffset, snippetStart, snippetEnd *int64
for e := range entries {
if graphstore.IsNodeFact(e) {
ticket := kytheuri.ToString(e.Source)
if node.Ticket != "" && node.Ticket != ticket {
nodeTicket := node.Ticket
node.Ticket = ""
var rec []byte
if len(node.Fact) > 0 {
rec, err = proto.Marshal(&node)
if err != nil {
return fmt.Errorf("error marshaling facts: %v", err)
}
}
if text != nil && textEncoding == nil {
textEncoding = proto.String(schema.DefaultTextEncoding)
}
if _, err := copyNode.Exec(
nodeTicket,
nodeKind,
subkind,
text,
textEncoding,
startOffset,
endOffset,
snippetStart,
snippetEnd,
len(node.Fact),
rec,
); err != nil {
return fmt.Errorf("error copying node: %v", err)
}
node.Fact, text = node.Fact[0:0], nil
nodeKind = ""
subkind, textEncoding = nil, nil
startOffset, endOffset, snippetStart, snippetEnd = nil, nil, nil, nil
}
if node.Ticket == "" {
node.Ticket = ticket
}
switch e.FactName {
case schema.NodeKindFact:
nodeKind = string(e.FactValue)
case schema.SubkindFact:
subkind = proto.String(string(e.FactValue))
case schema.TextFact:
text = &e.FactValue
case schema.TextEncodingFact:
textEncoding = proto.String(string(e.FactValue))
case schema.AnchorStartFact:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
startOffset = proto.Int64(n)
}
case schema.AnchorEndFact:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
endOffset = proto.Int64(n)
}
case schema.SnippetStartFact:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
snippetStart = proto.Int64(n)
}
case schema.SnippetEndFact:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
snippetEnd = proto.Int64(n)
}
default:
node.Fact = append(node.Fact, &cpb.Fact{
Name: e.FactName,
Value: e.FactValue,
})
}
} else if schema.EdgeDirection(e.EdgeKind) == schema.Forward {
kind, ordinal, _ := schema.ParseOrdinal(e.EdgeKind)
ticket := kytheuri.ToString(e.Source)
if _, err := copyEdge.Exec(ticket, kind, kytheuri.ToString(e.Target), ordinal); err != nil {
return fmt.Errorf("error copying edge: %v", err)
}
}
}
if _, err := copyNode.Exec(); err != nil {
return fmt.Errorf("error flushing nodes: %v", err)
} else if _, err := copyEdge.Exec(); err != nil {
return fmt.Errorf("error flushing edges: %v", err)
}
if err := nodesTx.Commit(); err != nil {
return fmt.Errorf("error committing Nodes transaction: %v", err)
} else if err := edgesTx.Commit(); err != nil {
return fmt.Errorf("error committing Edges transaction: %v", err)
}
return nil
}