From 6598deade66c8887514a1a6571fffb1bd7b16fd0 Mon Sep 17 00:00:00 2001 From: larkee <31196561+larkee@users.noreply.github.com> Date: Mon, 7 Jun 2021 16:20:05 +1000 Subject: [PATCH] feat: add query statistics package support (#129) * feat: add query statistics package support * style: fix lint * test: reorder env mocks Co-authored-by: larkee --- 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 | 30 ++++++++++++++++++++--------- 4 files changed, 55 insertions(+), 18 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..d5ccf39546 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..d33d9cc08a 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -179,30 +179,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_statistics_package" + ) @mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version") - def test_constructor_custom_query_options_env_config(self, mock_ver): + 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")