Skip to content

Commit 4adacff

Browse files
authored
Merge pull request #1521 from ydb-platform/pool-id
* Added `query.WithPoolID()` execute option
2 parents b958288 + 3e21558 commit 4adacff

File tree

13 files changed

+246
-244
lines changed

13 files changed

+246
-244
lines changed

CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
1+
* Added `query.WithPoolID()` execute option
2+
13
## v3.85.1
24
* Added `spans.Retry` constructor of `trace.Retry`
35

examples/go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,7 +51,7 @@ require (
5151
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
5252
github.com/syndtr/goleveldb v1.0.0 // indirect
5353
github.com/yandex-cloud/go-genproto v0.0.0-20220815090733-4c139c0154e2 // indirect
54-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136 // indirect
54+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b // indirect
5555
github.com/ydb-platform/ydb-go-yc-metadata v0.6.1 // indirect
5656
golang.org/x/crypto v0.24.0 // indirect
5757
golang.org/x/mod v0.17.0 // indirect

examples/go.sum

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1840,10 +1840,8 @@ github.com/ydb-platform/gorm-driver v0.1.3 h1:uewwScbRuCixNPC0LF7gDKvWcB13/iLj76
18401840
github.com/ydb-platform/gorm-driver v0.1.3/go.mod h1:49cSoG5J18muQTiKj4StL2dHs1/dB94OitnHOvetK24=
18411841
github.com/ydb-platform/xorm v0.0.3 h1:MXk42lANB6r/MMLg/XdJfyXJycGUDlCeLiMlLGDKVPw=
18421842
github.com/ydb-platform/xorm v0.0.3/go.mod h1:hFsU7EUF0o3S+l5c0eyP2yPVjJ0d4gsFdqCsyazzwBc=
1843-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240528144234-5d5a685e41f7 h1:nL8XwD6fSst7xFUirkaWJmE7kM0CdWRYgu6+YQer1d4=
1844-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240528144234-5d5a685e41f7/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
1845-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136 h1:MO32/Cba3XpNYWcoz3y13eHZG+RzDHmFPry3ren6BmE=
1846-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
1843+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b h1:8yiv/W+1xTdifJh1Stkck0gFJjys9kg0/r86Buljuss=
1844+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
18471845
github.com/ydb-platform/ydb-go-sdk-auth-environ v0.3.0 h1:JxSvw+Moont8qCmibP2MjSEIHfkWJLkw0fHZemAk+d0=
18481846
github.com/ydb-platform/ydb-go-sdk-auth-environ v0.3.0/go.mod h1:YzCPoNrTbrXZg9bO2YkbjI6eQLkaRIE9Bq8ponu0g8A=
18491847
github.com/ydb-platform/ydb-go-sdk-prometheus/v2 v2.0.1 h1:Lsir3AC2VQOTlp8UjZY9zQdCVfWvBNHT3hZn+jSGoo0=

go.mod

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ require (
66
github.com/golang-jwt/jwt/v4 v4.4.1
77
github.com/google/uuid v1.6.0
88
github.com/jonboulle/clockwork v0.3.0
9-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136
9+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b
1010
golang.org/x/net v0.23.0
1111
golang.org/x/sync v0.6.0
1212
golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1

go.sum

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,8 +65,8 @@ github.com/stretchr/testify v1.5.1/go.mod h1:5W2xD1RspED5o8YsWQXVCued0rvSQ+mT+I5
6565
github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
6666
github.com/stretchr/testify v1.7.1 h1:5TQK59W5E3v0r2duFAb7P95B6hEeOyEnHRa8MjYSMTY=
6767
github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=
68-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136 h1:MO32/Cba3XpNYWcoz3y13eHZG+RzDHmFPry3ren6BmE=
69-
github.com/ydb-platform/ydb-go-genproto v0.0.0-20240920120314-0fed943b0136/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
68+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b h1:8yiv/W+1xTdifJh1Stkck0gFJjys9kg0/r86Buljuss=
69+
github.com/ydb-platform/ydb-go-genproto v0.0.0-20241022174402-dd276c7f197b/go.mod h1:Er+FePu1dNUieD+XTMDduGpQuCPssK5Q4BjF+IIXJ3I=
7070
go.opentelemetry.io/proto/otlp v0.7.0/go.mod h1:PqfVotwruBrMGOCsRd/89rSnXhoiJIqeYNgFYFoEGnI=
7171
go.uber.org/mock v0.4.0 h1:VcM4ZOtdbR4f6VXfiOpwpVJDL6lCReaZ6mw31wqh7KU=
7272
go.uber.org/mock v0.4.0/go.mod h1:a6FSlNadKUHUa9IP5Vyt1zh4fC7uAwxMutEAscFbkZc=

internal/query/execute_query.go

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ type executeSettings interface {
3030
Params() *params.Parameters
3131
CallOptions() []grpc.CallOption
3232
RetryOpts() []retry.Option
33+
PoolID() string
3334
}
3435

3536
type executeScriptConfig interface {
@@ -56,6 +57,7 @@ func executeQueryScriptRequest(a *allocator.Allocator, q string, cfg executeScri
5657
Parameters: cfg.Params().ToYDB(a),
5758
StatsMode: Ydb_Query.StatsMode(cfg.StatsMode()),
5859
ResultsTtl: durationpb.New(cfg.ResultsTTL()),
60+
PoolId: cfg.PoolID(),
5961
}
6062

6163
return request, cfg.CallOptions()
@@ -74,6 +76,7 @@ func executeQueryRequest(a *allocator.Allocator, sessionID, q string, cfg execut
7476
request.Parameters = cfg.Params().ToYDB(a)
7577
request.StatsMode = Ydb_Query.StatsMode(cfg.StatsMode())
7678
request.ConcurrentResultSets = false
79+
request.PoolId = cfg.PoolID()
7780

7881
return request, cfg.CallOptions()
7982
}

internal/query/options/execute.go

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@ type (
3131
params params.Parameters
3232
execMode ExecMode
3333
statsMode StatsMode
34+
poolID string
3435
statsCallback func(queryStats stats.QueryStats)
3536
callOptions []grpc.CallOption
3637
txControl *tx.Control
@@ -49,6 +50,7 @@ type (
4950
// execute options
5051
callOptionsOption []grpc.CallOption
5152
txCommitOption struct{}
53+
poolID string
5254
parametersOption params.Parameters
5355
txControlOption tx.Control
5456
syntaxOption = Syntax
@@ -59,6 +61,10 @@ type (
5961
execModeOption = ExecMode
6062
)
6163

64+
func (poolID poolID) applyExecuteOption(s *executeSettings) {
65+
s.poolID = string(poolID)
66+
}
67+
6268
func (s *executeSettings) RetryOpts() []retry.Option {
6369
return s.retryOptions
6470
}
@@ -157,6 +163,10 @@ func (s *executeSettings) StatsMode() StatsMode {
157163
return s.statsMode
158164
}
159165

166+
func (s *executeSettings) PoolID() string {
167+
return s.poolID
168+
}
169+
160170
func (s *executeSettings) Params() *params.Parameters {
161171
if len(s.params) == 0 {
162172
return nil
@@ -176,12 +186,17 @@ var (
176186
_ Execute = StatsMode(0)
177187
_ Execute = txCommitOption{}
178188
_ Execute = (*txControlOption)(nil)
189+
_ Execute = poolID("")
179190
)
180191

181192
func WithCommit() txCommitOption {
182193
return txCommitOption{}
183194
}
184195

196+
func WithPoolID(id string) poolID {
197+
return poolID(id)
198+
}
199+
185200
func WithExecMode(mode ExecMode) execModeOption {
186201
return mode
187202
}
Lines changed: 153 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,153 @@
1+
package options
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/require"
7+
"google.golang.org/grpc"
8+
9+
"github.com/ydb-platform/ydb-go-sdk/v3/internal/allocator"
10+
"github.com/ydb-platform/ydb-go-sdk/v3/internal/params"
11+
internal "github.com/ydb-platform/ydb-go-sdk/v3/internal/query/tx"
12+
)
13+
14+
type txMock func() *internal.Control
15+
16+
func (tx txMock) txControl() *internal.Control {
17+
if tx == nil {
18+
return internal.NewControl(internal.WithTxID(""))
19+
}
20+
21+
return tx()
22+
}
23+
24+
func TestExecuteSettings(t *testing.T) {
25+
for _, tt := range []struct {
26+
name string
27+
tx txMock
28+
txOpts []Execute
29+
settings executeSettings
30+
}{
31+
{
32+
name: "WithTxID",
33+
tx: func() *internal.Control {
34+
return internal.NewControl(internal.WithTxID("test"))
35+
},
36+
settings: executeSettings{
37+
execMode: ExecModeExecute,
38+
statsMode: StatsModeNone,
39+
txControl: internal.NewControl(internal.WithTxID("test")),
40+
syntax: SyntaxYQL,
41+
},
42+
},
43+
{
44+
name: "WithStats",
45+
txOpts: []Execute{
46+
WithStatsMode(StatsModeFull, nil),
47+
},
48+
settings: executeSettings{
49+
execMode: ExecModeExecute,
50+
statsMode: StatsModeFull,
51+
txControl: internal.NewControl(internal.WithTxID("")),
52+
syntax: SyntaxYQL,
53+
},
54+
},
55+
{
56+
name: "WithExecMode",
57+
txOpts: []Execute{
58+
WithExecMode(ExecModeExplain),
59+
},
60+
settings: executeSettings{
61+
execMode: ExecModeExplain,
62+
statsMode: StatsModeNone,
63+
txControl: internal.NewControl(internal.WithTxID("")),
64+
syntax: SyntaxYQL,
65+
},
66+
},
67+
{
68+
name: "WithSyntax",
69+
txOpts: []Execute{
70+
WithSyntax(SyntaxPostgreSQL),
71+
},
72+
settings: executeSettings{
73+
execMode: ExecModeExecute,
74+
statsMode: StatsModeNone,
75+
txControl: internal.NewControl(internal.WithTxID("")),
76+
syntax: SyntaxPostgreSQL,
77+
},
78+
},
79+
{
80+
name: "WithGrpcOptions",
81+
txOpts: []Execute{
82+
WithCallOptions(grpc.CallContentSubtype("test")),
83+
},
84+
settings: executeSettings{
85+
execMode: ExecModeExecute,
86+
statsMode: StatsModeNone,
87+
txControl: internal.NewControl(internal.WithTxID("")),
88+
syntax: SyntaxYQL,
89+
callOptions: []grpc.CallOption{
90+
grpc.CallContentSubtype("test"),
91+
},
92+
},
93+
},
94+
{
95+
name: "WithParams",
96+
txOpts: []Execute{
97+
WithParameters(
98+
params.Builder{}.Param("$a").Text("A").Build(),
99+
),
100+
},
101+
settings: executeSettings{
102+
execMode: ExecModeExecute,
103+
statsMode: StatsModeNone,
104+
txControl: internal.NewControl(internal.WithTxID("")),
105+
syntax: SyntaxYQL,
106+
params: *params.Builder{}.Param("$a").Text("A").Build(),
107+
},
108+
},
109+
{
110+
name: "WithCommitTx",
111+
txOpts: []Execute{
112+
WithCommit(),
113+
},
114+
settings: executeSettings{
115+
execMode: ExecModeExecute,
116+
statsMode: StatsModeNone,
117+
txControl: internal.NewControl(internal.WithTxID(""), internal.CommitTx()),
118+
syntax: SyntaxYQL,
119+
params: nil,
120+
},
121+
},
122+
{
123+
name: "WithPoolID",
124+
txOpts: []Execute{
125+
WithPoolID("test-pool-id"),
126+
},
127+
settings: executeSettings{
128+
execMode: ExecModeExecute,
129+
statsMode: StatsModeNone,
130+
txControl: internal.NewControl(internal.WithTxID("")),
131+
syntax: SyntaxYQL,
132+
poolID: "test-pool-id",
133+
},
134+
},
135+
} {
136+
t.Run(tt.name, func(t *testing.T) {
137+
a := allocator.New()
138+
settings := ExecuteSettings(
139+
append(
140+
[]Execute{WithTxControl(tt.tx.txControl())},
141+
tt.txOpts...,
142+
)...,
143+
)
144+
require.Equal(t, tt.settings.Syntax(), settings.Syntax())
145+
require.Equal(t, tt.settings.ExecMode(), settings.ExecMode())
146+
require.Equal(t, tt.settings.StatsMode(), settings.StatsMode())
147+
require.Equal(t, tt.settings.PoolID(), settings.PoolID())
148+
require.Equal(t, tt.settings.TxControl().ToYDB(a).String(), settings.TxControl().ToYDB(a).String())
149+
require.Equal(t, tt.settings.Params().ToYDB(a), settings.Params().ToYDB(a))
150+
require.Equal(t, tt.settings.CallOptions(), settings.CallOptions())
151+
})
152+
}
153+
}

0 commit comments

Comments
 (0)