Skip to content

Commit 6598dea

Browse files
larkeelarkee
authored
feat: add query statistics package support (#129)
* feat: add query statistics package support * style: fix lint * test: reorder env mocks Co-authored-by: larkee <[email protected]>
1 parent 1c03dcc commit 6598dea

File tree

4 files changed

+55
-18
lines changed

4 files changed

+55
-18
lines changed

‎google/cloud/spanner_v1/_helpers.py

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -84,15 +84,19 @@ def _merge_query_options(base, merge):
8484
combined = base or ExecuteSqlRequest.QueryOptions()
8585
if type(combined) == dict:
8686
combined = ExecuteSqlRequest.QueryOptions(
87-
optimizer_version=combined.get("optimizer_version", "")
87+
optimizer_version=combined.get("optimizer_version", ""),
88+
optimizer_statistics_package=combined.get(
89+
"optimizer_statistics_package", ""
90+
),
8891
)
8992
merge = merge or ExecuteSqlRequest.QueryOptions()
9093
if type(merge) == dict:
9194
merge = ExecuteSqlRequest.QueryOptions(
92-
optimizer_version=merge.get("optimizer_version", "")
95+
optimizer_version=merge.get("optimizer_version", ""),
96+
optimizer_statistics_package=merge.get("optimizer_statistics_package", ""),
9397
)
9498
type(combined).pb(combined).MergeFrom(type(merge).pb(merge))
95-
if not combined.optimizer_version:
99+
if not combined.optimizer_version and not combined.optimizer_statistics_package:
96100
return None
97101
return combined
98102

‎google/cloud/spanner_v1/client.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,7 @@
6464
) % ((EMULATOR_ENV_VAR,) * 3)
6565
SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin"
6666
OPTIMIZER_VERSION_ENV_VAR = "SPANNER_OPTIMIZER_VERSION"
67+
OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR = "SPANNER_OPTIMIZER_STATISTICS_PACKAGE"
6768

6869

6970
def _get_spanner_emulator_host():
@@ -74,6 +75,10 @@ def _get_spanner_optimizer_version():
7475
return os.getenv(OPTIMIZER_VERSION_ENV_VAR, "")
7576

7677

78+
def _get_spanner_optimizer_statistics_package():
79+
return os.getenv(OPTIMIZER_STATISITCS_PACKAGE_ENV_VAR, "")
80+
81+
7782
class Client(ClientWithProject):
7883
"""Client for interacting with Cloud Spanner API.
7984
@@ -160,7 +165,8 @@ def __init__(
160165
self._client_info = client_info
161166

162167
env_query_options = ExecuteSqlRequest.QueryOptions(
163-
optimizer_version=_get_spanner_optimizer_version()
168+
optimizer_version=_get_spanner_optimizer_version(),
169+
optimizer_statistics_package=_get_spanner_optimizer_statistics_package(),
164170
)
165171

166172
# Environment flag config has higher precedence than application config.

‎tests/unit/test__helpers.py

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,15 @@ def test_base_none_and_merge_none(self):
3030
def test_base_dict_and_merge_none(self):
3131
from google.cloud.spanner_v1 import ExecuteSqlRequest
3232

33-
base = {"optimizer_version": "2"}
33+
base = {
34+
"optimizer_version": "2",
35+
"optimizer_statistics_package": "auto_20191128_14_47_22UTC",
36+
}
3437
merge = None
35-
expected = ExecuteSqlRequest.QueryOptions(optimizer_version="2")
38+
expected = ExecuteSqlRequest.QueryOptions(
39+
optimizer_version="2",
40+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
41+
)
3642
result = self._callFUT(base, merge)
3743
self.assertEqual(result, expected)
3844

@@ -48,7 +54,10 @@ def test_base_none_merge_object(self):
4854
from google.cloud.spanner_v1 import ExecuteSqlRequest
4955

5056
base = None
51-
merge = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
57+
merge = ExecuteSqlRequest.QueryOptions(
58+
optimizer_version="3",
59+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
60+
)
5261
result = self._callFUT(base, merge)
5362
self.assertEqual(result, merge)
5463

@@ -64,9 +73,15 @@ def test_base_none_merge_dict(self):
6473
def test_base_object_merge_dict(self):
6574
from google.cloud.spanner_v1 import ExecuteSqlRequest
6675

67-
base = ExecuteSqlRequest.QueryOptions(optimizer_version="1")
76+
base = ExecuteSqlRequest.QueryOptions(
77+
optimizer_version="1",
78+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
79+
)
6880
merge = {"optimizer_version": "3"}
69-
expected = ExecuteSqlRequest.QueryOptions(optimizer_version="3")
81+
expected = ExecuteSqlRequest.QueryOptions(
82+
optimizer_version="3",
83+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
84+
)
7085
result = self._callFUT(base, merge)
7186
self.assertEqual(result, expected)
7287

‎tests/unit/test_client.py

Lines changed: 21 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -179,30 +179,42 @@ def test_constructor_custom_query_options_client_config(self):
179179

180180
expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
181181
creds = _make_credentials()
182+
query_options = expected_query_options = ExecuteSqlRequest.QueryOptions(
183+
optimizer_version="1",
184+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
185+
)
182186
self._constructor_test_helper(
183187
expected_scopes,
184188
creds,
185-
query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"),
186-
expected_query_options=ExecuteSqlRequest.QueryOptions(
187-
optimizer_version="1"
188-
),
189+
query_options=query_options,
190+
expected_query_options=expected_query_options,
189191
)
190192

193+
@mock.patch(
194+
"google.cloud.spanner_v1.client._get_spanner_optimizer_statistics_package"
195+
)
191196
@mock.patch("google.cloud.spanner_v1.client._get_spanner_optimizer_version")
192-
def test_constructor_custom_query_options_env_config(self, mock_ver):
197+
def test_constructor_custom_query_options_env_config(self, mock_ver, mock_stats):
193198
from google.cloud.spanner_v1 import ExecuteSqlRequest
194199
from google.cloud.spanner_v1 import client as MUT
195200

196201
expected_scopes = (MUT.SPANNER_ADMIN_SCOPE,)
197202
creds = _make_credentials()
198203
mock_ver.return_value = "2"
204+
mock_stats.return_value = "auto_20191128_14_47_22UTC"
205+
query_options = ExecuteSqlRequest.QueryOptions(
206+
optimizer_version="1",
207+
optimizer_statistics_package="auto_20191128_10_47_22UTC",
208+
)
209+
expected_query_options = ExecuteSqlRequest.QueryOptions(
210+
optimizer_version="2",
211+
optimizer_statistics_package="auto_20191128_14_47_22UTC",
212+
)
199213
self._constructor_test_helper(
200214
expected_scopes,
201215
creds,
202-
query_options=ExecuteSqlRequest.QueryOptions(optimizer_version="1"),
203-
expected_query_options=ExecuteSqlRequest.QueryOptions(
204-
optimizer_version="2"
205-
),
216+
query_options=query_options,
217+
expected_query_options=expected_query_options,
206218
)
207219

208220
@mock.patch("google.cloud.spanner_v1.client._get_spanner_emulator_host")

0 commit comments

Comments
 (0)