Add basic auth option (#3)

* Add basic auth option

* Use pester.Do with optional basic auth header for POST and GET

* Refactoring

cmd/solrbulk/solrbulk.go

@@ -25,9 +25,11 @@ package main
 
 import (
 	"bufio"
+	b64 "encoding/base64"
 	"flag"
 	"fmt"
 	"io"
+	"net/http"
 	"os"
 	"runtime"
 	"runtime/pprof"
@@ -57,8 +59,21 @@ var (
 	purgePause               = flag.Duration("purge-pause", 2*time.Second, "insert a short pause after purge")
 	updateRequestHandlerName = flag.String("update-request-handler-name", "/update", "where solr.UpdateRequestHandler is mounted on the server, https://is.gd/s0eirv")
 	noFinalCommit            = flag.Bool("no-final-commit", false, "omit final commit")
+	basicAuth                = flag.String("auth", "", "username:password pair for basic auth")
 )
 
+func newGetRequest(url string, options solrbulk.Options) (*http.Request, error) {
+	req, err := http.NewRequest("GET", url, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	if options.BasicAuth != "" {
+		req.Header.Add("Authorization", "Basic "+b64.StdEncoding.EncodeToString([]byte(options.BasicAuth)))
+	}
+	return req, nil
+}
+
 func main() {
 	flag.Parse()
 	if *cpuprofile != "" {
@@ -79,6 +94,7 @@ func main() {
 		Verbose:                  *verbose,
 		UpdateRequestHandlerName: *updateRequestHandlerName,
 		Server:                   *server,
+		BasicAuth:                *basicAuth,
 	}
 	if !strings.HasPrefix(options.Server, "http") {
 		options.Server = fmt.Sprintf("http://%s", options.Server)
@@ -92,7 +108,12 @@ func main() {
 			}
 		)
 		for _, url := range urls {
-			resp, err := pester.Get(url)
+			req, err := newGetRequest(url, options)
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			resp, err := pester.Do(req)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -154,7 +175,12 @@ func main() {
 		queue <- line
 		i++
 		if i%options.CommitSize == 0 {
-			resp, err := pester.Get(commitURL)
+			req, err := newGetRequest(commitURL, options)
+			if err != nil {
+				log.Fatal(err)
+			}
+
+			resp, err := pester.Do(req)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -184,7 +210,13 @@ func main() {
 	if *optimize {
 		hostpath := fmt.Sprintf("%s%s", options.Server, options.UpdateRequestHandlerName)
 		url := fmt.Sprintf("%s?stream.body=<optimize/>", hostpath)
-		resp, err := pester.Get(url)
+
+		req, err := newGetRequest(url, options)
+		if err != nil {
+			log.Fatal(err)
+		}
+
+		resp, err := pester.Do(req)
 		if err != nil {
 			log.Fatal(err)
 		}

worker.go

@@ -23,10 +23,12 @@
 package solrbulk
 
 import (
+	b64 "encoding/base64"
 	"bytes"
 	"fmt"
 	"io"
 	"io/ioutil"
+	"net/http"
 	"strings"
 	"sync"
 	"time"
@@ -45,6 +47,20 @@ type Options struct {
 	Verbose                  bool
 	Server                   string
 	UpdateRequestHandlerName string
+	BasicAuth                string
+}
+
+func newPostRequest(url string, body string, options Options) (*http.Request, error) {
+	req, err := http.NewRequest("POST", url, strings.NewReader(body))
+	if err != nil {
+		return nil, err
+	}
+
+	if options.BasicAuth  != "" {
+		req.Header.Add("Authorization", "Basic "+ b64.StdEncoding.EncodeToString([]byte(options.BasicAuth)))
+	}
+	req.Header.Set("Content-Type", "application/json")
+	return req, nil
 }
 
 // BulkIndex takes a set of documents as strings and indexes them into SOLR.
@@ -59,7 +75,13 @@ func BulkIndex(docs []string, options Options) error {
 		lines = append(lines, doc)
 	}
 	body := fmt.Sprintf("[%s]\n", strings.Join(lines, ","))
-	resp, err := pester.Post(link, "application/json", strings.NewReader(body))
+
+	req, err := newPostRequest(link, body, options)
+	if err != nil {
+		return err
+	}
+
+	resp, err := pester.Do(req)
 	if err != nil {
 		return err
 	}