Skip to content

Commit

Permalink
Merge pull request #1521 from ydb-platform/pool-id
Browse files Browse the repository at this point in the history
* Added `query.WithPoolID()` execute option
  • Loading branch information
asmyasnikov authored Oct 22, 2024
2 parents b958288 + 3e21558 commit 4adacff
Show file tree
Hide file tree
Showing 13 changed files with 246 additions and 244 deletions.
2 changes: 2 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
* Added `query.WithPoolID()` execute option

## v3.85.1
* Added `spans.Retry` constructor of `trace.Retry`

Expand Down
2 changes: 1 addition & 1 deletion examples/go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -51,7 +51,7 @@ require (
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
github.com/syndtr/goleveldb v1.0.0 // indirect
github.com/yandex-cloud/go-genproto v0.0.0-20220815090733-4c139c0154e2 // indirect
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136 // indirect
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b // indirect
github.com/ydb-platform/ydb-go-yc-metadata v0.6.1 // indirect
golang.org/x/crypto v0.24.0 // indirect
golang.org/x/mod v0.17.0 // indirect
Expand Down
6 changes: 2 additions & 4 deletions examples/go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -1840,10 +1840,8 @@ github.com/ydb-platform/gorm-driver v0.1.3 h1:uewwScbRuCixNPC0LF7gDKvWcB13/iLj76
github.com/ydb-platform/gorm-driver v0.1.3/go.mod h1:49cSoG5J18muQTiKj4StL2dHs1/dB94OitnHOvetK24=
github.com/ydb-platform/xorm v0.0.3 h1:MXk42lANB6r/MMLg/XdJfyXJycGUDlCeLiMlLGDKVPw=
github.com/ydb-platform/xorm v0.0.3/go.mod h1:hFsU7EUF0o3S+l5c0eyP2yPVjJ0d4gsFdqCsyazzwBc=
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240528144234-5d5a685e41f7 h1:nL8XwD6fSst7xFUirkaWJmE7kM0CdWRYgu6+YQer1d4=
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240528144234-5d5a685e41f7/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
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=
github.com/ydb-platform/ydb-go-sdk-auth-environ v0.3.0 h1:JxSvw+Moont8qCmibP2MjSEIHfkWJLkw0fHZemAk+d0=
github.com/ydb-platform/ydb-go-sdk-auth-environ v0.3.0/go.mod h1:YzCPoNrTbrXZg9bO2YkbjI6eQLkaRIE9Bq8ponu0g8A=
github.com/ydb-platform/ydb-go-sdk-prometheus/v2 v2.0.1 h1:Lsir3AC2VQOTlp8UjZY9zQdCVfWvBNHT3hZn+jSGoo0=
Expand Down
2 changes: 1 addition & 1 deletion go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
4 changes: 2 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
3 changes: 3 additions & 0 deletions internal/query/execute_query.go
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ type executeSettings interface {
Params() *params.Parameters
CallOptions() []grpc.CallOption
RetryOpts() []retry.Option
PoolID() string
}

type executeScriptConfig interface {
Expand All @@ -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()
Expand All @@ -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()
}
Expand Down
15 changes: 15 additions & 0 deletions internal/query/options/execute.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -49,6 +50,7 @@ type (
// execute options
callOptionsOption []grpc.CallOption
txCommitOption struct{}
poolID string
parametersOption params.Parameters
txControlOption tx.Control
syntaxOption = Syntax
Expand All @@ -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
}
Expand Down Expand Up @@ -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
Expand All @@ -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
}
Expand Down
153 changes: 153 additions & 0 deletions internal/query/options/execute_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
package options

import (
"testing"

"github.com/stretchr/testify/require"
"google.golang.org/grpc"

"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
internal "github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
)

type txMock func() *internal.Control

func (tx txMock) txControl() *internal.Control {
if tx == nil {
return internal.NewControl(internal.WithTxID(""))
}

return tx()
}

func TestExecuteSettings(t *testing.T) {
for _, tt := range []struct {
name string
tx txMock
txOpts []Execute
settings executeSettings
}{
{
name: "WithTxID",
tx: func() *internal.Control {
return internal.NewControl(internal.WithTxID("test"))
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("test")),
syntax: SyntaxYQL,
},
},
{
name: "WithStats",
txOpts: []Execute{
WithStatsMode(StatsModeFull, nil),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeFull,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxYQL,
},
},
{
name: "WithExecMode",
txOpts: []Execute{
WithExecMode(ExecModeExplain),
},
settings: executeSettings{
execMode: ExecModeExplain,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxYQL,
},
},
{
name: "WithSyntax",
txOpts: []Execute{
WithSyntax(SyntaxPostgreSQL),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxPostgreSQL,
},
},
{
name: "WithGrpcOptions",
txOpts: []Execute{
WithCallOptions(grpc.CallContentSubtype("test")),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxYQL,
callOptions: []grpc.CallOption{
grpc.CallContentSubtype("test"),
},
},
},
{
name: "WithParams",
txOpts: []Execute{
WithParameters(
params.Builder{}.Param("$a").Text("A").Build(),
),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxYQL,
params: *params.Builder{}.Param("$a").Text("A").Build(),
},
},
{
name: "WithCommitTx",
txOpts: []Execute{
WithCommit(),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID(""), internal.CommitTx()),
syntax: SyntaxYQL,
params: nil,
},
},
{
name: "WithPoolID",
txOpts: []Execute{
WithPoolID("test-pool-id"),
},
settings: executeSettings{
execMode: ExecModeExecute,
statsMode: StatsModeNone,
txControl: internal.NewControl(internal.WithTxID("")),
syntax: SyntaxYQL,
poolID: "test-pool-id",
},
},
} {
t.Run(tt.name, func(t *testing.T) {
a := allocator.New()
settings := ExecuteSettings(
append(
[]Execute{WithTxControl(tt.tx.txControl())},
tt.txOpts...,
)...,
)
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())
})
}
}
Loading

0 comments on commit 4adacff

Please sign in to comment.