Merge pull request #515 from ydb-platform/disable-server-balancer

* Added `ydb.WithDisableServerBalancer()` database/sql connector option

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `ydb.WithDisableServerBalancer()` database/sql connector option
+
 ## v3.42.3
 * Added `credentials.NewStaticCredentials()` static credentials constructor
 * Changed `internal/credentials.NewStaticCredentials()` signature and behaviour for create grpc connection on each call to auth service 

internal/xsql/connector.go

@@ -51,6 +51,13 @@ func WithTrace(t trace.DatabaseSQL, opts ...trace.DatabaseSQLComposeOption) Conn
 	}
 }
 
+func WithDisableServerBalancer() ConnectorOption {
+	return func(c *Connector) error {
+		c.disableServerBalancer = true
+		return nil
+	}
+}
+
 func Open(d Driver, connection connection, opts ...ConnectorOption) (_ *Connector, err error) {
 	c := &Connector{
 		driver:           d,
@@ -90,10 +97,11 @@ type Connector struct {
 	driver     Driver
 	connection connection
 
-	defaultTxControl     *table.TransactionControl
-	defaultQueryMode     QueryMode
-	defaultDataQueryOpts []options.ExecuteDataQueryOption
-	defaultScanQueryOpts []options.ExecuteScanQueryOption
+	defaultTxControl      *table.TransactionControl
+	defaultQueryMode      QueryMode
+	defaultDataQueryOpts  []options.ExecuteDataQueryOption
+	defaultScanQueryOpts  []options.ExecuteScanQueryOption
+	disableServerBalancer bool
 
 	trace trace.DatabaseSQL
 }
@@ -117,11 +125,10 @@ func (c *Connector) Connect(ctx context.Context) (_ driver.Conn, err error) {
 	defer func() {
 		onDone(err)
 	}()
-	s, err := c.connection.Table().CreateSession( //nolint:staticcheck // SA1019
-		meta.WithAllowFeatures(ctx,
-			metaHeaders.HintSessionBalancer,
-		),
-	)
+	if !c.disableServerBalancer {
+		ctx = meta.WithAllowFeatures(ctx, metaHeaders.HintSessionBalancer)
+	}
+	s, err := c.connection.Table().CreateSession(ctx) //nolint:staticcheck // SA1019
 	if err != nil {
 		return nil, xerrors.WithStackTrace(err)
 	}

sql.go

@@ -117,6 +117,10 @@ func WithDatabaseSQLTrace(t trace.DatabaseSQL, opts ...trace.DatabaseSQLComposeO
 	return xsql.WithTrace(t, opts...)
 }
 
+func WithDisableServerBalancer() ConnectorOption {
+	return xsql.WithDisableServerBalancer()
+}
+
 func Connector(db Connection, opts ...ConnectorOption) (*xsql.Connector, error) {
 	if c, ok := db.(*connection); ok {
 		opts = append(opts, c.databaseSQLOptions...)