Skip to content

Commit 29c7247

Browse files
hengfenglihengfengli
authored
feat(spanner): add the support of optimizer statistics package (#2717)
1 parent 980126a commit 29c7247

5 files changed

Lines changed: 63 additions & 35 deletions

File tree

‎spanner/batch_test.go‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@ package spanner
1818

1919
import (
2020
"context"
21-
"os"
2221
"sync"
2322
"testing"
2423
"time"
@@ -80,8 +79,8 @@ func TestPartitionQuery_QueryOptions(t *testing.T) {
8079
for _, tt := range queryOptionsTestCases() {
8180
t.Run(tt.name, func(t *testing.T) {
8281
if tt.env.Options != nil {
83-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
84-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
82+
unset := setQueryOptionsEnvVars(tt.env.Options)
83+
defer unset()
8584
}
8685

8786
ctx := context.Background()
@@ -113,6 +112,9 @@ func TestPartitionQuery_QueryOptions(t *testing.T) {
113112
if got, want := p.qreq.QueryOptions.OptimizerVersion, tt.want.Options.OptimizerVersion; got != want {
114113
t.Fatalf("Incorrect optimizer version: got %v, want %v", got, want)
115114
}
115+
if got, want := p.qreq.QueryOptions.OptimizerStatisticsPackage, tt.want.Options.OptimizerStatisticsPackage; got != want {
116+
t.Fatalf("Incorrect optimizer statistics package: got %v, want %v", got, want)
117+
}
116118
}
117119
})
118120
}

‎spanner/client.go‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -227,13 +227,17 @@ func allClientOpts(numChannels int, userOpts ...option.ClientOption) []option.Cl
227227
// via application-level configuration. If the environment variables are set,
228228
// this will return the overwritten query options.
229229
func getQueryOptions(opts QueryOptions) QueryOptions {
230+
if opts.Options == nil {
231+
opts.Options = &sppb.ExecuteSqlRequest_QueryOptions{}
232+
}
230233
opv := os.Getenv("SPANNER_OPTIMIZER_VERSION")
231234
if opv != "" {
232-
if opts.Options == nil {
233-
opts.Options = &sppb.ExecuteSqlRequest_QueryOptions{}
234-
}
235235
opts.Options.OptimizerVersion = opv
236236
}
237+
opsp := os.Getenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE")
238+
if opsp != "" {
239+
opts.Options.OptimizerStatisticsPackage = opsp
240+
}
237241
return opts
238242
}
239243

‎spanner/client_test.go‎

Lines changed: 45 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -397,8 +397,8 @@ func TestClient_Single_QueryOptions(t *testing.T) {
397397
for _, tt := range queryOptionsTestCases() {
398398
t.Run(tt.name, func(t *testing.T) {
399399
if tt.env.Options != nil {
400-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
401-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
400+
unset := setQueryOptionsEnvVars(tt.env.Options)
401+
defer unset()
402402
}
403403

404404
ctx := context.Background()
@@ -458,6 +458,9 @@ func checkReqsForQueryOptions(t *testing.T, server InMemSpannerServer, qo QueryO
458458
if got, want := reqQueryOptions.OptimizerVersion, qo.Options.OptimizerVersion; got != want {
459459
t.Fatalf("Optimizer version mismatch, got %v, want %v", got, want)
460460
}
461+
if got, want := reqQueryOptions.OptimizerStatisticsPackage, qo.Options.OptimizerStatisticsPackage; got != want {
462+
t.Fatalf("Optimizer statistics package mismatch, got %v, want %v", got, want)
463+
}
461464
}
462465

463466
func testSingleQuery(t *testing.T, serverError error) error {
@@ -623,8 +626,8 @@ func TestClient_ReadOnlyTransaction_QueryOptions(t *testing.T) {
623626
for _, tt := range queryOptionsTestCases() {
624627
t.Run(tt.name, func(t *testing.T) {
625628
if tt.env.Options != nil {
626-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
627-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
629+
unset := setQueryOptionsEnvVars(tt.env.Options)
630+
defer unset()
628631
}
629632

630633
ctx := context.Background()
@@ -645,6 +648,15 @@ func TestClient_ReadOnlyTransaction_QueryOptions(t *testing.T) {
645648
}
646649
}
647650

651+
func setQueryOptionsEnvVars(opts *sppb.ExecuteSqlRequest_QueryOptions) func() {
652+
os.Setenv("SPANNER_OPTIMIZER_VERSION", opts.OptimizerVersion)
653+
os.Setenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE", opts.OptimizerStatisticsPackage)
654+
return func() {
655+
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
656+
defer os.Setenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE", "")
657+
}
658+
}
659+
648660
func testReadOnlyTransaction(t *testing.T, executionTimes map[string]SimulatedExecutionTime) error {
649661
server, client, teardown := setupMockedTestServer(t)
650662
defer teardown()
@@ -784,8 +796,8 @@ func TestClient_ReadWriteTransaction_Query_QueryOptions(t *testing.T) {
784796
for _, tt := range queryOptionsTestCases() {
785797
t.Run(tt.name, func(t *testing.T) {
786798
if tt.env.Options != nil {
787-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
788-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
799+
unset := setQueryOptionsEnvVars(tt.env.Options)
800+
defer unset()
789801
}
790802

791803
ctx := context.Background()
@@ -813,8 +825,8 @@ func TestClient_ReadWriteTransaction_Update_QueryOptions(t *testing.T) {
813825
for _, tt := range queryOptionsTestCases() {
814826
t.Run(tt.name, func(t *testing.T) {
815827
if tt.env.Options != nil {
816-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
817-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
828+
unset := setQueryOptionsEnvVars(tt.env.Options)
829+
defer unset()
818830
}
819831

820832
ctx := context.Background()
@@ -2240,7 +2252,10 @@ func TestClient_EmulatorWithCredentialsFile(t *testing.T) {
22402252

22412253
func TestBatchReadOnlyTransaction_QueryOptions(t *testing.T) {
22422254
ctx := context.Background()
2243-
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
2255+
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
2256+
OptimizerVersion: "1",
2257+
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
2258+
}}
22442259
_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{QueryOptions: qo})
22452260
defer teardown()
22462261

@@ -2256,7 +2271,10 @@ func TestBatchReadOnlyTransaction_QueryOptions(t *testing.T) {
22562271
}
22572272

22582273
func TestBatchReadOnlyTransactionFromID_QueryOptions(t *testing.T) {
2259-
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
2274+
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
2275+
OptimizerVersion: "1",
2276+
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
2277+
}}
22602278
_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{QueryOptions: qo})
22612279
defer teardown()
22622280

@@ -2276,48 +2294,49 @@ type QueryOptionsTestCase struct {
22762294
}
22772295

22782296
func queryOptionsTestCases() []QueryOptionsTestCase {
2297+
statsPkg := "auto_20191128_14_47_22UTC"
22792298
return []QueryOptionsTestCase{
22802299
{
22812300
"Client level",
2282-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2301+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
22832302
QueryOptions{Options: nil},
22842303
QueryOptions{Options: nil},
2285-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2304+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
22862305
},
22872306
{
22882307
"Environment level",
22892308
QueryOptions{Options: nil},
2290-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2309+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
22912310
QueryOptions{Options: nil},
2292-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2311+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
22932312
},
22942313
{
22952314
"Query level",
22962315
QueryOptions{Options: nil},
22972316
QueryOptions{Options: nil},
2298-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2299-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2317+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
2318+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
23002319
},
23012320
{
23022321
"Environment level has precedence",
2303-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2304-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
2322+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
2323+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
23052324
QueryOptions{Options: nil},
2306-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
2325+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
23072326
},
23082327
{
23092328
"Query level has precedence than client level",
2310-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2329+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
23112330
QueryOptions{Options: nil},
2312-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
2313-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
2331+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
2332+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
23142333
},
23152334
{
23162335
"Query level has highest precedence",
2317-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
2318-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
2319-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
2320-
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
2336+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
2337+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
2338+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
2339+
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
23212340
},
23222341
}
23232342
}

‎spanner/integration_test.go‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -635,7 +635,10 @@ func TestIntegration_SingleUse_WithQueryOptions(t *testing.T) {
635635
t.Fatal(err)
636636
}
637637
}
638-
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
638+
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
639+
OptimizerVersion: "1",
640+
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
641+
}}
639642
got, err := readAll(client.Single().QueryWithOptions(ctx, Statement{
640643
"SELECT SingerId, FirstName, LastName FROM Singers WHERE SingerId IN (@id1, @id3, @id4)",
641644
map[string]interface{}{"id1": int64(1), "id3": int64(3), "id4": int64(4)},

‎spanner/pdml_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -145,8 +145,8 @@ func TestPartitionedUpdate_QueryOptions(t *testing.T) {
145145
for _, tt := range queryOptionsTestCases() {
146146
t.Run(tt.name, func(t *testing.T) {
147147
if tt.env.Options != nil {
148-
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
149-
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
148+
unset := setQueryOptionsEnvVars(tt.env.Options)
149+
defer unset()
150150
}
151151

152152
ctx := context.Background()

0 commit comments

Comments
 (0)