Ejemplo n.º 1
0
func main() {
	log.SetPrefix("write_entries: ")

	flag.Parse()
	if *numWorkers < 1 {
		flagutil.UsageErrorf("Invalid number of --workers %d (must be ≥ 1)", *numWorkers)
	} else if *batchSize < 1 {
		flagutil.UsageErrorf("Invalid --batch_size %d (must be ≥ 1)", *batchSize)
	} else if gs == nil {
		flagutil.UsageError("Missing --graphstore")
	}

	defer gsutil.LogClose(context.Background(), gs)
	gsutil.EnsureGracefulExit(gs)

	if *profCPU != "" {
		f, err := vfs.Create(context.Background(), *profCPU)
		if err != nil {
			log.Fatal(err)
		}
		defer f.Close()
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	writes := graphstore.BatchWrites(stream.ReadEntries(os.Stdin), *batchSize)

	var (
		wg         sync.WaitGroup
		numEntries uint64
	)
	wg.Add(*numWorkers)
	for i := 0; i < *numWorkers; i++ {
		go func() {
			defer wg.Done()
			num, err := writeEntries(context.Background(), gs, writes)
			if err != nil {
				log.Fatal(err)
			}

			atomic.AddUint64(&numEntries, num)
		}()
	}
	wg.Wait()

	log.Printf("Wrote %d entries", numEntries)
}
Ejemplo n.º 2
0
func main() {
	flag.Parse()

	if len(flag.Args()) > 2 || (gs != nil && len(flag.Args()) > 1) {
		fmt.Fprintf(os.Stderr, "ERROR: too many arguments %v\n", flag.Args())
		flag.Usage()
		os.Exit(1)
	}

	if gs != nil {
		defer gsutil.LogClose(context.Background(), gs)
	}

	var in io.ReadCloser = os.Stdin
	if gs == nil && len(flag.Args()) > 0 {
		file, err := vfs.Open(context.Background(), flag.Arg(0))
		if err != nil {
			log.Fatalf("Failed to open input file %q: %v", flag.Arg(0), err)
		}
		defer file.Close()
		in = file
	}

	outIdx := 1
	if gs != nil {
		outIdx = 0
	}

	var out io.WriteCloser = os.Stdout
	if len(flag.Args()) > outIdx {
		file, err := vfs.Create(context.Background(), flag.Arg(outIdx))
		if err != nil {
			log.Fatalf("Failed to create output file %q: %v", flag.Arg(outIdx), err)
		}
		defer file.Close()
		out = file
	}

	var (
		entries      <-chan *spb.Entry
		reverseEdges int
		triples      int
	)

	if gs == nil {
		entries = stream.ReadEntries(in)
	} else {
		ch := make(chan *spb.Entry)
		entries = ch
		go func() {
			defer close(ch)
			if err := gs.Scan(context.Background(), &spb.ScanRequest{}, func(e *spb.Entry) error {
				ch <- e
				return nil
			}); err != nil {
				log.Fatalf("Error scanning graphstore: %v", err)
			}
		}()
	}

	for entry := range entries {
		if schema.EdgeDirection(entry.EdgeKind) == schema.Reverse && !*keepReverseEdges {
			reverseEdges++
			continue
		}

		t, err := toTriple(entry)
		if err != nil {
			log.Fatal(err)
		}
		fmt.Fprintln(out, t)
		triples++
	}

	if !*quiet {
		if !*keepReverseEdges {
			log.Printf("Skipped %d reverse edges", reverseEdges)
		}
		log.Printf("Wrote %d triples", triples)
	}
}