From 29c724771f0b19849c76e62d4bc8e9342922bf75 Mon Sep 17 00:00:00 2001 From: Hengfeng Li Date: Tue, 8 Jun 2021 16:30:21 +1000 Subject: [PATCH] feat(spanner): add the support of optimizer statistics package (#2717) --- 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 18f136cfa90..9659bb566fc 100644 --- a/spanner/client.go +++ b/spanner/client.go @@ -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 } diff --git a/spanner/client_test.go b/spanner/client_test.go index d1e271c7319..2b09d1b5ee2 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() @@ -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 { @@ -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() @@ -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() @@ -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() @@ -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() @@ -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() @@ -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() @@ -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}}, }, } } diff --git a/spanner/integration_test.go b/spanner/integration_test.go index 6f415ebb171..511788ed819 100644 --- a/spanner/integration_test.go +++ b/spanner/integration_test.go @@ -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)}, 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()