* Added query.WithPoolID() execute option

Author: asmyasnikov

CHANGELOG.md

@@ -1,3 +1,5 @@
+* Added `query.WithPoolID()` execute option
+
 ## v3.85.1
 * Added `spans.Retry` constructor of `trace.Retry`
 

go.mod

@@ -6,7 +6,7 @@ require (
 	github.com/golang-jwt/jwt/v4 v4.4.1
 	github.com/google/uuid v1.6.0
 	github.com/jonboulle/clockwork v0.3.0
-	github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136
+	github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b
 	golang.org/x/net v0.23.0
 	golang.org/x/sync v0.6.0
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1

go.sum

@@ -65,8 +65,8 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
 github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
 github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
-github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136 h1:MO32/Cba3XpNYWcoz3y13eHZG+RzDHmFPry3ren6BmE=
-github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
+github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b h1:8yiv/W+1xTdifJh1Stkck0gFJjys9kg0/r86Buljuss=
+github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
 go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
 go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
 go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=

internal/query/execute_query.go

@@ -30,6 +30,7 @@ type executeSettings interface {
 	Params() *params.Parameters
 	CallOptions() []grpc.CallOption
 	RetryOpts() []retry.Option
+	PoolID() string
 }
 
 type executeScriptConfig interface {
@@ -56,6 +57,7 @@ func executeQueryScriptRequest(a *allocator.Allocator, q string, cfg executeScri
 		Parameters:    cfg.Params().ToYDB(a),
 		StatsMode:     Ydb_Query.StatsMode(cfg.StatsMode()),
 		ResultsTtl:    durationpb.New(cfg.ResultsTTL()),
+		PoolId:        cfg.PoolID(),
 	}
 
 	return request, cfg.CallOptions()
@@ -74,6 +76,7 @@ func executeQueryRequest(a *allocator.Allocator, sessionID, q string, cfg execut
 	request.Parameters = cfg.Params().ToYDB(a)
 	request.StatsMode = Ydb_Query.StatsMode(cfg.StatsMode())
 	request.ConcurrentResultSets = false
+	request.PoolId = cfg.PoolID()
 
 	return request, cfg.CallOptions()
 }

internal/query/options/execute.go

@@ -31,6 +31,7 @@ type (
 		params        params.Parameters
 		execMode      ExecMode
 		statsMode     StatsMode
+		poolID        string
 		statsCallback func(queryStats stats.QueryStats)
 		callOptions   []grpc.CallOption
 		txControl     *tx.Control
@@ -49,6 +50,7 @@ type (
 	// execute options
 	callOptionsOption []grpc.CallOption
 	txCommitOption    struct{}
+	poolID            string
 	parametersOption  params.Parameters
 	txControlOption   tx.Control
 	syntaxOption      = Syntax
@@ -59,6 +61,10 @@ type (
 	execModeOption = ExecMode
 )
 
+func (poolID poolID) applyExecuteOption(s *executeSettings) {
+	s.poolID = string(poolID)
+}
+
 func (s *executeSettings) RetryOpts() []retry.Option {
 	return s.retryOptions
 }
@@ -157,6 +163,10 @@ func (s *executeSettings) StatsMode() StatsMode {
 	return s.statsMode
 }
 
+func (s *executeSettings) PoolID() string {
+	return s.poolID
+}
+
 func (s *executeSettings) Params() *params.Parameters {
 	if len(s.params) == 0 {
 		return nil
@@ -176,12 +186,17 @@ var (
 	_ Execute = StatsMode(0)
 	_ Execute = txCommitOption{}
 	_ Execute = (*txControlOption)(nil)
+	_ Execute = poolID("")
 )
 
 func WithCommit() txCommitOption {
 	return txCommitOption{}
 }
 
+func WithPoolID(id string) poolID {
+	return poolID(id)
+}
+
 func WithExecMode(mode ExecMode) execModeOption {
 	return mode
 }

internal/query/transaction_test.go

@@ -502,9 +502,14 @@ type testExecuteSettings struct {
 	txControl   *query.TransactionControl
 	syntax      options.Syntax
 	params      *params.Parameters
+	poolID      string
 	callOptions []grpc.CallOption
 }
 
+func (s testExecuteSettings) PoolID() string {
+	return s.poolID
+}
+
 func (s testExecuteSettings) RetryOpts() []retry.Option {
 	return nil
 }
@@ -647,6 +652,19 @@ func TestExecuteSettings(t *testing.T) {
 				params:    nil,
 			},
 		},
+		{
+			name: "WithPoolID",
+			txOpts: []options.Execute{
+				options.WithPoolID("test-pool-id"),
+			},
+			settings: testExecuteSettings{
+				execMode:  options.ExecModeExecute,
+				statsMode: options.StatsModeNone,
+				txControl: query.TxControl(query.WithTxID("")),
+				syntax:    options.SyntaxYQL,
+				poolID:    "test-pool-id",
+			},
+		},
 	} {
 		t.Run(tt.name, func(t *testing.T) {
 			a := allocator.New()
@@ -659,6 +677,7 @@ func TestExecuteSettings(t *testing.T) {
 			require.Equal(t, tt.settings.Syntax(), settings.Syntax())
 			require.Equal(t, tt.settings.ExecMode(), settings.ExecMode())
 			require.Equal(t, tt.settings.StatsMode(), settings.StatsMode())
+			require.Equal(t, tt.settings.PoolID(), settings.PoolID())
 			require.Equal(t, tt.settings.TxControl().ToYDB(a).String(), settings.TxControl().ToYDB(a).String())
 			require.Equal(t, tt.settings.Params().ToYDB(a), settings.Params().ToYDB(a))
 			require.Equal(t, tt.settings.CallOptions(), settings.CallOptions())

query/execute_options.go

@@ -0,0 +1,64 @@
+package query
+
+import (
+	"google.golang.org/grpc"
+
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
+	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
+)
+
+const (
+	SyntaxYQL        = options.SyntaxYQL
+	SyntaxPostgreSQL = options.SyntaxPostgreSQL
+)
+
+const (
+	ExecModeParse    = options.ExecModeParse
+	ExecModeValidate = options.ExecModeValidate
+	ExecModeExplain  = options.ExecModeExplain
+	ExecModeExecute  = options.ExecModeExecute
+)
+
+const (
+	StatsModeBasic   = options.StatsModeBasic
+	StatsModeNone    = options.StatsModeNone
+	StatsModeFull    = options.StatsModeFull
+	StatsModeProfile = options.StatsModeProfile
+)
+
+func WithParameters(parameters *params.Parameters) options.Execute {
+	return options.WithParameters(parameters)
+}
+
+func WithTxControl(txControl *tx.Control) options.Execute {
+	return options.WithTxControl(txControl)
+}
+
+func WithTxSettings(txSettings tx.Settings) options.DoTxOption {
+	return options.WithTxSettings(txSettings)
+}
+
+func WithCommit() options.Execute {
+	return options.WithCommit()
+}
+
+func WithExecMode(mode options.ExecMode) options.Execute {
+	return options.WithExecMode(mode)
+}
+
+func WithSyntax(syntax options.Syntax) options.Execute {
+	return options.WithSyntax(syntax)
+}
+
+func WithStatsMode(mode options.StatsMode, callback func(Stats)) options.Execute {
+	return options.WithStatsMode(mode, callback)
+}
+
+func WithCallOptions(opts ...grpc.CallOption) options.Execute {
+	return options.WithCallOptions(opts...)
+}
+
+func WithPoolID(id string) options.Execute {
+	return options.WithPoolID(id)
+}

query/session.go

@@ -3,11 +3,6 @@ package query
 import (
 	"context"
 
-	"google.golang.org/grpc"
-
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/options"
-	"github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
 	"github.com/ydb-platform/ydb-go-sdk/v3/internal/stats"
 )
 
@@ -25,54 +20,3 @@ type (
 	}
 	Stats = stats.QueryStats
 )
-
-const (
-	SyntaxYQL        = options.SyntaxYQL
-	SyntaxPostgreSQL = options.SyntaxPostgreSQL
-)
-
-const (
-	ExecModeParse    = options.ExecModeParse
-	ExecModeValidate = options.ExecModeValidate
-	ExecModeExplain  = options.ExecModeExplain
-	ExecModeExecute  = options.ExecModeExecute
-)
-
-const (
-	StatsModeBasic   = options.StatsModeBasic
-	StatsModeNone    = options.StatsModeNone
-	StatsModeFull    = options.StatsModeFull
-	StatsModeProfile = options.StatsModeProfile
-)
-
-func WithParameters(parameters *params.Parameters) options.Execute {
-	return options.WithParameters(parameters)
-}
-
-func WithTxControl(txControl *tx.Control) options.Execute {
-	return options.WithTxControl(txControl)
-}
-
-func WithTxSettings(txSettings tx.Settings) options.DoTxOption {
-	return options.WithTxSettings(txSettings)
-}
-
-func WithCommit() options.Execute {
-	return options.WithCommit()
-}
-
-func WithExecMode(mode options.ExecMode) options.Execute {
-	return options.WithExecMode(mode)
-}
-
-func WithSyntax(syntax options.Syntax) options.Execute {
-	return options.WithSyntax(syntax)
-}
-
-func WithStatsMode(mode options.StatsMode, callback func(Stats)) options.Execute {
-	return options.WithStatsMode(mode, callback)
-}
-
-func WithCallOptions(opts ...grpc.CallOption) options.Execute {
-	return options.WithCallOptions(opts...)
-}