Skip to content
Permalink
Browse files
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>
  • Loading branch information
larkee and larkee committed Jun 7, 2021
1 parent 1c03dcc commit 6598deade66c8887514a1a6571fffb1bd7b16fd0
Showing with 55 additions and 18 deletions.
  1. +7 −3 google/cloud/spanner_v1/_helpers.py
  2. +7 −1 google/cloud/spanner_v1/client.py
  3. +20 −5 tests/unit/test__helpers.py
  4. +21 −9 tests/unit/test_client.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

@@ -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.
@@ -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)

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

0 comments on commit 6598dea

Please sign in to comment.