feat: add influxv1 auth.

kavinli · kavinli · commit 8f9492de4e69 · 2023-02-02T14:03:33.000+08:00
diff --git a/bulk_load/load.go b/bulk_load/load.go
@@ -40,6 +40,9 @@ type BulkLoad interface {
 
 type LoadRunner struct {
 	DbName                 string
+	User                   string
+	Password               string
+	BasicAuthentication    string
 	Workers                int
 	ItemLimit              int64
 	BatchSize              int
@@ -110,6 +113,9 @@ func (r *LoadRunner) notifyHandler(arg int) (int, error) {
 
 func (r *LoadRunner) Init(defaultBatchSize int) {
 	flag.StringVar(&r.DbName, "db", "benchmark_db", "Database name.")
+	flag.StringVar(&r.User, "user", "", "User name, credentials as query parameters.")
+	flag.StringVar(&r.Password, "password", "", "User password, credentials as query parameters.")
+	flag.StringVar(&r.BasicAuthentication, "basic-authentication", "", "Authenticate with basic authentication. format [user:password].")
 	flag.IntVar(&r.BatchSize, "batch-size", defaultBatchSize, "Batch size (1 line of input = 1 item).")
 	flag.IntVar(&r.Workers, "workers", 1, "Number of parallel requests to make.")
 	flag.Int64Var(&r.ItemLimit, "item-limit", -1, "Number of items to read from stdin before quitting. (1 item per 1 line of input.)")
diff --git a/bulk_query_gen/database_config.go b/bulk_query_gen/database_config.go
@@ -4,4 +4,6 @@ type DatabaseConfig map[string]string
 
 const (
 	DatabaseName = "database-name"
+	UserName = "user-name"
+	Password = "password"
 )
diff --git a/bulk_query_gen/influxdb/influx_bareagg_common.go b/bulk_query_gen/influxdb/influx_bareagg_common.go
@@ -18,7 +18,7 @@ func NewInfluxBareAggregateQuery(agg Aggregate, lang Language, dbConfig bulkQuer
 	}
 
 	return &InfluxBareAggregateQuery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 		aggregate:    agg,
 	}
 }
diff --git a/bulk_query_gen/influxdb/influx_common.go b/bulk_query_gen/influxdb/influx_common.go
@@ -45,13 +45,17 @@ type InfluxCommon struct {
 	bulkQuerygen.CommonParams
 	language     Language
 	DatabaseName string
+	UserName     string
+	Password     string
 }
 
-func newInfluxCommon(lang Language, dbName string, interval bulkQuerygen.TimeInterval, scaleVar int) *InfluxCommon {
+func newInfluxCommon(lang Language, dbName string, userName string, password string, interval bulkQuerygen.TimeInterval, scaleVar int) *InfluxCommon {
 	return &InfluxCommon{
 		CommonParams: *bulkQuerygen.NewCommonParams(interval, scaleVar),
 		language:     lang,
-		DatabaseName: dbName}
+		DatabaseName: dbName,
+		UserName: userName,
+		Password: password}
 }
 
 // getHttpQuery gets the right kind of http request based on the language being used
@@ -64,6 +68,10 @@ func (d *InfluxCommon) getHttpQuery(humanLabel, intervalStart, query string, q *
 		getValues := url.Values{}
 		getValues.Set("db", d.DatabaseName)
 		getValues.Set("q", query)
+		if len(d.UserName) != 0 && len(d.Password) != 0 {
+			getValues.Set("u", d.UserName)
+			getValues.Set("p", d.Password)
+		}
 		q.Method = []byte("GET")
 		q.Path = []byte(fmt.Sprintf("/query?%s", getValues.Encode()))
 		q.Body = nil
diff --git a/bulk_query_gen/influxdb/influx_dashboard_common.go b/bulk_query_gen/influxdb/influx_dashboard_common.go
@@ -25,7 +25,7 @@ func newInfluxDashboard(lang Language, dbConfig bulkQuerygen.DatabaseConfig, int
 		clustersCount = 1
 	}
 	return &InfluxDashboard{
-		InfluxCommon:  *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], interval, scaleVar),
+		InfluxCommon:  *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], interval, scaleVar),
 		ClustersCount: clustersCount,
 		TimeWindow:    bulkQuerygen.TimeWindow{interval.Start, duration},
 	}
diff --git a/bulk_query_gen/influxdb/influx_devops_common.go b/bulk_query_gen/influxdb/influx_devops_common.go
@@ -21,7 +21,7 @@ func newInfluxDevopsCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfig,
 	}
 
 	return &InfluxDevops{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], interval, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], interval, scaleVar),
 	}
 }
 
diff --git a/bulk_query_gen/influxdb/influx_group_window_transpose_common.go b/bulk_query_gen/influxdb/influx_group_window_transpose_common.go
@@ -20,7 +20,7 @@ func NewInfluxGroupWindowTransposeQuery(agg Aggregate, card Cardinality, lang La
 	}
 
 	return &InfluxGroupWindowTransposeQuery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 		aggregate:    agg,
 		interval:     queryInterval,
 		cardinality:  card,
diff --git a/bulk_query_gen/influxdb/influx_groupagg_common.go b/bulk_query_gen/influxdb/influx_groupagg_common.go
@@ -18,7 +18,7 @@ func NewInfluxGroupAggregateQuery(agg Aggregate, lang Language, dbConfig bulkQue
 	}
 
 	return &InfluxGroupAggregateQuery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 		aggregate:    agg,
 	}
 }
diff --git a/bulk_query_gen/influxdb/influx_iot_common.go b/bulk_query_gen/influxdb/influx_iot_common.go
@@ -22,7 +22,7 @@ func NewInfluxIotCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfig, que
 	}
 
 	return &InfluxIot{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 	}
 }
 
diff --git a/bulk_query_gen/influxdb/influx_metaquery_common.go b/bulk_query_gen/influxdb/influx_metaquery_common.go
@@ -18,7 +18,7 @@ func NewInfluxMetaqueryCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfi
 	}
 
 	return &InfluxMetaquery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 	}
 }
 
diff --git a/bulk_query_gen/influxdb/influx_multi_measurement_common.go b/bulk_query_gen/influxdb/influx_multi_measurement_common.go
@@ -19,7 +19,7 @@ func InfluxMultiMeasurementCommon(lang Language, dbConfig bulkQuerygen.DatabaseC
 	}
 
 	return &InfluxMultiMeasurement{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 	}
 }
 
diff --git a/bulk_query_gen/influxdb/influx_ungroupedagg_common.go b/bulk_query_gen/influxdb/influx_ungroupedagg_common.go
@@ -18,7 +18,7 @@ func NewInfluxUngroupedAggregateQuery(agg Aggregate, lang Language, dbConfig bul
 	}
 
 	return &InfluxUngroupedAggregateQuery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 		aggregate:    agg,
 	}
 }
diff --git a/bulk_query_gen/influxdb/influx_windowagg_common.go b/bulk_query_gen/influxdb/influx_windowagg_common.go
@@ -19,7 +19,7 @@ func NewInfluxWindowAggregateQuery(agg Aggregate, lang Language, dbConfig bulkQu
 	}
 
 	return &InfluxWindowAggregateQuery{
-		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),
+		InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),
 		aggregate:    agg,
 		interval:     queryInterval,
 	}
diff --git a/cmd/bulk_load_influx/http_writer.go b/cmd/bulk_load_influx/http_writer.go
@@ -7,6 +7,7 @@ import (
 	"fmt"
 	"log"
 	"time"
+	"encoding/base64"
 
 	"github.com/valyala/fasthttp"
 )
@@ -32,6 +33,10 @@ type HTTPWriterConfig struct {
 	// Name of the target database into which points will be written.
 	Database string
 
+	User string
+	Password string
+	BasicAuthentication string 
+
 	// Id of the target bucket into which points will be written. (InfluxDB v2)
 	BucketId string
 
@@ -86,6 +91,9 @@ func (w *HTTPWriter) WriteLineProtocol(body []byte, isGzip bool) (int64, error)
 	if w.c.AuthToken != "" {
 		req.Header.Add("Authorization", fmt.Sprintf("Token %s", w.c.AuthToken))
 	}
+	if w.c.BasicAuthentication != "" {
+		req.Header.Add("Authorization", fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(w.c.BasicAuthentication))))
+	}
 	req.SetBody(body)
 
 	resp := fasthttp.AcquireResponse()
diff --git a/cmd/bulk_load_influx/main.go b/cmd/bulk_load_influx/main.go
@@ -290,10 +290,16 @@ func (l *InfluxBulkLoad) PrepareProcess(i int) {
 			DebugInfo:      fmt.Sprintf("worker #%d, dest url: %s", i, l.configs[i].url),
 			Host:           l.configs[i].url,
 			Database:       bulk_load.Runner.DbName,
+			User:           bulk_load.Runner.User,
+			Password:       bulk_load.Runner.Password,
+			BasicAuthentication: bulk_load.Runner.BasicAuthentication,
 			BackingOffChan: l.configs[i].backingOffChan,
 			BackingOffDone: l.configs[i].backingOffDone,
 		}
 		url = c.Host + "/write?consistency=" + l.consistency + "&db=" + neturl.QueryEscape(c.Database)
+		if len(c.User) != 0 && len(c.Password) != 0 {
+			url = url + "&u=" + neturl.QueryEscape(c.User) + "&p=" + neturl.QueryEscape(c.Password)
+		}
 	}
 	l.configs[i].writer = NewHTTPWriter(*c, url)
 }
diff --git a/cmd/bulk_query_gen/main.go b/cmd/bulk_query_gen/main.go
@@ -431,6 +431,8 @@ var (
 	queryCount int
 
 	dbName string // TODO(rw): make this a map[string]string -> DatabaseConfig
+	userName string 
+	password string
 
 	timestampStartStr string
 	timestampEndStr   string
@@ -475,6 +477,8 @@ func init() {
 	flag.IntVar(&scaleVar, "scale-var", 1, "Scaling variable (must be the equal to the scale-var used for data generation).")
 	flag.IntVar(&queryCount, "queries", 1000, "Number of queries to generate.")
 	flag.StringVar(&dbName, "db", "benchmark_db", "Database to use (ignored for ElasticSearch).")
+	flag.StringVar(&userName, "user", "", "User name, credentials as query parameters.")
+	flag.StringVar(&password, "password", "", "User password, credentials as query parameters.")
 
 	flag.StringVar(&timestampStartStr, "timestamp-start", common.DefaultDateTimeStart, "Beginning timestamp (RFC3339).")
 	flag.StringVar(&timestampEndStr, "timestamp-end", common.DefaultDateTimeEnd, "Ending timestamp (RFC3339).")
@@ -587,6 +591,8 @@ func main() {
 
 	dbConfig := bulkQueryGen.DatabaseConfig{
 		bulkQueryGen.DatabaseName: dbName,
+		bulkQueryGen.UserName: userName,
+		bulkQueryGen.Password: password,
 	}
 
 	// Make the query generator:

Original file line number	Diff line number	Diff line change
`@@ -4,4 +4,6 @@ type DatabaseConfig map[string]string`
`4`	`4`
`5`	`5`	`const (`
`6`	`6`	`DatabaseName = "database-name"`
	`7`	`+ UserName = "user-name"`
	`8`	`+ Password = "password"`
`7`	`9`	`)`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ func NewInfluxBareAggregateQuery(agg Aggregate, lang Language, dbConfig bulkQuer`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`return &InfluxBareAggregateQuery{`
`21`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),`
	`21`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),`
`22`	`22`	`aggregate: agg,`
`23`	`23`	`}`
`24`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -25,7 +25,7 @@ func newInfluxDashboard(lang Language, dbConfig bulkQuerygen.DatabaseConfig, int`
`25`	`25`	`clustersCount = 1`
`26`	`26`	`}`
`27`	`27`	`return &InfluxDashboard{`
`28`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], interval, scaleVar),`
	`28`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], interval, scaleVar),`
`29`	`29`	`ClustersCount: clustersCount,`
`30`	`30`	`TimeWindow: bulkQuerygen.TimeWindow{interval.Start, duration},`
`31`	`31`	`}`
Original file line number	Diff line number	Diff line change
`@@ -21,7 +21,7 @@ func newInfluxDevopsCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfig,`
`21`	`21`	`}`
`22`	`22`
`23`	`23`	`return &InfluxDevops{`
`24`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], interval, scaleVar),`
	`24`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], interval, scaleVar),`
`25`	`25`	`}`
`26`	`26`	`}`
`27`	`27`
Original file line number	Diff line number	Diff line change
`@@ -20,7 +20,7 @@ func NewInfluxGroupWindowTransposeQuery(agg Aggregate, card Cardinality, lang La`
`20`	`20`	`}`
`21`	`21`
`22`	`22`	`return &InfluxGroupWindowTransposeQuery{`
`23`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),`
	`23`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),`
`24`	`24`	`aggregate: agg,`
`25`	`25`	`interval: queryInterval,`
`26`	`26`	`cardinality: card,`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ func NewInfluxGroupAggregateQuery(agg Aggregate, lang Language, dbConfig bulkQue`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`return &InfluxGroupAggregateQuery{`
`21`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),`
	`21`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),`
`22`	`22`	`aggregate: agg,`
`23`	`23`	`}`
`24`	`24`	`}`
Original file line number	Diff line number	Diff line change
`@@ -22,7 +22,7 @@ func NewInfluxIotCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfig, que`
`22`	`22`	`}`
`23`	`23`
`24`	`24`	`return &InfluxIot{`
`25`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),`
	`25`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),`
`26`	`26`	`}`
`27`	`27`	`}`
`28`	`28`
Original file line number	Diff line number	Diff line change
`@@ -18,7 +18,7 @@ func NewInfluxMetaqueryCommon(lang Language, dbConfig bulkQuerygen.DatabaseConfi`
`18`	`18`	`}`
`19`	`19`
`20`	`20`	`return &InfluxMetaquery{`
`21`		`- InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], queriesFullRange, scaleVar),`
	`21`	`+ InfluxCommon: *newInfluxCommon(lang, dbConfig[bulkQuerygen.DatabaseName], dbConfig[bulkQuerygen.UserName], dbConfig[bulkQuerygen.Password], queriesFullRange, scaleVar),`
`22`	`22`	`}`
`23`	`23`	`}`
`24`	`24`