diff --git a/spanner/batch_test.go b/spanner/batch_test.go index d3396749adbd..0c77576b94f5 100644 --- a/spanner/batch_test.go +++ b/spanner/batch_test.go @@ -18,7 +18,6 @@ package spanner import ( "context" - "os" "sync" "testing" "time" @@ -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() @@ -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) + } } }) } diff --git a/spanner/client.go b/spanner/client.go index c4e9f2603b69..5d6430abe2e3 100644 --- a/spanner/client.go +++ b/spanner/client.go @@ -229,13 +229,17 @@ func NewClientWithConfig(ctx context.Context, database string, config ClientConf // 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 } diff --git a/spanner/client_test.go b/spanner/client_test.go index c410c9513853..f82324201afe 100644 --- a/spanner/client_test.go +++ b/spanner/client_test.go @@ -349,8 +349,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() @@ -397,6 +397,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 { @@ -562,8 +565,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() @@ -584,6 +587,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() @@ -723,8 +735,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() @@ -752,8 +764,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() @@ -2010,7 +2022,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() @@ -2026,7 +2041,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() @@ -2046,48 +2064,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}}, }, } } diff --git a/spanner/integration_test.go b/spanner/integration_test.go index de5feccb6fd9..81810624acdc 100644 --- a/spanner/integration_test.go +++ b/spanner/integration_test.go @@ -568,7 +568,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)}, diff --git a/spanner/pdml_test.go b/spanner/pdml_test.go index 430ad674e89e..9ea5ca8a806e 100644 --- a/spanner/pdml_test.go +++ b/spanner/pdml_test.go @@ -146,8 +146,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()