From 7f926f78e10f060ea29c7f84df9ce310a3b8fdaf Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 6 Aug 2020 15:17:38 +1000 Subject: [PATCH 1/3] feat: add query statistics package support --- google/cloud/spanner_v1/_helpers.py | 10 ++++++--- google/cloud/spanner_v1/client.py | 8 ++++++- tests/unit/test__helpers.py | 25 ++++++++++++++++----- tests/unit/test_client.py | 35 ++++++++++++++++++++--------- 4 files changed, 59 insertions(+), 19 deletions(-) diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index 1385809162..2d1bf322bf 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -84,15 +84,19 @@ def _merge_query_options(base, merge): combined = base or ExecuteSqlRequest.QueryOptions() if type(combined) == dict: combined = ExecuteSqlRequest.QueryOptions( - optimizer_version=combined.get("optimizer_version", "") + optimizer_version=combined.get("optimizer_version", ""), + optimizer_statistics_package=combined.get( + "optimizer_statistics_package", "" + ), ) merge = merge or ExecuteSqlRequest.QueryOptions() if type(merge) == dict: merge = ExecuteSqlRequest.QueryOptions( - optimizer_version=merge.get("optimizer_version", "") + optimizer_version=merge.get("optimizer_version", ""), + optimizer_statistics_package=merge.get("optimizer_statistics_package", ""), ) type(combined).pb(combined).MergeFrom(type(merge).pb(merge)) - if not combined.optimizer_version: + if not combined.optimizer_version and not combined.optimizer_statistics_package: return None return combined diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 1b447cbfa8..e74ae09b36 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -64,6 +64,7 @@ ) % ((EMULATOR_ENV_VAR,) * 3) SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin" OPTIMIZER_VERSION_ENV_VAR = "SPANNER_OPTIMIZER_VERSION" +OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR = "SPANNER_OPTIMIZER_STATISTICS_PACKAGE" def _get_spanner_emulator_host(): @@ -74,6 +75,10 @@ def _get_spanner_optimizer_version(): return os.getenv(OPTIMIZER_VERSION_ENV_VAR, "") +def _get_spanner_optimizer_statistics_package(): + return os.getenv(OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR, "") + + class Client(ClientWithProject): """Client for interacting with Cloud Spanner API. @@ -160,7 +165,8 @@ def __init__( self._client_info = client_info env_query_options = ExecuteSqlRequest.QueryOptions( - optimizer_version=_get_spanner_optimizer_version() + optimizer_version=_get_spanner_optimizer_version(), + optimizer_statistics_package=_get_spanner_optimizer_statistics_package() ) # Environment flag config has higher precedence than application config. diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index 661a2c0472..2ee66ed154 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -30,9 +30,15 @@ def test_base_none_and_merge_none(self): def test_base_dict_and_merge_none(self): from google.cloud.spanner_v1 import ExecuteSqlRequest - base = {"optimizer_version": "2"} + base = { + "optimizer_version": "2", + "optimizer_statistics_package": "auto_20191128_14_47_22UTC", + } merge = None - expected = ExecuteSqlRequest.QueryOptions(optimizer_version="2") + expected = ExecuteSqlRequest.QueryOptions( + optimizer_version="2", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) result = self._callFUT(base, merge) self.assertEqual(result, expected) @@ -48,7 +54,10 @@ def test_base_none_merge_object(self): from google.cloud.spanner_v1 import ExecuteSqlRequest base = None - merge = ExecuteSqlRequest.QueryOptions(optimizer_version="3") + merge = ExecuteSqlRequest.QueryOptions( + optimizer_version="3", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) result = self._callFUT(base, merge) self.assertEqual(result, merge) @@ -64,9 +73,15 @@ def test_base_none_merge_dict(self): def test_base_object_merge_dict(self): from google.cloud.spanner_v1 import ExecuteSqlRequest - base = ExecuteSqlRequest.QueryOptions(optimizer_version="1") + base = ExecuteSqlRequest.QueryOptions( + optimizer_version="1", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) merge = {"optimizer_version": "3"} - expected = ExecuteSqlRequest.QueryOptions(optimizer_version="3") + expected = ExecuteSqlRequest.QueryOptions( + optimizer_version="3", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) result = self._callFUT(base, merge) self.assertEqual(result, expected) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 40d10de9df..b2301554f6 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -103,7 +103,10 @@ def _constructor_test_helper( expected_client_options.api_endpoint, ) if expected_query_options is not None: - self.assertEqual(client._query_options, expected_query_options) + self.assertEqual( + client._query_options, + expected_query_options + ) @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host") @mock.patch("warnings.warn") @@ -179,30 +182,42 @@ def test_constructor_custom_query_options_client_config(self): expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,) creds = _make_credentials() + query_options = expected_query_options = ExecuteSqlRequest.QueryOptions( + optimizer_version="1", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) self._constructor_test_helper( expected_scopes, creds, - query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"), - expected_query_options=ExecuteSqlRequest.QueryOptions( - optimizer_version="1" - ), + query_options=query_options, + expected_query_options=expected_query_options, ) @mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version") - def test_constructor_custom_query_options_env_config(self, mock_ver): + @mock.patch( + "google.cloud.spanner_v1.client._get_spanner_optimizer_statistics_package" + ) + def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats): from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import client as MUT expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,) creds = _make_credentials() mock_ver.return_value = "2" + mock_stats.return_value = "auto_20191128_14_47_22UTC" + query_options = ExecuteSqlRequest.QueryOptions( + optimizer_version="1", + optimizer_statistics_package="auto_20191128_10_47_22UTC", + ) + expected_query_options = ExecuteSqlRequest.QueryOptions( + optimizer_version="2", + optimizer_statistics_package="auto_20191128_14_47_22UTC", + ) self._constructor_test_helper( expected_scopes, creds, - query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"), - expected_query_options=ExecuteSqlRequest.QueryOptions( - optimizer_version="2" - ), + query_options=query_options, + expected_query_options=expected_query_options, ) @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host") From ae6a8ba3d13ec6cd6d5c0829fb9bd5b4251a7089 Mon Sep 17 00:00:00 2001 From: larkee Date: Thu, 8 Apr 2021 16:01:56 +1000 Subject: [PATCH 2/3] style: fix lint --- google/cloud/spanner_v1/client.py | 2 +- tests/unit/test_client.py | 5 +---- 2 files changed, 2 insertions(+), 5 deletions(-) diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index e74ae09b36..d5ccf39546 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -166,7 +166,7 @@ def __init__( env_query_options = ExecuteSqlRequest.QueryOptions( optimizer_version=_get_spanner_optimizer_version(), - optimizer_statistics_package=_get_spanner_optimizer_statistics_package() + optimizer_statistics_package=_get_spanner_optimizer_statistics_package(), ) # Environment flag config has higher precedence than application config. diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index b2301554f6..81bb207d5d 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -103,10 +103,7 @@ def _constructor_test_helper( expected_client_options.api_endpoint, ) if expected_query_options is not None: - self.assertEqual( - client._query_options, - expected_query_options - ) + self.assertEqual(client._query_options, expected_query_options) @mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host") @mock.patch("warnings.warn") From b249ebb629c99c6a1c4bfd134ac9d5e4187a35d5 Mon Sep 17 00:00:00 2001 From: larkee Date: Mon, 7 Jun 2021 15:20:06 +1000 Subject: [PATCH 3/3] test: reorder env mocks --- tests/unit/test_client.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 81bb207d5d..d33d9cc08a 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -190,10 +190,10 @@ def test_constructor_custom_query_options_client_config(self): expected_query_options=expected_query_options, ) - @mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version") @mock.patch( "google.cloud.spanner_v1.client._get_spanner_optimizer_statistics_package" ) + @mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version") def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats): from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import client as MUT