Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add query statistics package support #129

Merged
merged 3 commits into from Jun 7, 2021
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
10 changes: 7 additions & 3 deletions google/cloud/spanner_v1/_helpers.py
Expand Up @@ -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

Expand Down
8 changes: 7 additions & 1 deletion google/cloud/spanner_v1/client.py
Expand Up @@ -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():
Expand All @@ -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.

Expand Down Expand Up @@ -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.
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 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 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 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
30 changes: 21 additions & 9 deletions tests/unit/test_client.py
Expand Up @@ -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")
Expand Down