// BatchWrites returns a channel of WriteRequests for the given entries.
// Consecutive entries with the same Source will be collected in the same
// WriteRequest, with each request containing up to maxSize updates.
func BatchWrites(entries <-chan *spb.Entry, maxSize int) <-chan *spb.WriteRequest {
ch := make(chan *spb.WriteRequest)
go func() {
defer close(ch)
var req *spb.WriteRequest
for entry := range entries {
update := &spb.WriteRequest_Update{
EdgeKind: entry.EdgeKind,
Target: entry.Target,
FactName: entry.FactName,
FactValue: entry.FactValue,
}
if req != nil && (!compare.VNamesEqual(req.Source, entry.Source) || len(req.Update) >= maxSize) {
ch <- req
req = nil
}
if req == nil {
req = &spb.WriteRequest{
Source: entry.Source,
Update: []*spb.WriteRequest_Update{update},
}
} else {
req.Update = append(req.Update, update)
}
}
if req != nil {
ch <- req
}
}()
return ch
}
// Write implements part of the graphstore.Service interface.
func (db *DB) Write(ctx context.Context, req *spb.WriteRequest) error {
if req.GetSource() == nil {
return fmt.Errorf("missing Source in WriteRequest: {%v}", req)
}
tx, err := db.Begin()
if err != nil {
return fmt.Errorf("sql transaction begin error: %v", err)
}
writeStmt := tx.Stmt(db.writeStmt)
defer func() {
if err := writeStmt.Close(); err != nil {
log.Printf("error closing SQL Stmt: %v", err)
}
}()
for _, update := range req.Update {
if update.Target == nil {
update.Target = emptyVName
}
_, err := writeStmt.Exec(
req.Source.Signature, req.Source.Corpus, req.Source.Root, req.Source.Path, req.Source.Language,
update.EdgeKind,
update.FactName,
update.Target.Signature, update.Target.Corpus, update.Target.Root, update.Target.Path, update.Target.Language,
update.FactValue)
if err != nil {
tx.Rollback()
return fmt.Errorf("sql insertion error: %v", err)
}
}
if err := tx.Commit(); err != nil {
return fmt.Errorf("sql commit error: %v", err)
}
return nil
}
func writeEdges(ctx context.Context, t table.Proto, edges <-chan *spb.Entry) error {
tempDir, err := ioutil.TempDir("", "reverse.edges")
if err != nil {
return fmt.Errorf("failed to create temporary directory: %v", err)
}
defer func() {
drainEntries(edges) // ensure channel is drained on errors
log.Println("Removing temporary edges table", tempDir)
if err := os.RemoveAll(tempDir); err != nil {
log.Printf("Failed to remove temporary directory %q: %v", tempDir, err)
}
}()
gs, err := leveldb.OpenGraphStore(tempDir, nil)
if err != nil {
return fmt.Errorf("failed to create temporary GraphStore: %v", err)
}
defer gs.Close(ctx)
log.Println("Writing temporary reverse edges table")
var writeReq *spb.WriteRequest
for e := range edges {
if writeReq != nil && !compare.VNamesEqual(e.Source, writeReq.Source) {
if err := writeWithReverses(ctx, gs, writeReq); err != nil {
return err
}
writeReq = nil
}
if writeReq == nil {
writeReq = &spb.WriteRequest{Source: e.Source}
}
writeReq.Update = append(writeReq.Update, &spb.WriteRequest_Update{
Target: e.Target,
EdgeKind: e.EdgeKind,
FactName: e.FactName,
FactValue: e.FactValue,
})
}
if writeReq != nil {
if err := writeWithReverses(ctx, gs, writeReq); err != nil {
return err
}
}
return writeEdgePages(ctx, t, gs)
}
