feat: support greptimedb auth (#22)

* feat: support user & password

Signed-off-by: evenyag <realevenyag@gmail.com>

* chore: use && instead of ||

Signed-off-by: evenyag <realevenyag@gmail.com>

---------

Signed-off-by: evenyag <realevenyag@gmail.com>

Author: evenyag

cmd/tsbs_load_greptime/creator.go

@@ -1,6 +1,7 @@
 package main
 
 import (
+	"encoding/base64"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -17,6 +18,15 @@ func (d *dbCreator) Init() {
 	d.daemonURL = daemonURLs[0] // pick first one since it always exists
 }
 
+// addAuthHeader adds Basic authentication header to the request if credentials are provided
+func (d *dbCreator) addAuthHeader(req *http.Request) {
+	if username != "" && password != "" {
+		credentials := username + ":" + password
+		encoded := base64.StdEncoding.EncodeToString([]byte(credentials))
+		req.Header.Set("Authorization", "Basic "+encoded)
+	}
+}
+
 func (d *dbCreator) DBExists(dbName string) bool {
 	dbs, err := d.listDatabases()
 	if err != nil {
@@ -33,7 +43,14 @@ func (d *dbCreator) DBExists(dbName string) bool {
 
 func (d *dbCreator) listDatabases() ([]string, error) {
 	u := fmt.Sprintf("%s/v1/sql?sql=show%%20databases", d.daemonURL)
-	resp, err := http.Get(u)
+	req, err := http.NewRequest("GET", u, nil)
+	if err != nil {
+		return nil, fmt.Errorf("listDatabases error: %s", err.Error())
+	}
+	d.addAuthHeader(req)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
 	if err != nil {
 		return nil, fmt.Errorf("listDatabases error: %s", err.Error())
 	}
@@ -71,10 +88,20 @@ func (d *dbCreator) listDatabases() ([]string, error) {
 
 func (d *dbCreator) RemoveOldDB(dbName string) error {
 	u := fmt.Sprintf("%s/v1/sql?sql=drop+database+%s", d.daemonURL, dbName)
-	resp, err := http.Post(u, "text/plain", nil)
+	req, err := http.NewRequest("POST", u, nil)
 	if err != nil {
 		return fmt.Errorf("drop db error: %s", err.Error())
 	}
+	req.Header.Set("Content-Type", "text/plain")
+	d.addAuthHeader(req)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		return fmt.Errorf("drop db error: %s", err.Error())
+	}
+	defer resp.Body.Close()
+
 	if resp.StatusCode != 200 {
 		return fmt.Errorf("drop db returned non-200 code: %d", resp.StatusCode)
 	}
@@ -84,7 +111,14 @@ func (d *dbCreator) RemoveOldDB(dbName string) error {
 
 func (d *dbCreator) CreateDB(dbName string) error {
 	u := fmt.Sprintf("%s/v1/sql?sql=create%%20database%%20%s", d.daemonURL, dbName)
-	resp, err := http.Get(u)
+	req, err := http.NewRequest("GET", u, nil)
+	if err != nil {
+		return fmt.Errorf("create db error: %s", err.Error())
+	}
+	d.addAuthHeader(req)
+
+	client := &http.Client{}
+	resp, err := client.Do(req)
 	if err != nil {
 		return fmt.Errorf("create db error: %s", err.Error())
 	}

cmd/tsbs_load_greptime/http_writer.go

@@ -4,6 +4,7 @@ package main
 
 import (
 	"bytes"
+	"encoding/base64"
 	"fmt"
 	"net/url"
 	"time"
@@ -38,26 +39,42 @@ type HTTPWriterConfig struct {
 
 	// Debug label for more informative errors.
 	DebugInfo string
+
+	// Username for authentication.
+	Username string
+
+	// Password for authentication.
+	Password string
 }
 
 // HTTPWriter is a Writer that writes to an InfluxDB HTTP server.
 type HTTPWriter struct {
 	client fasthttp.Client
 
-	c   HTTPWriterConfig
-	url []byte
+	c          HTTPWriterConfig
+	url        []byte
+	authHeader []byte
 }
 
 // NewHTTPWriter returns a new HTTPWriter from the supplied HTTPWriterConfig.
 func NewHTTPWriter(c HTTPWriterConfig, consistency string) *HTTPWriter {
-	return &HTTPWriter{
+	w := &HTTPWriter{
 		client: fasthttp.Client{
 			Name: httpClientName,
 		},
 
 		c:   c,
 		url: []byte(c.Host + "/v1/influxdb/write?consistency=" + consistency + "&db=" + url.QueryEscape(c.Database)),
 	}
+
+	// Generate Basic auth header if credentials are provided
+	if c.Username != "" && c.Password != "" {
+		credentials := c.Username + ":" + c.Password
+		encoded := base64.StdEncoding.EncodeToString([]byte(credentials))
+		w.authHeader = []byte("Basic " + encoded)
+	}
+
+	return w
 }
 
 var (
@@ -72,6 +89,9 @@ func (w *HTTPWriter) initializeReq(req *fasthttp.Request, body []byte, isGzip bo
 	if isGzip {
 		req.Header.Add(headerContentEncoding, headerGzip)
 	}
+	if len(w.authHeader) > 0 {
+		req.Header.SetBytesKV([]byte("Authorization"), w.authHeader)
+	}
 	req.SetBody(body)
 }
 

cmd/tsbs_load_greptime/main.go

@@ -30,6 +30,8 @@ var (
 	useGzip           bool
 	doAbortOnExist    bool
 	consistency       string
+	username          string
+	password          string
 )
 
 // Global vars
@@ -75,6 +77,8 @@ func init() {
 	consistency = viper.GetString("consistency")
 	backoff = viper.GetDuration("backoff")
 	useGzip = viper.GetBool("gzip")
+	username = viper.GetString("user")
+	password = viper.GetString("password")
 
 	if _, ok := consistencyChoices[consistency]; !ok {
 		log.Fatalf("invalid consistency settings")

cmd/tsbs_load_greptime/process.go

@@ -26,6 +26,8 @@ func (p *processor) Init(numWorker int, _, _ bool) {
 		DebugInfo: fmt.Sprintf("worker #%d, dest url: %s", numWorker, daemonURL),
 		Host:      daemonURL,
 		Database:  loader.DatabaseName(),
+		Username:  username,
+		Password:  password,
 	}
 	w := NewHTTPWriter(cfg, consistency)
 	p.initWithHTTPWriter(numWorker, w)

pkg/targets/influx/implemented_target.go

@@ -23,6 +23,8 @@ func (t *influxTarget) TargetSpecificFlags(flagPrefix string, flagSet *pflag.Fla
 	flagSet.String(flagPrefix+"consistency", "all", "Write consistency. Must be one of: any, one, quorum, all.")
 	flagSet.Duration(flagPrefix+"backoff", time.Second, "Time to sleep between requests when server indicates backpressure is needed.")
 	flagSet.Bool(flagPrefix+"gzip", true, "Whether to gzip encode requests (default true).")
+	flagSet.String(flagPrefix+"user", "", "Username for authentication.")
+	flagSet.String(flagPrefix+"password", "", "Password for authentication.")
 }
 
 func (t *influxTarget) TargetName() string {