Example #1
0
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
}
Example #2
0
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
}