Generalize exported WriteStream vs WriteStdin (#6)

README.md

@@ -9,7 +9,7 @@ Go package with utilities for interacting with [Badger](https://github.com/dgrap
 
 
 - [Getting Started](#getting-started)
-  - [Stream Stdin to Badger](#stream-stdin-to-badger)
+  - [IO Stream to Badger](#io-stream-to-badger)
     - [Example](#example)
 - [Development](#development)
   - [Dependency Management](#dependency-management)
@@ -20,21 +20,27 @@ Go package with utilities for interacting with [Badger](https://github.com/dgrap
 
 ## Getting Started
 
-### Stream Stdin to Badger
+### IO Stream to Badger
 
-To create a CLI to stream stdin into Badger, use `badgerutils.WriteStdin`. It takes a function `lineToKeyed` as a parameter, which converts a `string` into a struct that implements the `Keyed` interface.
+To stream data Badger, use `badgerutils.WriteStream`.
 
 #### Example
 
+Creates a CLI tool that streams data from stdin.
+
 ```Go
-// examples/writer.go
+// examples/writer_cli.go
 package main
 
 import (
+	"errors"
+	"flag"
 	"fmt"
-	"github.com/Surfline/badgerutils"
 	"log"
+	"os"
 	"strings"
+
+	"github.com/Surfline/badgerutils"
 )
 
 type sampleRecord struct {
@@ -53,7 +59,18 @@ func lineToKeyed(line string) (badgerutils.Keyed, error) {
 }
 
 func main() {
-	if err := badgerutils.WriteStdin(lineToKeyed); err != nil {
+	dir := flag.String("dir", "", "Directory to save DB files")
+	batchSize := flag.Int("batch-size", 1000, "Number of records to write per transaction")
+	flag.Parse()
+
+	if *dir == "" {
+		log.Fatal(errors.New("dir flag is required"))
+	}
+
+	log.Printf("Directory: %v", *dir)
+	log.Printf("Batch Size: %v", *batchSize)
+
+	if err := badgerutils.WriteStream(os.Stdin, *dir, *batchSize, lineToKeyed); err != nil {
 		log.Fatal(err)
 	}
 }

examples/writer.go → examples/writer_cli.go

@@ -1,10 +1,14 @@
 package main
 
 import (
+	"errors"
+	"flag"
 	"fmt"
-	"github.com/Surfline/badgerutils"
 	"log"
+	"os"
 	"strings"
+
+	"github.com/Surfline/badgerutils"
 )
 
 type sampleRecord struct {
@@ -23,7 +27,18 @@ func lineToKeyed(line string) (badgerutils.Keyed, error) {
 }
 
 func main() {
-	if err := badgerutils.WriteStdin(lineToKeyed); err != nil {
+	dir := flag.String("dir", "", "Directory to save DB files")
+	batchSize := flag.Int("batch-size", 1000, "Number of records to write per transaction")
+	flag.Parse()
+
+	if *dir == "" {
+		log.Fatal(errors.New("dir flag is required"))
+	}
+
+	log.Printf("Directory: %v", *dir)
+	log.Printf("Batch Size: %v", *batchSize)
+
+	if err := badgerutils.WriteStream(os.Stdin, *dir, *batchSize, lineToKeyed); err != nil {
 		log.Fatal(err)
 	}
 }

writer.go

@@ -1,23 +1,22 @@
-// Package badgerutils provides functions for interacting with the underlying database.
+// Package badgerutils provides functions for interacting with Badger.
 package badgerutils
 
 import (
 	"bufio"
 	"bytes"
 	"encoding/gob"
-	"errors"
-	"flag"
 	"fmt"
-	"github.com/dgraph-io/badger"
 	"io"
 	"log"
-	"os"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
+
+	"github.com/dgraph-io/badger"
 )
 
+// Keyed interface defines a Key method for defining a key from a struct.
 type Keyed interface {
 	Key() string
 }
@@ -72,10 +71,10 @@ func writeBatch(kvs []keyValue, db *badger.DB, cherr chan error, done func(int32
 	})
 }
 
-func writeInput(reader io.Reader, dir string, batchSize int, lineToKeyed func(string) (Keyed, error)) error {
-	log.Printf("Directory: %v", dir)
-	log.Printf("Batch Size: %v", batchSize)
-
+// WriteStream translates io.Reader stream into key/value pairs that are written into the Badger.
+// lineToKeyed function parameter defines how stdin is translated to a value and how to define a key
+// from that value.
+func WriteStream(reader io.Reader, dir string, batchSize int, lineToKeyed func(string) (Keyed, error)) error {
 	// Open Badger database from directory
 	opts := badger.DefaultOptions
 	opts.Dir = dir
@@ -100,7 +99,7 @@ func writeInput(reader io.Reader, dir string, batchSize int, lineToKeyed func(st
 	kvBatch := make([]keyValue, 0)
 	cherr := make(chan error)
 
-	// Read from stdin and write key/values in batches
+	// Read from stream and write key/values in batches
 	scanner := bufio.NewScanner(reader)
 	for scanner.Scan() {
 		kv, err := stringToKeyValue(scanner.Text(), lineToKeyed)
@@ -121,7 +120,7 @@ func writeInput(reader io.Reader, dir string, batchSize int, lineToKeyed func(st
 		writeBatch(kvBatch, db, cherr, done)
 	}
 
-	// Read and handle errors streaming from stdin
+	// Read and handle errors from stream
 	if err = scanner.Err(); err != nil {
 		return err
 	}
@@ -144,18 +143,3 @@ func writeInput(reader io.Reader, dir string, batchSize int, lineToKeyed func(st
 	log.Printf("Inserted %v records in %v", kvCount.get(), elapsed)
 	return nil
 }
-
-// WriteStdin translates stdin into key/value pairs that are written into the Badger.
-// lineToKeyed function parameter defines how stdin is translated to a value and how to define a key
-// from that value.
-func WriteStdin(lineToKeyed func(string) (Keyed, error)) error {
-	dir := flag.String("dir", "", "Directory to save DB files")
-	batchSize := flag.Int("batch-size", 1000, "Number of records to write per transaction")
-	flag.Parse()
-
-	if *dir == "" {
-		return errors.New("dir flag is required")
-	}
-
-	return writeInput(os.Stdin, *dir, *batchSize, lineToKeyed)
-}

writer_test.go

@@ -4,12 +4,13 @@ import (
 	"bytes"
 	"encoding/gob"
 	"fmt"
-	"github.com/dgraph-io/badger"
-	"github.com/stretchr/testify/require"
 	"io/ioutil"
 	"os"
 	"strings"
 	"testing"
+
+	"github.com/dgraph-io/badger"
+	"github.com/stretchr/testify/require"
 )
 
 type sampleRecord struct {
@@ -76,7 +77,7 @@ func readDB(dir string) ([]sampleRecord, error) {
 	return sampleRecords, nil
 }
 
-func TestWriteInput(t *testing.T) {
+func TestWriteStream(t *testing.T) {
 	dir, err := os.Getwd()
 	require.Nil(t, err)
 	tmpDir, err := ioutil.TempDir(dir, "temp")
@@ -86,7 +87,7 @@ func TestWriteInput(t *testing.T) {
 	reader := strings.NewReader(`field11,field12,field13
 field21,field22,field23
 field31,field32,field33`)
-	err = writeInput(reader, tmpDir, 2, csvToSampleRecord)
+	err = WriteStream(reader, tmpDir, 2, csvToSampleRecord)
 	require.Nil(t, err)
 
 	writtenSampleRecords, err := readDB(tmpDir)