Skip to content

Commit 6eecfdc

Browse files
crwilcoxtseaver
crwilcox
authored and
tseaver
committed
Add timeout + retry settings to Sessions/Snapshots (#6536)
1 parent 74dc48d commit 6eecfdc

File tree

4 files changed

+58
-8
lines changed

4 files changed

+58
-8
lines changed

‎spanner/google/cloud/spanner_v1/session.py‎

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
# pylint: disable=ungrouped-imports
2323
from google.api_core.exceptions import Aborted, GoogleAPICallError, NotFound
24+
import google.api_core.gapic_v1.method
2425
from google.cloud.spanner_v1._helpers import _metadata_with_prefix
2526
from google.cloud.spanner_v1.batch import Batch
2627
from google.cloud.spanner_v1.snapshot import Snapshot
@@ -197,7 +198,9 @@ def read(self, table, columns, keyset, index='', limit=0):
197198
"""
198199
return self.snapshot().read(table, columns, keyset, index, limit)
199200

200-
def execute_sql(self, sql, params=None, param_types=None, query_mode=None):
201+
def execute_sql(self, sql, params=None, param_types=None, query_mode=None,
202+
retry=google.api_core.gapic_v1.method.DEFAULT,
203+
timeout=google.api_core.gapic_v1.method.DEFAULT):
201204
"""Perform an ``ExecuteStreamingSql`` API request.
202205
203206
:type sql: str
@@ -222,7 +225,7 @@ def execute_sql(self, sql, params=None, param_types=None, query_mode=None):
222225
:returns: a result set instance which can be used to consume rows.
223226
"""
224227
return self.snapshot().execute_sql(
225-
sql, params, param_types, query_mode)
228+
sql, params, param_types, query_mode, retry=retry, timeout=timeout)
226229

227230
def batch(self):
228231
"""Factory to create a batch for this session.

‎spanner/google/cloud/spanner_v1/snapshot.py‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.cloud.spanner_v1.proto.transaction_pb2 import TransactionSelector
2222

2323
from google.api_core.exceptions import ServiceUnavailable
24+
import google.api_core.gapic_v1.method
2425
from google.cloud._helpers import _datetime_to_pb_timestamp
2526
from google.cloud._helpers import _timedelta_to_duration_pb
2627
from google.cloud.spanner_v1._helpers import _make_value_pb
@@ -143,7 +144,9 @@ def read(self, table, columns, keyset, index='', limit=0, partition=None):
143144
return StreamedResultSet(iterator)
144145

145146
def execute_sql(self, sql, params=None, param_types=None,
146-
query_mode=None, partition=None):
147+
query_mode=None, partition=None,
148+
retry=google.api_core.gapic_v1.method.DEFAULT,
149+
timeout=google.api_core.gapic_v1.method.DEFAULT):
147150
"""Perform an ``ExecuteStreamingSql`` API request.
148151
149152
:type sql: str
@@ -204,7 +207,9 @@ def execute_sql(self, sql, params=None, param_types=None,
204207
query_mode=query_mode,
205208
partition_token=partition,
206209
seqno=self._execute_sql_count,
207-
metadata=metadata)
210+
metadata=metadata,
211+
retry=retry,
212+
timeout=timeout)
208213

209214
iterator = _restart_on_unavailable(restart)
210215

‎spanner/tests/unit/test_session.py‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515

1616
import unittest
17-
17+
import google.api_core.gapic_v1.method
1818
import mock
1919

2020

@@ -371,6 +371,36 @@ def test_execute_sql_defaults(self):
371371
None,
372372
None,
373373
None,
374+
timeout=google.api_core.gapic_v1.method.DEFAULT,
375+
retry=google.api_core.gapic_v1.method.DEFAULT,
376+
)
377+
378+
def test_execute_sql_non_default_retry(self):
379+
from google.protobuf.struct_pb2 import Struct, Value
380+
from google.cloud.spanner_v1.proto.type_pb2 import STRING
381+
382+
SQL = 'SELECT first_name, age FROM citizens'
383+
database = self._make_database()
384+
session = self._make_one(database)
385+
session._session_id = 'DEADBEEF'
386+
387+
params = Struct(fields={'foo': Value(string_value='bar')})
388+
param_types = {'foo': STRING}
389+
390+
with mock.patch(
391+
'google.cloud.spanner_v1.session.Snapshot') as snapshot:
392+
found = session.execute_sql(
393+
SQL, params, param_types, 'PLAN', retry=None, timeout=None)
394+
395+
self.assertIs(found, snapshot().execute_sql.return_value)
396+
397+
snapshot().execute_sql.assert_called_once_with(
398+
SQL,
399+
params,
400+
param_types,
401+
'PLAN',
402+
timeout=None,
403+
retry=None
374404
)
375405

376406
def test_execute_sql_explicit(self):
@@ -397,6 +427,8 @@ def test_execute_sql_explicit(self):
397427
params,
398428
param_types,
399429
'PLAN',
430+
timeout=google.api_core.gapic_v1.method.DEFAULT,
431+
retry=google.api_core.gapic_v1.method.DEFAULT,
400432
)
401433

402434
def test_batch_not_created(self):

‎spanner/tests/unit/test_snapshot.py‎

Lines changed: 13 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515

1616
import unittest
17-
17+
import google.api_core.gapic_v1.method
1818
import mock
1919

2020

@@ -333,7 +333,9 @@ def test_execute_sql_w_params_wo_param_types(self):
333333
derived.execute_sql(SQL_QUERY_WITH_PARAM, PARAMS)
334334

335335
def _execute_sql_helper(
336-
self, multi_use, first=True, count=0, partition=None, sql_count=0):
336+
self, multi_use, first=True, count=0, partition=None, sql_count=0,
337+
timeout=google.api_core.gapic_v1.method.DEFAULT,
338+
retry=google.api_core.gapic_v1.method.DEFAULT):
337339
from google.protobuf.struct_pb2 import Struct
338340
from google.cloud.spanner_v1.proto.result_set_pb2 import (
339341
PartialResultSet, ResultSetMetadata, ResultSetStats)
@@ -380,7 +382,7 @@ def _execute_sql_helper(
380382

381383
result_set = derived.execute_sql(
382384
SQL_QUERY_WITH_PARAM, PARAMS, PARAM_TYPES,
383-
query_mode=MODE, partition=partition)
385+
query_mode=MODE, partition=partition, retry=retry, timeout=timeout)
384386

385387
self.assertEqual(derived._read_request_count, count + 1)
386388

@@ -417,6 +419,8 @@ def _execute_sql_helper(
417419
partition_token=partition,
418420
seqno=sql_count,
419421
metadata=[('google-cloud-resource-prefix', database.name)],
422+
timeout=timeout,
423+
retry=retry,
420424
)
421425

422426
self.assertEqual(derived._execute_sql_count, sql_count + 1)
@@ -441,6 +445,12 @@ def test_execute_sql_w_multi_use_w_first_w_count_gt_0(self):
441445
with self.assertRaises(ValueError):
442446
self._execute_sql_helper(multi_use=True, first=True, count=1)
443447

448+
def test_execute_sql_w_retry(self):
449+
self._execute_sql_helper(multi_use=False, retry=None)
450+
451+
def test_execute_sql_w_timeout(self):
452+
self._execute_sql_helper(multi_use=False, timeout=None)
453+
444454
def _partition_read_helper(
445455
self, multi_use, w_txn,
446456
size=None, max_partitions=None, index=None):

0 commit comments

Comments
 (0)