From 048355f690dfd3ecc57d8689973ee472ad72e926 Mon Sep 17 00:00:00 2001 From: Hengfeng Li Date: Tue, 11 Aug 2020 14:12:14 +1000 Subject: [PATCH] feat(spanner): add the support of optimizer statistics package --- spanner/batch_test.go | 8 +++-- spanner/client.go | 10 ++++-- spanner/client_test.go | 71 +++++++++++++++++++++++-------------- spanner/integration_test.go | 5 ++- spanner/pdml_test.go | 4 +-- 5 files changed, 63 insertions(+), 35 deletions(-) diff --git a/spanner/batch_test.go b/spanner/batch_test.go index d3396749adb..0c77576b94f 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 afa35315e29..f40ddf4ad75 100644 --- a/spanner/client.go +++ b/spanner/client.go @@ -221,13 +221,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 } diff --git a/spanner/client_test.go b/spanner/client_test.go index c7ec16bc535..56d8a73cde4 100644 --- a/spanner/client_test.go +++ b/spanner/client_test.go @@ -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() @@ -445,6 +445,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 { @@ -610,8 +613,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() @@ -632,6 +635,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() @@ -771,8 +783,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() @@ -800,8 +812,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() @@ -2206,7 +2218,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() @@ -2222,7 +2237,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() @@ -2242,48 +2260,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 73287c4b8d6..d5108d0e9c4 100644 --- a/spanner/integration_test.go +++ b/spanner/integration_test.go @@ -620,7 +620,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 bbbc2df3187..ca4ac323044 100644 --- a/spanner/pdml_test.go +++ b/spanner/pdml_test.go @@ -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()