func collectNodes(nodeEntries <-chan *spb.Entry) <-chan *srvpb.Node {
nodes := make(chan *srvpb.Node)
go func() {
var (
node *srvpb.Node
vname *spb.VName
)
for e := range nodeEntries {
if node != nil && !compare.VNamesEqual(e.Source, vname) {
nodes <- node
node = nil
vname = nil
}
if node == nil {
vname = e.Source
ticket := kytheuri.ToString(vname)
node = &srvpb.Node{Ticket: ticket}
}
node.Fact = append(node.Fact, &srvpb.Node_Fact{
Name: e.FactName,
Value: e.FactValue,
})
}
if node != nil {
nodes <- node
}
close(nodes)
}()
return nodes
}
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(facts.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 facts.NodeKind:
nodeKind = string(e.FactValue)
case facts.Subkind:
subkind = proto.String(string(e.FactValue))
case facts.Text:
text = &e.FactValue
case facts.TextEncoding:
textEncoding = proto.String(string(e.FactValue))
case facts.AnchorStart:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
startOffset = proto.Int64(n)
}
case facts.AnchorEnd:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
endOffset = proto.Int64(n)
}
case facts.SnippetStart:
n, err := strconv.ParseInt(string(e.FactValue), 10, 64)
if err == nil {
snippetStart = proto.Int64(n)
}
case facts.SnippetEnd:
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 edges.IsForward(e.EdgeKind) {
kind, ordinal, _ := edges.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
}