feat: add query statistics package support (#129)

* feat: add query statistics package support * style: fix lint * test: reorder env mocks Co-authored-by: larkee <larkee@users.noreply.github.com>

Author: larkee

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
 

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.

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)
 

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")