Skip to content

Commit

Permalink
feat(spanner): add the support of optimizer statistics package (#2717)
Browse files Browse the repository at this point in the history
  • Loading branch information
@hengfengli
hengfengli committed Jun 8, 2021
1 parent 980126a commit 29c7247
Show file tree
Hide file tree
Showing 5 changed files with 63 additions and 35 deletions.
8 changes: 5 additions & 3 deletions spanner/batch_test.go
View file
Expand Up @@ -18,7 +18,6 @@ package spanner

import (
"context"
"os"
"sync"
"testing"
"time"
Expand Down Expand Up @@ -80,8 +79,8 @@ func TestPartitionQuery_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

ctx := context.Background()
Expand Down Expand Up @@ -113,6 +112,9 @@ func TestPartitionQuery_QueryOptions(t *testing.T) {
if got, want := p.qreq.QueryOptions.OptimizerVersion, tt.want.Options.OptimizerVersion; got != want {
t.Fatalf("Incorrect optimizer version: got %v, want %v", got, want)
}
if got, want := p.qreq.QueryOptions.OptimizerStatisticsPackage, tt.want.Options.OptimizerStatisticsPackage; got != want {
t.Fatalf("Incorrect optimizer statistics package: got %v, want %v", got, want)
}
}
})
}
Expand Down
10 changes: 7 additions & 3 deletions spanner/client.go
View file
Expand Up @@ -227,13 +227,17 @@ func allClientOpts(numChannels int, userOpts ...option.ClientOption) []option.Cl
// via application-level configuration. If the environment variables are set,
// this will return the overwritten query options.
func getQueryOptions(opts QueryOptions) QueryOptions {
if opts.Options == nil {
opts.Options = &sppb.ExecuteSqlRequest_QueryOptions{}
}
opv := os.Getenv("SPANNER_OPTIMIZER_VERSION")
if opv != "" {
if opts.Options == nil {
opts.Options = &sppb.ExecuteSqlRequest_QueryOptions{}
}
opts.Options.OptimizerVersion = opv
}
opsp := os.Getenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE")
if opsp != "" {
opts.Options.OptimizerStatisticsPackage = opsp
}
return opts
}

Expand Down
71 changes: 45 additions & 26 deletions spanner/client_test.go
View file
Expand Up @@ -397,8 +397,8 @@ func TestClient_Single_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

ctx := context.Background()
Expand Down Expand Up @@ -458,6 +458,9 @@ func checkReqsForQueryOptions(t *testing.T, server InMemSpannerServer, qo QueryO
if got, want := reqQueryOptions.OptimizerVersion, qo.Options.OptimizerVersion; got != want {
t.Fatalf("Optimizer version mismatch, got %v, want %v", got, want)
}
if got, want := reqQueryOptions.OptimizerStatisticsPackage, qo.Options.OptimizerStatisticsPackage; got != want {
t.Fatalf("Optimizer statistics package mismatch, got %v, want %v", got, want)
}
}

func testSingleQuery(t *testing.T, serverError error) error {
Expand Down Expand Up @@ -623,8 +626,8 @@ func TestClient_ReadOnlyTransaction_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

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

func setQueryOptionsEnvVars(opts *sppb.ExecuteSqlRequest_QueryOptions) func() {
os.Setenv("SPANNER_OPTIMIZER_VERSION", opts.OptimizerVersion)
os.Setenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE", opts.OptimizerStatisticsPackage)
return func() {
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
defer os.Setenv("SPANNER_OPTIMIZER_STATISTICS_PACKAGE", "")
}
}

func testReadOnlyTransaction(t *testing.T, executionTimes map[string]SimulatedExecutionTime) error {
server, client, teardown := setupMockedTestServer(t)
defer teardown()
Expand Down Expand Up @@ -784,8 +796,8 @@ func TestClient_ReadWriteTransaction_Query_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

ctx := context.Background()
Expand Down Expand Up @@ -813,8 +825,8 @@ func TestClient_ReadWriteTransaction_Update_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

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

func TestBatchReadOnlyTransaction_QueryOptions(t *testing.T) {
ctx := context.Background()
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
OptimizerVersion: "1",
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
}}
_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{QueryOptions: qo})
defer teardown()

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

func TestBatchReadOnlyTransactionFromID_QueryOptions(t *testing.T) {
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
OptimizerVersion: "1",
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
}}
_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{QueryOptions: qo})
defer teardown()

Expand All @@ -2276,48 +2294,49 @@ type QueryOptionsTestCase struct {
}

func queryOptionsTestCases() []QueryOptionsTestCase {
statsPkg := "auto_20191128_14_47_22UTC"
return []QueryOptionsTestCase{
{
"Client level",
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: nil},
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
},
{
"Environment level",
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
},
{
"Query level",
QueryOptions{Options: nil},
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
},
{
"Environment level has precedence",
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
},
{
"Query level has precedence than client level",
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: nil},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
},
{
"Query level has highest precedence",
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3"}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "2", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "3", OptimizerStatisticsPackage: statsPkg}},
},
}
}
Expand Down
5 changes: 4 additions & 1 deletion spanner/integration_test.go
View file
Expand Up @@ -635,7 +635,10 @@ func TestIntegration_SingleUse_WithQueryOptions(t *testing.T) {
t.Fatal(err)
}
}
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{OptimizerVersion: "1"}}
qo := QueryOptions{Options: &sppb.ExecuteSqlRequest_QueryOptions{
OptimizerVersion: "1",
OptimizerStatisticsPackage: "auto_20191128_14_47_22UTC",
}}
got, err := readAll(client.Single().QueryWithOptions(ctx, Statement{
"SELECT SingerId, FirstName, LastName FROM Singers WHERE SingerId IN (@id1, @id3, @id4)",
map[string]interface{}{"id1": int64(1), "id3": int64(3), "id4": int64(4)},
Expand Down
4 changes: 2 additions & 2 deletions spanner/pdml_test.go
View file
Expand Up @@ -145,8 +145,8 @@ func TestPartitionedUpdate_QueryOptions(t *testing.T) {
for _, tt := range queryOptionsTestCases() {
t.Run(tt.name, func(t *testing.T) {
if tt.env.Options != nil {
os.Setenv("SPANNER_OPTIMIZER_VERSION", tt.env.Options.OptimizerVersion)
defer os.Setenv("SPANNER_OPTIMIZER_VERSION", "")
unset := setQueryOptionsEnvVars(tt.env.Options)
defer unset()
}

ctx := context.Background()
Expand Down

0 comments on commit 29c7247

Please sign in to comment.