add WithConnectionsPerEndpoint and WithLocationPreference options

Author: adwski

client.go

@@ -14,10 +14,6 @@ import (
 	qq "github.com/adwski/ydb-go-query/query"
 )
 
-const (
-	defaultConnectionsPerEndpoint = 2
-)
-
 var (
 	ErrNoInitialNodes           = errors.New("no initial nodes was provided")
 	ErrDBEmpty                  = errors.New("db is empty")
@@ -127,8 +123,8 @@ func newClient(ctx context.Context, cfg *Config, opts ...Option) (*Client, error
 		InitNodes: cfg.InitialNodes,
 		DB:        cfg.DB,
 		Balancer: balancer.Config{
-			LocationPreference: []string{"TODO"},
-			ConnsPerEndpoint:   defaultConnectionsPerEndpoint,
+			LocationPreference: cfg.locationPreference,
+			ConnsPerEndpoint:   cfg.connectionsPerEndpoint,
 			IgnoreLocations:    false,
 		},
 		TransportCredentials: cfg.transportCredentials,

config.go

@@ -3,6 +3,7 @@ package ydbgoquery
 import (
 	"context"
 	"errors"
+	"strings"
 	"time"
 
 	"github.com/adwski/ydb-go-query/internal/logger"
@@ -23,9 +24,10 @@ import (
 )
 
 const (
-	defaultSessionCreateTimeout = 3 * time.Second
-	defaultQueryTimeout         = 5 * time.Minute
-	defaultSessionPoolSize      = 10
+	defaultSessionCreateTimeout   = 3 * time.Second
+	defaultQueryTimeout           = 5 * time.Minute
+	defaultSessionPoolSize        = 10
+	defaultConnectionsPerEndpoint = 2
 )
 
 var (
@@ -43,10 +45,14 @@ type (
 		DB           string
 		InitialNodes []string
 
+		locationPreference []string
+
 		poolSize    uint
 		poolReadyHi uint
 		poolReadyLo uint
 
+		connectionsPerEndpoint int
+
 		sessionCreateTimeout time.Duration
 		queryTimeout         time.Duration
 	}
@@ -58,6 +64,7 @@ func (cfg *Config) setDefaults() {
 	cfg.sessionCreateTimeout = defaultSessionCreateTimeout
 	cfg.queryTimeout = defaultQueryTimeout
 	cfg.poolSize = defaultSessionPoolSize
+	cfg.connectionsPerEndpoint = defaultConnectionsPerEndpoint
 	cfg.transportCredentials = transportCreds.Insecure()
 	cfg.txSettings = txsettings.SerializableReadWrite()
 }
@@ -123,6 +130,27 @@ func WithSessionPoolReadyThresholds(high, low uint) Option {
 	}
 }
 
+func WithLocationPreference(pref string) Option {
+	return func(ctx context.Context, cfg *Config) error {
+		cfg.locationPreference = strings.Split(pref, ",")
+		for idx := range cfg.locationPreference {
+			cfg.locationPreference[idx] = strings.TrimSpace(cfg.locationPreference[idx])
+		}
+
+		return nil
+	}
+}
+
+func WithConnectionsPerEndpoint(connections int) Option {
+	return func(ctx context.Context, cfg *Config) error {
+		if connections > 0 {
+			cfg.connectionsPerEndpoint = connections
+		}
+
+		return nil
+	}
+}
+
 func withTransportSecurity(credentials credentials.TransportCredentials) Option {
 	return func(ctx context.Context, cfg *Config) error {
 		cfg.transportCredentials = credentials