diff --git a/CHANGELOG.md b/CHANGELOG.md index b3599e0dc..025f94d81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,5 @@ +* Added `query.WithPoolID()` execute option + ## v3.85.1 * Added `spans.Retry` constructor of `trace.Retry` diff --git a/go.mod b/go.mod index 26b51dee5..f17f69a36 100644 --- a/go.mod +++ b/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-20241002120727-5acc94bcb119 golang.org/x/net v0.23.0 golang.org/x/sync v0.6.0 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 diff --git a/go.sum b/go.sum index 496e4ea15..5df1cd592 100644 --- a/go.sum +++ b/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-20241002120727-5acc94bcb119 h1:ExSVPjuxuGuu91L0cTD2EZnMOr7VIq1vuA2dVYG0+Xc= +github.com/ydb-platform/ydb-go-genproto v0.0.0-20241002120727-5acc94bcb119/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= diff --git a/internal/query/execute_query.go b/internal/query/execute_query.go index 25b223582..969ac5b93 100644 --- a/internal/query/execute_query.go +++ b/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() } diff --git a/internal/query/options/execute.go b/internal/query/options/execute.go index f6e228fb6..4cfe1957a 100644 --- a/internal/query/options/execute.go +++ b/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 } diff --git a/internal/query/transaction_test.go b/internal/query/transaction_test.go index 30e2b2643..f7e4c2327 100644 --- a/internal/query/transaction_test.go +++ b/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()) diff --git a/query/execute_options.go b/query/execute_options.go new file mode 100644 index 000000000..031648c62 --- /dev/null +++ b/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) +} diff --git a/query/session.go b/query/session.go index 6207bbc42..1891c159c 100644 --- a/query/session.go +++ b/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...) -}