fix: support same config options in connection string and env vars

InfluxCommunity · Aug 31, 2023 · 1e4b43a · 1e4b43a
1 parent 50fb45b
commit 1e4b43a
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 45 deletions.
diff --git a/influxdb3/client.go b/influxdb3/client.go
@@ -113,12 +113,14 @@ func New(config ClientConfig) (*Client, error) {
 // NewFromConnectionString creates new Client from the specified connection string.
 // Parameters:
 //   - connectionString: connection string in URL format.
+//
 // Supported query parameters:
 //   - token (required)
 //   - org
 //   - database
 //   - precision
 //   - gzipThreshold
+//
 // Example: https://us-east-1-1.aws.cloud2.influxdata.com/?token=my-token&database=my-database
 func NewFromConnectionString(connectionString string) (*Client, error) {
 	cfg := ClientConfig{}
@@ -135,6 +137,8 @@ func NewFromConnectionString(connectionString string) (*Client, error) {
 //   - INFLUX_TOKEN (required)
 //   - INFLUX_ORG
 //   - INFLUX_DATABASE
+//   - INFLUX_PRECISION
+//   - INFLUX_GZIP_THRESHOLD
 func NewFromEnv() (*Client, error) {
 	cfg := ClientConfig{}
 	err := cfg.env()

diff --git a/influxdb3/client_test.go b/influxdb3/client_test.go
@@ -148,7 +148,8 @@ func TestNewFromConnectionString(t *testing.T) {
 				Organization: "my-org",
 				Database:     "my-db",
 				WriteOptions: &WriteOptions{
-					Precision: lineprotocol.Millisecond,
+					Precision:     lineprotocol.Millisecond,
+					GzipThreshold: 1000, // default
 				},
 			},
 		},
@@ -224,12 +225,35 @@ func TestNewFromEnv(t *testing.T) {
 				WriteOptions: &DefaultWriteOptions,
 			},
 		},
+		{
+			name: "with write options",
+			vars: map[string]string{
+				"INFLUX_HOST":           "http://host:8086",
+				"INFLUX_TOKEN":          "abc",
+				"INFLUX_ORG":            "my-org",
+				"INFLUX_DATABASE":       "my-db",
+				"INFLUX_PRECISION":      "ms",
+				"INFLUX_GZIP_THRESHOLD": "64",
+			},
+			cfg: &ClientConfig{
+				Host:         "http://host:8086",
+				Token:        "abc",
+				Organization: "my-org",
+				Database:     "my-db",
+				WriteOptions: &WriteOptions{
+					Precision:     lineprotocol.Millisecond,
+					GzipThreshold: 64,
+				},
+			},
+		},
 	}
 	clearEnv := func() {
-		os.Setenv(envInfluxHost, "")
-		os.Setenv(envInfluxToken, "")
-		os.Setenv(envInfluxOrg, "")
-		os.Setenv(envInfluxDatabase, "")
+		os.Unsetenv(envInfluxHost)
+		os.Unsetenv(envInfluxToken)
+		os.Unsetenv(envInfluxOrg)
+		os.Unsetenv(envInfluxDatabase)
+		os.Unsetenv(envInfluxPrecision)
+		os.Unsetenv(envInfluxGzipThreshold)
 	}
 	setEnv := func(vars map[string]string) {
 		for k, v := range vars {

diff --git a/influxdb3/config.go b/influxdb3/config.go
@@ -33,10 +33,12 @@ import (
 )
 
 const (
-	envInfluxHost     = "INFLUX_HOST"
-	envInfluxToken    = "INFLUX_TOKEN"
-	envInfluxOrg      = "INFLUX_ORG"
-	envInfluxDatabase = "INFLUX_DATABASE"
+	envInfluxHost          = "INFLUX_HOST"
+	envInfluxToken         = "INFLUX_TOKEN"
+	envInfluxOrg           = "INFLUX_ORG"
+	envInfluxDatabase      = "INFLUX_DATABASE"
+	envInfluxPrecision     = "INFLUX_PRECISION"
+	envInfluxGzipThreshold = "INFLUX_GZIP_THRESHOLD"
 )
 
 // ClientConfig holds the parameters for creating a new client.
@@ -102,61 +104,93 @@ func (c *ClientConfig) parse(connectionString string) error {
 	u.RawQuery = ""
 	c.Host = u.String()
 
-	token, ok := values["token"]
-	if ok {
+	if token, ok := values["token"]; ok {
 		c.Token = token[0]
 	}
-	org, ok := values["org"]
-	if ok {
+	if org, ok := values["org"]; ok {
 		c.Organization = org[0]
 	}
-	database, ok := values["database"]
-	if ok {
+	if database, ok := values["database"]; ok {
 		c.Database = database[0]
 	}
-	var writeOptions *WriteOptions
-	precision, ok := values["precision"]
-	if ok {
-		if writeOptions == nil {
-			writeOptions = &WriteOptions{}
+	if precision, ok := values["precision"]; ok {
+		if err := c.parsePrecision(precision[0]); err != nil {
+			return err
 		}
-		switch precision[0] {
-		case "ns":
-			writeOptions.Precision = lineprotocol.Nanosecond
-		case "us":
-			writeOptions.Precision = lineprotocol.Microsecond
-		case "ms":
-			writeOptions.Precision = lineprotocol.Millisecond
-		case "s":
-			writeOptions.Precision = lineprotocol.Second
-		default:
-			return fmt.Errorf("unsupported precision %s", precision[0])
+	}
+	if gzipThreshold, ok := values["gzipThreshold"]; ok {
+		if err := c.parseGzipThreshold(gzipThreshold[0]); err != nil {
+			return err
 		}
 	}
-	gzipThreshold, ok := values["gzipThreshold"]
-	if ok {
-		if writeOptions == nil {
-			writeOptions = &WriteOptions{}
+
+	return nil
+}
+
+// env initializes the client config from environment variables.
+func (c *ClientConfig) env() error {
+	if host, ok := os.LookupEnv(envInfluxHost); ok {
+		c.Host = host
+	}
+	if token, ok := os.LookupEnv(envInfluxToken); ok {
+		c.Token = token
+	}
+	if org, ok := os.LookupEnv(envInfluxOrg); ok {
+		c.Organization = org
+	}
+	if database, ok := os.LookupEnv(envInfluxDatabase); ok {
+		c.Database = database
+	}
+	if precision, ok := os.LookupEnv(envInfluxPrecision); ok {
+		if err := c.parsePrecision(precision); err != nil {
+			return err
 		}
-		writeOptions.GzipThreshold, err = strconv.Atoi(gzipThreshold[0])
-		if err != nil {
+	}
+	if gzipThreshold, ok := os.LookupEnv(envInfluxGzipThreshold); ok {
+		if err := c.parseGzipThreshold(gzipThreshold); err != nil {
 			return err
 		}
 	}
 
-	if writeOptions != nil {
-		c.WriteOptions = writeOptions
+	return nil
+}
+
+// parsePrecision parses and sets precision
+func (c *ClientConfig) parsePrecision(precision string) error {
+	if c.WriteOptions == nil {
+		options := DefaultWriteOptions
+		c.WriteOptions = &options
+	}
+
+	switch precision {
+	case "ns":
+		c.WriteOptions.Precision = lineprotocol.Nanosecond
+	case "us":
+		c.WriteOptions.Precision = lineprotocol.Microsecond
+	case "ms":
+		c.WriteOptions.Precision = lineprotocol.Millisecond
+	case "s":
+		c.WriteOptions.Precision = lineprotocol.Second
+	default:
+		return fmt.Errorf("unsupported precision '%s'", precision)
 	}
 
 	return nil
 }
 
-// env initializes the client config from environment variables.
-func (c *ClientConfig) env() error {
-	c.Host = os.Getenv(envInfluxHost)
-	c.Token = os.Getenv(envInfluxToken)
-	c.Organization = os.Getenv(envInfluxOrg)
-	c.Database = os.Getenv(envInfluxDatabase)
+// parseGzipThreshold parses and sets gzip threshold
+func (c *ClientConfig) parseGzipThreshold(threshold string) error {
+	if c.WriteOptions == nil {
+		options := DefaultWriteOptions
+		c.WriteOptions = &options
+	}
+
+	value, err := strconv.Atoi(threshold)
+	if err != nil {
+		return err
+	}
+
+	c.WriteOptions.GzipThreshold = value
 
 	return nil
 }