Skip to content

Commit

Permalink
feat: add query statistics package support
Browse files Browse the repository at this point in the history
  • Loading branch information
larkee committed Aug 6, 2020
1 parent 4529ad3 commit a468d3e
Show file tree
Hide file tree
Showing 4 changed files with 59 additions and 19 deletions.
10 changes: 7 additions & 3 deletions google/cloud/spanner_v1/_helpers.py
Expand Up @@ -73,15 +73,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", ""),
)
combined.MergeFrom(merge)
if not combined.optimizer_version:
if not combined.optimizer_version and not combined.optimizer_statistics_package:
return None
return combined

Expand Down
8 changes: 7 additions & 1 deletion google/cloud/spanner_v1/client.py
Expand Up @@ -65,6 +65,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"
_USER_AGENT_DEPRECATED = (
"The 'user_agent' argument to 'Client' is deprecated / unused. "
"Please pass an appropriate 'client_info' instead."
Expand All @@ -79,6 +80,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 InstanceConfig(object):
"""Named configurations for Spanner instances.
Expand Down Expand Up @@ -197,7 +202,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.
Expand Down
25 changes: 20 additions & 5 deletions tests/unit/test__helpers.py
Expand Up @@ -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.proto.spanner_pb2 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)

Expand All @@ -48,7 +54,10 @@ def test_base_none_merge_object(self):
from google.cloud.spanner_v1.proto.spanner_pb2 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)

Expand All @@ -64,9 +73,15 @@ def test_base_none_merge_dict(self):
def test_base_object_merge_dict(self):
from google.cloud.spanner_v1.proto.spanner_pb2 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)

Expand Down
35 changes: 25 additions & 10 deletions tests/unit/test_client.py
Expand Up @@ -104,7 +104,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")
Expand Down Expand Up @@ -192,30 +195,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_stati)stics_package"
)
def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats):
from google.cloud.spanner_v1.proto.spanner_pb2 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")
Expand Down

0 comments on commit a468d3e

Please sign in to comment.