Skip to content

Commit 5147921

Browse files
authored
feat: implement query options versioning support (#30)
* feat: implement query options versioning support * refactor _merge_query_options to use MergeFrom protobuf function * address comments * fix assignment Co-authored-by: larkee <[email protected]>
1 parent 23916c5 commit 5147921

File tree

12 files changed

+385
-41
lines changed

12 files changed

+385
-41
lines changed

google/cloud/spanner_v1/_helpers.py

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
from google.cloud._helpers import _date_from_iso8601_date
2727
from google.cloud._helpers import _datetime_to_rfc3339
2828
from google.cloud.spanner_v1.proto import type_pb2
29+
from google.cloud.spanner_v1.proto.spanner_pb2 import ExecuteSqlRequest
2930

3031

3132
def _try_to_coerce_bytes(bytestring):
@@ -47,6 +48,44 @@ def _try_to_coerce_bytes(bytestring):
4748
)
4849

4950

51+
def _merge_query_options(base, merge):
52+
"""Merge higher precedence QueryOptions with current QueryOptions.
53+
54+
:type base:
55+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
56+
or :class:`dict` or None
57+
:param base: The current QueryOptions that is intended for use.
58+
59+
:type merge:
60+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
61+
or :class:`dict` or None
62+
:param merge:
63+
The QueryOptions that have a higher priority than base. These options
64+
should overwrite the fields in base.
65+
66+
:rtype:
67+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
68+
or None
69+
:returns:
70+
QueryOptions object formed by merging the two given QueryOptions.
71+
If the resultant object only has empty fields, returns None.
72+
"""
73+
combined = base or ExecuteSqlRequest.QueryOptions()
74+
if type(combined) == dict:
75+
combined = ExecuteSqlRequest.QueryOptions(
76+
optimizer_version=combined.get("optimizer_version", "")
77+
)
78+
merge = merge or ExecuteSqlRequest.QueryOptions()
79+
if type(merge) == dict:
80+
merge = ExecuteSqlRequest.QueryOptions(
81+
optimizer_version=merge.get("optimizer_version", "")
82+
)
83+
combined.MergeFrom(merge)
84+
if not combined.optimizer_version:
85+
return None
86+
return combined
87+
88+
5089
# pylint: disable=too-many-return-statements,too-many-branches
5190
def _make_value_pb(value):
5291
"""Helper for :func:`_make_list_value_pbs`.

google/cloud/spanner_v1/client.py

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -50,9 +50,10 @@
5050

5151
from google.cloud.client import ClientWithProject
5252
from google.cloud.spanner_v1 import __version__
53-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
53+
from google.cloud.spanner_v1._helpers import _merge_query_options, _metadata_with_prefix
5454
from google.cloud.spanner_v1.instance import DEFAULT_NODE_COUNT
5555
from google.cloud.spanner_v1.instance import Instance
56+
from google.cloud.spanner_v1.proto.spanner_pb2 import ExecuteSqlRequest
5657

5758
_CLIENT_INFO = client_info.ClientInfo(client_library_version=__version__)
5859
EMULATOR_ENV_VAR = "SPANNER_EMULATOR_HOST"
@@ -62,6 +63,7 @@
6263
"without a scheme: ex %s=localhost:8080."
6364
) % ((EMULATOR_ENV_VAR,) * 3)
6465
SPANNER_ADMIN_SCOPE = "https://www.googleapis.com/auth/spanner.admin"
66+
OPTIMIZER_VERSION_ENV_VAR = "SPANNER_OPTIMIZER_VERSION"
6567
_USER_AGENT_DEPRECATED = (
6668
"The 'user_agent' argument to 'Client' is deprecated / unused. "
6769
"Please pass an appropriate 'client_info' instead."
@@ -72,6 +74,10 @@ def _get_spanner_emulator_host():
7274
return os.getenv(EMULATOR_ENV_VAR)
7375

7476

77+
def _get_spanner_optimizer_version():
78+
return os.getenv(OPTIMIZER_VERSION_ENV_VAR, "")
79+
80+
7581
class InstanceConfig(object):
7682
"""Named configurations for Spanner instances.
7783
@@ -132,11 +138,20 @@ class Client(ClientWithProject):
132138
:param user_agent:
133139
(Deprecated) The user agent to be used with API request.
134140
Not used.
141+
135142
:type client_options: :class:`~google.api_core.client_options.ClientOptions`
136143
or :class:`dict`
137144
:param client_options: (Optional) Client options used to set user options
138145
on the client. API Endpoint should be set through client_options.
139146
147+
:type query_options:
148+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
149+
or :class:`dict`
150+
:param query_options:
151+
(Optional) Query optimizer configuration to use for the given query.
152+
If a dict is provided, it must be of the same form as the protobuf
153+
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
154+
140155
:raises: :class:`ValueError <exceptions.ValueError>` if both ``read_only``
141156
and ``admin`` are :data:`True`
142157
"""
@@ -157,6 +172,7 @@ def __init__(
157172
client_info=_CLIENT_INFO,
158173
user_agent=None,
159174
client_options=None,
175+
query_options=None,
160176
):
161177
# NOTE: This API has no use for the _http argument, but sending it
162178
# will have no impact since the _http() @property only lazily
@@ -172,6 +188,13 @@ def __init__(
172188
else:
173189
self._client_options = client_options
174190

191+
env_query_options = ExecuteSqlRequest.QueryOptions(
192+
optimizer_version=_get_spanner_optimizer_version()
193+
)
194+
195+
# Environment flag config has higher precedence than application config.
196+
self._query_options = _merge_query_options(query_options, env_query_options)
197+
175198
if user_agent is not None:
176199
warnings.warn(_USER_AGENT_DEPRECATED, DeprecationWarning, stacklevel=2)
177200
self.user_agent = user_agent

google/cloud/spanner_v1/database.py

Lines changed: 36 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -30,8 +30,11 @@
3030
import six
3131

3232
# pylint: disable=ungrouped-imports
33-
from google.cloud.spanner_v1._helpers import _make_value_pb
34-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
33+
from google.cloud.spanner_v1._helpers import (
34+
_make_value_pb,
35+
_merge_query_options,
36+
_metadata_with_prefix,
37+
)
3538
from google.cloud.spanner_v1.batch import Batch
3639
from google.cloud.spanner_v1.gapic.spanner_client import SpannerClient
3740
from google.cloud.spanner_v1.gapic.transports import spanner_grpc_transport
@@ -350,7 +353,9 @@ def drop(self):
350353
metadata = _metadata_with_prefix(self.name)
351354
api.drop_database(self.name, metadata=metadata)
352355

353-
def execute_partitioned_dml(self, dml, params=None, param_types=None):
356+
def execute_partitioned_dml(
357+
self, dml, params=None, param_types=None, query_options=None
358+
):
354359
"""Execute a partitionable DML statement.
355360
356361
:type dml: str
@@ -365,9 +370,20 @@ def execute_partitioned_dml(self, dml, params=None, param_types=None):
365370
(Optional) maps explicit types for one or more param values;
366371
required if parameters are passed.
367372
373+
:type query_options:
374+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
375+
or :class:`dict`
376+
:param query_options:
377+
(Optional) Query optimizer configuration to use for the given query.
378+
If a dict is provided, it must be of the same form as the protobuf
379+
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
380+
368381
:rtype: int
369382
:returns: Count of rows affected by the DML statement.
370383
"""
384+
query_options = _merge_query_options(
385+
self._instance._client._query_options, query_options
386+
)
371387
if params is not None:
372388
if param_types is None:
373389
raise ValueError("Specify 'param_types' when passing 'params'.")
@@ -398,6 +414,7 @@ def execute_partitioned_dml(self, dml, params=None, param_types=None):
398414
transaction=txn_selector,
399415
params=params_pb,
400416
param_types=param_types,
417+
query_options=query_options,
401418
metadata=metadata,
402419
)
403420

@@ -748,6 +765,7 @@ def generate_query_batches(
748765
param_types=None,
749766
partition_size_bytes=None,
750767
max_partitions=None,
768+
query_options=None,
751769
):
752770
"""Start a partitioned query operation.
753771
@@ -783,6 +801,14 @@ def generate_query_batches(
783801
service uses this as a hint, the actual number of partitions may
784802
differ.
785803
804+
:type query_options:
805+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
806+
or :class:`dict`
807+
:param query_options:
808+
(Optional) Query optimizer configuration to use for the given query.
809+
If a dict is provided, it must be of the same form as the protobuf
810+
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
811+
786812
:rtype: iterable of dict
787813
:returns:
788814
mappings of information used peform actual partitioned reads via
@@ -801,6 +827,13 @@ def generate_query_batches(
801827
query_info["params"] = params
802828
query_info["param_types"] = param_types
803829

830+
# Query-level options have higher precedence than client-level and
831+
# environment-level options
832+
default_query_options = self._database._instance._client._query_options
833+
query_info["query_options"] = _merge_query_options(
834+
default_query_options, query_options
835+
)
836+
804837
for partition in partitions:
805838
yield {"partition": partition, "query": query_info}
806839

google/cloud/spanner_v1/session.py

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,7 @@ def execute_sql(
202202
params=None,
203203
param_types=None,
204204
query_mode=None,
205+
query_options=None,
205206
retry=google.api_core.gapic_v1.method.DEFAULT,
206207
timeout=google.api_core.gapic_v1.method.DEFAULT,
207208
):
@@ -225,11 +226,22 @@ def execute_sql(
225226
:param query_mode: Mode governing return of results / query plan. See
226227
https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode1
227228
229+
:type query_options:
230+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
231+
or :class:`dict`
232+
:param query_options: (Optional) Options that are provided for query plan stability.
233+
228234
:rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
229235
:returns: a result set instance which can be used to consume rows.
230236
"""
231237
return self.snapshot().execute_sql(
232-
sql, params, param_types, query_mode, retry=retry, timeout=timeout
238+
sql,
239+
params,
240+
param_types,
241+
query_mode,
242+
query_options=query_options,
243+
retry=retry,
244+
timeout=timeout,
233245
)
234246

235247
def batch(self):

google/cloud/spanner_v1/snapshot.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from google.api_core.exceptions import ServiceUnavailable
2424
import google.api_core.gapic_v1.method
2525
from google.cloud._helpers import _datetime_to_pb_timestamp
26+
from google.cloud.spanner_v1._helpers import _merge_query_options
2627
from google.cloud._helpers import _timedelta_to_duration_pb
2728
from google.cloud.spanner_v1._helpers import _make_value_pb
2829
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
@@ -157,6 +158,7 @@ def execute_sql(
157158
params=None,
158159
param_types=None,
159160
query_mode=None,
161+
query_options=None,
160162
partition=None,
161163
retry=google.api_core.gapic_v1.method.DEFAULT,
162164
timeout=google.api_core.gapic_v1.method.DEFAULT,
@@ -180,6 +182,14 @@ def execute_sql(
180182
:param query_mode: Mode governing return of results / query plan. See
181183
https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode1
182184
185+
:type query_options:
186+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
187+
or :class:`dict`
188+
:param query_options:
189+
(Optional) Query optimizer configuration to use for the given query.
190+
If a dict is provided, it must be of the same form as the protobuf
191+
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
192+
183193
:type partition: bytes
184194
:param partition: (Optional) one of the partition tokens returned
185195
from :meth:`partition_query`.
@@ -211,6 +221,11 @@ def execute_sql(
211221
transaction = self._make_txn_selector()
212222
api = database.spanner_api
213223

224+
# Query-level options have higher precedence than client-level and
225+
# environment-level options
226+
default_query_options = database._instance._client._query_options
227+
query_options = _merge_query_options(default_query_options, query_options)
228+
214229
restart = functools.partial(
215230
api.execute_streaming_sql,
216231
self._session.name,
@@ -221,6 +236,7 @@ def execute_sql(
221236
query_mode=query_mode,
222237
partition_token=partition,
223238
seqno=self._execute_sql_count,
239+
query_options=query_options,
224240
metadata=metadata,
225241
retry=retry,
226242
timeout=timeout,

google/cloud/spanner_v1/transaction.py

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@
1717
from google.protobuf.struct_pb2 import Struct
1818

1919
from google.cloud._helpers import _pb_timestamp_to_datetime
20-
from google.cloud.spanner_v1._helpers import _make_value_pb
21-
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
20+
from google.cloud.spanner_v1._helpers import (
21+
_make_value_pb,
22+
_merge_query_options,
23+
_metadata_with_prefix,
24+
)
2225
from google.cloud.spanner_v1.proto.transaction_pb2 import TransactionSelector
2326
from google.cloud.spanner_v1.proto.transaction_pb2 import TransactionOptions
2427
from google.cloud.spanner_v1.snapshot import _SnapshotBase
@@ -162,7 +165,9 @@ def _make_params_pb(params, param_types):
162165

163166
return None
164167

165-
def execute_update(self, dml, params=None, param_types=None, query_mode=None):
168+
def execute_update(
169+
self, dml, params=None, param_types=None, query_mode=None, query_options=None
170+
):
166171
"""Perform an ``ExecuteSql`` API request with DML.
167172
168173
:type dml: str
@@ -182,6 +187,11 @@ def execute_update(self, dml, params=None, param_types=None, query_mode=None):
182187
:param query_mode: Mode governing return of results / query plan. See
183188
https://cloud.google.com/spanner/reference/rpc/google.spanner.v1#google.spanner.v1.ExecuteSqlRequest.QueryMode1
184189
190+
:type query_options:
191+
:class:`google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
192+
or :class:`dict`
193+
:param query_options: (Optional) Options that are provided for query plan stability.
194+
185195
:rtype: int
186196
:returns: Count of rows affected by the DML statement.
187197
"""
@@ -191,13 +201,19 @@ def execute_update(self, dml, params=None, param_types=None, query_mode=None):
191201
transaction = self._make_txn_selector()
192202
api = database.spanner_api
193203

204+
# Query-level options have higher precedence than client-level and
205+
# environment-level options
206+
default_query_options = database._instance._client._query_options
207+
query_options = _merge_query_options(default_query_options, query_options)
208+
194209
response = api.execute_sql(
195210
self._session.name,
196211
dml,
197212
transaction=transaction,
198213
params=params_pb,
199214
param_types=param_types,
200215
query_mode=query_mode,
216+
query_options=query_options,
201217
seqno=self._execute_sql_count,
202218
metadata=metadata,
203219
)

0 commit comments

Comments
 (0)