def begin(self):
"""Begin a transaction on the database.
:rtype: bytes
:returns: the ID for the newly-begun transaction.
:raises ValueError:
if the transaction is already begun, committed, or rolled back.
"""
if self._transaction_id is not None:
raise ValueError("Transaction already begun")
if self.committed is not None:
raise ValueError("Transaction already committed")
if self._rolled_back:
raise ValueError("Transaction is already rolled back")
database = self._session._database
api = database.spanner_api
metadata = _metadata_with_prefix(database.name)
txn_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite())
response = api.begin_transaction(self._session.name,
txn_options,
metadata=metadata)
self._transaction_id = response.id
return self._transaction_id
def test_begin_ok_exact_staleness(self):
from google.protobuf.duration_pb2 import Duration
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB, TransactionOptions)
transaction_pb = TransactionPB(id=TXN_ID)
database = _Database()
api = database.spanner_api = self._make_spanner_api()
api.begin_transaction.return_value = transaction_pb
duration = self._makeDuration(seconds=SECONDS, microseconds=MICROS)
session = _Session(database)
snapshot = self._make_one(session,
exact_staleness=duration,
multi_use=True)
txn_id = snapshot.begin()
self.assertEqual(txn_id, TXN_ID)
self.assertEqual(snapshot._transaction_id, TXN_ID)
expected_duration = Duration(seconds=SECONDS, nanos=MICROS * 1000)
expected_txn_options = TransactionOptions(
read_only=TransactionOptions.ReadOnly(
exact_staleness=expected_duration))
api.begin_transaction.assert_called_once_with(
session.name,
expected_txn_options,
metadata=[('google-cloud-resource-prefix', database.name)])
def execute_partitioned_dml(self, dml, params=None, param_types=None):
"""Execute a partitionable DML statement.
:type dml: str
:param dml: DML statement
:type params: dict, {str -> column value}
:param params: values for parameter replacement. Keys must match
the names used in ``dml``.
:type param_types: dict[str -> Union[dict, .types.Type]]
:param param_types:
(Optional) maps explicit types for one or more param values;
required if parameters are passed.
:rtype: int
:returns: Count of rows affected by the DML statement.
"""
if params is not None:
if param_types is None:
raise ValueError(
"Specify 'param_types' when passing 'params'.")
params_pb = Struct(fields={
key: _make_value_pb(value)
for key, value in params.items()
})
else:
params_pb = None
api = self.spanner_api
txn_options = TransactionOptions(
partitioned_dml=TransactionOptions.PartitionedDml())
metadata = _metadata_with_prefix(self.name)
with SessionCheckout(self._pool) as session:
txn = api.begin_transaction(session.name,
txn_options,
metadata=metadata)
txn_selector = TransactionSelector(id=txn.id)
restart = functools.partial(
api.execute_streaming_sql,
session.name,
dml,
transaction=txn_selector,
params=params_pb,
param_types=param_types,
metadata=metadata,
)
iterator = _restart_on_unavailable(restart)
result_set = StreamedResultSet(iterator)
list(result_set) # consume all partials
return result_set.stats.row_count_lower_bound
def test_begin_ok_exact_strong(self):
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
transaction_pb = TransactionPB(id=TXN_ID)
database = _Database()
api = database.spanner_api = self._make_spanner_api()
api.begin_transaction.return_value = transaction_pb
session = _Session(database)
snapshot = self._make_one(session, multi_use=True)
txn_id = snapshot.begin()
self.assertEqual(txn_id, TXN_ID)
self.assertEqual(snapshot._transaction_id, TXN_ID)
expected_txn_options = TransactionOptions(
read_only=TransactionOptions.ReadOnly(strong=True))
api.begin_transaction.assert_called_once_with(
session.name,
expected_txn_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def _make_txn_selector(self):
"""Helper for :meth:`read`."""
if self._transaction_id is not None:
return TransactionSelector(id=self._transaction_id)
if self._read_timestamp:
key = "read_timestamp"
value = _datetime_to_pb_timestamp(self._read_timestamp)
elif self._min_read_timestamp:
key = "min_read_timestamp"
value = _datetime_to_pb_timestamp(self._min_read_timestamp)
elif self._max_staleness:
key = "max_staleness"
value = _timedelta_to_duration_pb(self._max_staleness)
elif self._exact_staleness:
key = "exact_staleness"
value = _timedelta_to_duration_pb(self._exact_staleness)
else:
key = "strong"
value = True
options = TransactionOptions(read_only=TransactionOptions.ReadOnly(
**{key: value}))
if self._multi_use:
return TransactionSelector(begin=options)
else:
return TransactionSelector(single_use=options)
def test_run_in_transaction_w_args_w_kwargs_wo_abort(self):
import datetime
from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB, TransactionOptions)
from google.cloud._helpers import UTC
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = 'citizens'
COLUMNS = ['email', 'first_name', 'last_name', 'age']
VALUES = [
['*****@*****.**', 'Phred', 'Phlyntstone', 32],
['*****@*****.**', 'Bharney', 'Rhubble', 31],
]
TRANSACTION_ID = b'FACEDACE'
transaction_pb = TransactionPB(id=TRANSACTION_ID)
now = datetime.datetime.utcnow().replace(tzinfo=UTC)
now_pb = _datetime_to_pb_timestamp(now)
response = CommitResponse(commit_timestamp=now_pb)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.return_value = response
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
return 42
return_value = session.run_in_transaction(unit_of_work,
'abc',
some_arg='def')
self.assertIsNone(session._transaction)
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIsInstance(txn, Transaction)
self.assertEqual(return_value, 42)
self.assertEqual(args, ('abc', ))
self.assertEqual(kw, {'some_arg': 'def'})
expected_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite(), )
gax_api.begin_transaction.assert_called_once_with(
self.SESSION_NAME,
expected_options,
metadata=[('google-cloud-resource-prefix', database.name)],
)
gax_api.commit.assert_called_once_with(
self.SESSION_NAME,
txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[('google-cloud-resource-prefix', database.name)],
)
def test_run_in_transaction_callback_raises_non_gax_error(self):
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
transaction_pb = TransactionPB(id=TRANSACTION_ID)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.rollback.return_value = None
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
class Testing(Exception):
pass
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
raise Testing()
with self.assertRaises(Testing):
session.run_in_transaction(unit_of_work)
self.assertIsNone(session._transaction)
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIsInstance(txn, Transaction)
self.assertIsNone(txn.committed)
self.assertTrue(txn._rolled_back)
self.assertEqual(args, ())
self.assertEqual(kw, {})
expected_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite())
gax_api.begin_transaction.assert_called_once_with(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
gax_api.rollback.assert_called_once_with(
self.SESSION_NAME,
TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def _make_txn_selector(self):
from google.cloud.spanner_v1.proto.transaction_pb2 import (
TransactionOptions, TransactionSelector)
if self._transaction_id:
return TransactionSelector(id=self._transaction_id)
options = TransactionOptions(
read_only=TransactionOptions.ReadOnly(strong=True))
if self._multi_use:
return TransactionSelector(begin=options)
return TransactionSelector(single_use=options)
def test_run_in_transaction_callback_raises_non_abort_rpc_error(self):
from google.api_core.exceptions import Cancelled
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB, TransactionOptions)
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = 'citizens'
COLUMNS = ['email', 'first_name', 'last_name', 'age']
VALUES = [
['*****@*****.**', 'Phred', 'Phlyntstone', 32],
['*****@*****.**', 'Bharney', 'Rhubble', 31],
]
TRANSACTION_ID = b'FACEDACE'
transaction_pb = TransactionPB(id=TRANSACTION_ID)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.rollback.return_value = None
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
raise Cancelled('error')
with self.assertRaises(Cancelled):
session.run_in_transaction(unit_of_work)
self.assertIsNone(session._transaction)
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIsInstance(txn, Transaction)
self.assertIsNone(txn.committed)
self.assertFalse(txn._rolled_back)
self.assertEqual(args, ())
self.assertEqual(kw, {})
expected_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite(), )
gax_api.begin_transaction.assert_called_once_with(
self.SESSION_NAME,
expected_options,
metadata=[('google-cloud-resource-prefix', database.name)],
)
gax_api.rollback.assert_not_called()
def commit(self):
"""Commit mutations to the database.
:rtype: datetime
:returns: timestamp of the committed changes.
"""
self._check_state()
database = self._session._database
api = database.spanner_api
metadata = _metadata_with_prefix(database.name)
txn_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite())
response = api.commit(self._session.name, self._mutations,
single_use_transaction=txn_options,
metadata=metadata)
self.committed = _pb_timestamp_to_datetime(
response.commit_timestamp)
return self.committed
def commit(self):
"""Commit mutations to the database.
:rtype: datetime
:returns: timestamp of the committed changes.
"""
self._check_state()
database = self._session._database
api = database.spanner_api
metadata = _metadata_with_prefix(database.name)
txn_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
trace_attributes = {"num_mutations": len(self._mutations)}
with trace_call("CloudSpanner.Commit", self._session, trace_attributes):
response = api.commit(
self._session.name,
mutations=self._mutations,
single_use_transaction=txn_options,
metadata=metadata,
)
self.committed = _pb_timestamp_to_datetime(response.commit_timestamp)
return self.committed
def test_run_in_transaction_w_callback_raises_abort_wo_metadata(self):
import datetime
from google.api_core.exceptions import Aborted
from google.protobuf.duration_pb2 import Duration
from google.rpc.error_details_pb2 import RetryInfo
from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud._helpers import UTC
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
RETRY_SECONDS = 1
RETRY_NANOS = 3456
transaction_pb = TransactionPB(id=TRANSACTION_ID)
now = datetime.datetime.utcnow().replace(tzinfo=UTC)
now_pb = _datetime_to_pb_timestamp(now)
response = CommitResponse(commit_timestamp=now_pb)
retry_info = RetryInfo(
retry_delay=Duration(seconds=RETRY_SECONDS, nanos=RETRY_NANOS)
)
trailing_metadata = [
("google.rpc.retryinfo-bin", retry_info.SerializeToString())
]
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = [response]
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
if len(called_with) < 2:
raise _make_rpc_error(Aborted, trailing_metadata)
txn.insert(TABLE_NAME, COLUMNS, VALUES)
with mock.patch("time.sleep") as sleep_mock:
session.run_in_transaction(unit_of_work)
sleep_mock.assert_called_once_with(RETRY_SECONDS + RETRY_NANOS / 1.0e9)
self.assertEqual(len(called_with), 2)
for index, (txn, args, kw) in enumerate(called_with):
self.assertIsInstance(txn, Transaction)
if index == 0:
self.assertIsNone(txn.committed)
else:
self.assertEqual(txn.committed, now)
self.assertEqual(args, ())
self.assertEqual(kw, {})
expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
self.assertEqual(
gax_api.begin_transaction.call_args_list,
[
mock.call(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
]
* 2,
)
gax_api.commit.assert_called_once_with(
self.SESSION_NAME,
mutations=txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def test_run_in_transaction_w_abort_no_retry_metadata(self):
import datetime
from google.api_core.exceptions import Aborted
from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud._helpers import UTC
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
transaction_pb = TransactionPB(id=TRANSACTION_ID)
now = datetime.datetime.utcnow().replace(tzinfo=UTC)
now_pb = _datetime_to_pb_timestamp(now)
aborted = _make_rpc_error(Aborted, trailing_metadata=[])
response = CommitResponse(commit_timestamp=now_pb)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = [aborted, response]
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
return "answer"
return_value = session.run_in_transaction(unit_of_work, "abc", some_arg="def")
self.assertEqual(len(called_with), 2)
for index, (txn, args, kw) in enumerate(called_with):
self.assertIsInstance(txn, Transaction)
self.assertEqual(return_value, "answer")
self.assertEqual(args, ("abc",))
self.assertEqual(kw, {"some_arg": "def"})
expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
self.assertEqual(
gax_api.begin_transaction.call_args_list,
[
mock.call(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
]
* 2,
)
self.assertEqual(
gax_api.commit.call_args_list,
[
mock.call(
self.SESSION_NAME,
mutations=txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
]
* 2,
)
def test_run_in_transaction_w_timeout(self):
from google.api_core.exceptions import Aborted
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
transaction_pb = TransactionPB(id=TRANSACTION_ID)
aborted = _make_rpc_error(Aborted, trailing_metadata=[])
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = aborted
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
# retry several times to check backoff
def _time(_results=[1, 2, 4, 8]):
return _results.pop(0)
with mock.patch("time.time", _time):
if HAS_OPENTELEMETRY_INSTALLED:
with mock.patch("opentelemetry.util.time", _ConstantTime()):
with mock.patch("time.sleep") as sleep_mock:
with self.assertRaises(Aborted):
session.run_in_transaction(unit_of_work, timeout_secs=8)
else:
with mock.patch("time.sleep") as sleep_mock:
with self.assertRaises(Aborted):
session.run_in_transaction(unit_of_work, timeout_secs=8)
# unpacking call args into list
call_args = [call_[0][0] for call_ in sleep_mock.call_args_list]
call_args = list(map(int, call_args))
assert call_args == [2, 4]
assert sleep_mock.call_count == 2
self.assertEqual(len(called_with), 3)
for txn, args, kw in called_with:
self.assertIsInstance(txn, Transaction)
self.assertIsNone(txn.committed)
self.assertEqual(args, ())
self.assertEqual(kw, {})
expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
self.assertEqual(
gax_api.begin_transaction.call_args_list,
[
mock.call(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
]
* 3,
)
self.assertEqual(
gax_api.commit.call_args_list,
[
mock.call(
self.SESSION_NAME,
mutations=txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
]
* 3,
)
def test_run_in_transaction_w_abort_w_retry_metadata_deadline(self):
import datetime
from google.api_core.exceptions import Aborted
from google.protobuf.duration_pb2 import Duration
from google.rpc.error_details_pb2 import RetryInfo
from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud.spanner_v1.transaction import Transaction
from google.cloud._helpers import UTC
from google.cloud._helpers import _datetime_to_pb_timestamp
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
RETRY_SECONDS = 1
RETRY_NANOS = 3456
transaction_pb = TransactionPB(id=TRANSACTION_ID)
now = datetime.datetime.utcnow().replace(tzinfo=UTC)
now_pb = _datetime_to_pb_timestamp(now)
response = CommitResponse(commit_timestamp=now_pb)
retry_info = RetryInfo(
retry_delay=Duration(seconds=RETRY_SECONDS, nanos=RETRY_NANOS)
)
trailing_metadata = [
("google.rpc.retryinfo-bin", retry_info.SerializeToString())
]
aborted = _make_rpc_error(Aborted, trailing_metadata=trailing_metadata)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = [aborted, response]
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
# retry once w/ timeout_secs=1
def _time(_results=[1, 1.5]):
return _results.pop(0)
with mock.patch("time.time", _time):
if HAS_OPENTELEMETRY_INSTALLED:
with mock.patch("opentelemetry.util.time", _ConstantTime()):
with mock.patch("time.sleep") as sleep_mock:
with self.assertRaises(Aborted):
session.run_in_transaction(
unit_of_work, "abc", timeout_secs=1
)
else:
with mock.patch("time.sleep") as sleep_mock:
with self.assertRaises(Aborted):
session.run_in_transaction(unit_of_work, "abc", timeout_secs=1)
sleep_mock.assert_not_called()
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIsInstance(txn, Transaction)
self.assertIsNone(txn.committed)
self.assertEqual(args, ("abc",))
self.assertEqual(kw, {})
expected_options = TransactionOptions(read_write=TransactionOptions.ReadWrite())
gax_api.begin_transaction.assert_called_once_with(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
gax_api.commit.assert_called_once_with(
self.SESSION_NAME,
mutations=txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def execute_partitioned_dml(self,
dml,
params=None,
param_types=None,
query_options=None):
"""Execute a partitionable DML statement.
:type dml: str
:param dml: DML statement
:type params: dict, {str -> column value}
:param params: values for parameter replacement. Keys must match
the names used in ``dml``.
:type param_types: dict[str -> Union[dict, .types.Type]]
:param param_types:
(Optional) maps explicit types for one or more param values;
required if parameters are passed.
:type query_options:
:class:`~google.cloud.spanner_v1.proto.ExecuteSqlRequest.QueryOptions`
or :class:`dict`
:param query_options:
(Optional) Query optimizer configuration to use for the given query.
If a dict is provided, it must be of the same form as the protobuf
message :class:`~google.cloud.spanner_v1.types.QueryOptions`
:rtype: int
:returns: Count of rows affected by the DML statement.
"""
query_options = _merge_query_options(
self._instance._client._query_options, query_options)
if params is not None:
if param_types is None:
raise ValueError(
"Specify 'param_types' when passing 'params'.")
params_pb = Struct(fields={
key: _make_value_pb(value)
for key, value in params.items()
})
else:
params_pb = None
api = self.spanner_api
txn_options = TransactionOptions(
partitioned_dml=TransactionOptions.PartitionedDml())
metadata = _metadata_with_prefix(self.name)
def execute_pdml():
with SessionCheckout(self._pool) as session:
txn = api.begin_transaction(session.name,
txn_options,
metadata=metadata)
txn_selector = TransactionSelector(id=txn.id)
restart = functools.partial(
api.execute_streaming_sql,
session.name,
dml,
transaction=txn_selector,
params=params_pb,
param_types=param_types,
query_options=query_options,
metadata=metadata,
)
iterator = _restart_on_unavailable(restart)
result_set = StreamedResultSet(iterator)
list(result_set) # consume all partials
return result_set.stats.row_count_lower_bound
retry_config = api._method_configs["ExecuteStreamingSql"].retry
return _retry_on_aborted(execute_pdml, retry_config)()
def _read_helper(self, multi_use, first=True, count=0, partition=None):
from google.protobuf.struct_pb2 import Struct
from google.cloud.spanner_v1.proto.result_set_pb2 import (
PartialResultSet, ResultSetMetadata, ResultSetStats)
from google.cloud.spanner_v1.proto.transaction_pb2 import (
TransactionSelector, TransactionOptions)
from google.cloud.spanner_v1.proto.type_pb2 import Type, StructType
from google.cloud.spanner_v1.proto.type_pb2 import STRING, INT64
from google.cloud.spanner_v1.keyset import KeySet
from google.cloud.spanner_v1._helpers import _make_value_pb
VALUES = [
[u'bharney', 31],
[u'phred', 32],
]
VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES]
struct_type_pb = StructType(fields=[
StructType.Field(name='name', type=Type(code=STRING)),
StructType.Field(name='age', type=Type(code=INT64)),
])
metadata_pb = ResultSetMetadata(row_type=struct_type_pb)
stats_pb = ResultSetStats(query_stats=Struct(
fields={
'rows_returned': _make_value_pb(2),
}))
result_sets = [
PartialResultSet(values=VALUE_PBS[0], metadata=metadata_pb),
PartialResultSet(values=VALUE_PBS[1], stats=stats_pb),
]
KEYS = [['*****@*****.**'], ['*****@*****.**']]
keyset = KeySet(keys=KEYS)
INDEX = 'email-address-index'
LIMIT = 20
database = _Database()
api = database.spanner_api = self._make_spanner_api()
api.streaming_read.return_value = _MockIterator(*result_sets)
session = _Session(database)
derived = self._makeDerived(session)
derived._multi_use = multi_use
derived._read_request_count = count
if not first:
derived._transaction_id = TXN_ID
if partition is not None: # 'limit' and 'partition' incompatible
result_set = derived.read(TABLE_NAME,
COLUMNS,
keyset,
index=INDEX,
partition=partition)
else:
result_set = derived.read(TABLE_NAME,
COLUMNS,
keyset,
index=INDEX,
limit=LIMIT)
self.assertEqual(derived._read_request_count, count + 1)
if multi_use:
self.assertIs(result_set._source, derived)
else:
self.assertIsNone(result_set._source)
self.assertEqual(list(result_set), VALUES)
self.assertEqual(result_set.metadata, metadata_pb)
self.assertEqual(result_set.stats, stats_pb)
txn_options = TransactionOptions(read_only=TransactionOptions.ReadOnly(
strong=True))
if multi_use:
if first:
expected_transaction = TransactionSelector(begin=txn_options)
else:
expected_transaction = TransactionSelector(id=TXN_ID)
else:
expected_transaction = TransactionSelector(single_use=txn_options)
if partition is not None:
expected_limit = 0
else:
expected_limit = LIMIT
api.streaming_read.assert_called_once_with(
self.SESSION_NAME,
TABLE_NAME,
COLUMNS,
keyset._to_pb(),
transaction=expected_transaction,
index=INDEX,
limit=expected_limit,
partition_token=partition,
metadata=[('google-cloud-resource-prefix', database.name)],
)
def test_run_in_transaction_w_timeout(self):
from google.api_core.exceptions import Aborted
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB,
TransactionOptions,
)
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = "citizens"
COLUMNS = ["email", "first_name", "last_name", "age"]
VALUES = [
["*****@*****.**", "Phred", "Phlyntstone", 32],
["*****@*****.**", "Bharney", "Rhubble", 31],
]
TRANSACTION_ID = b"FACEDACE"
transaction_pb = TransactionPB(id=TRANSACTION_ID)
aborted = _make_rpc_error(Aborted, trailing_metadata=[])
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = aborted
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
# retry once w/ timeout_secs=1
def _time(_results=[1, 1.5, 2.5]):
return _results.pop(0)
with mock.patch("time.time", _time):
with mock.patch("time.sleep") as sleep_mock:
with self.assertRaises(Aborted):
session.run_in_transaction(unit_of_work, timeout_secs=1)
sleep_mock.assert_not_called()
self.assertEqual(len(called_with), 2)
for txn, args, kw in called_with:
self.assertIsInstance(txn, Transaction)
self.assertIsNone(txn.committed)
self.assertEqual(args, ())
self.assertEqual(kw, {})
expected_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite())
self.assertEqual(
gax_api.begin_transaction.call_args_list,
[
mock.call(
self.SESSION_NAME,
expected_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
] * 2,
)
self.assertEqual(
gax_api.commit.call_args_list,
[
mock.call(
self.SESSION_NAME,
txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[("google-cloud-resource-prefix", database.name)],
)
] * 2,
)
def test_run_in_transaction_w_abort_w_retry_metadata(self):
import datetime
from google.api_core.exceptions import Aborted
from google.protobuf.duration_pb2 import Duration
from google.rpc.error_details_pb2 import RetryInfo
from google.cloud.spanner_v1.proto.spanner_pb2 import CommitResponse
from google.cloud.spanner_v1.proto.transaction_pb2 import (
Transaction as TransactionPB, TransactionOptions)
from google.cloud._helpers import UTC
from google.cloud._helpers import _datetime_to_pb_timestamp
from google.cloud.spanner_v1.transaction import Transaction
TABLE_NAME = 'citizens'
COLUMNS = ['email', 'first_name', 'last_name', 'age']
VALUES = [
['*****@*****.**', 'Phred', 'Phlyntstone', 32],
['*****@*****.**', 'Bharney', 'Rhubble', 31],
]
TRANSACTION_ID = b'FACEDACE'
RETRY_SECONDS = 12
RETRY_NANOS = 3456
retry_info = RetryInfo(
retry_delay=Duration(seconds=RETRY_SECONDS, nanos=RETRY_NANOS))
trailing_metadata = [
('google.rpc.retryinfo-bin', retry_info.SerializeToString()),
]
aborted = _make_rpc_error(
Aborted,
trailing_metadata=trailing_metadata,
)
transaction_pb = TransactionPB(id=TRANSACTION_ID)
now = datetime.datetime.utcnow().replace(tzinfo=UTC)
now_pb = _datetime_to_pb_timestamp(now)
response = CommitResponse(commit_timestamp=now_pb)
gax_api = self._make_spanner_api()
gax_api.begin_transaction.return_value = transaction_pb
gax_api.commit.side_effect = [aborted, response]
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
with mock.patch('time.sleep') as sleep_mock:
session.run_in_transaction(unit_of_work, 'abc', some_arg='def')
sleep_mock.assert_called_once_with(RETRY_SECONDS + RETRY_NANOS / 1.0e9)
self.assertEqual(len(called_with), 2)
for index, (txn, args, kw) in enumerate(called_with):
self.assertIsInstance(txn, Transaction)
if index == 1:
self.assertEqual(txn.committed, now)
else:
self.assertIsNone(txn.committed)
self.assertEqual(args, ('abc', ))
self.assertEqual(kw, {'some_arg': 'def'})
expected_options = TransactionOptions(
read_write=TransactionOptions.ReadWrite(), )
self.assertEqual(gax_api.begin_transaction.call_args_list, [
mock.call(
self.SESSION_NAME,
expected_options,
metadata=[('google-cloud-resource-prefix', database.name)],
)
] * 2)
self.assertEqual(gax_api.commit.call_args_list, [
mock.call(
self.SESSION_NAME,
txn._mutations,
transaction_id=TRANSACTION_ID,
metadata=[('google-cloud-resource-prefix', database.name)],
)
] * 2)
def _execute_sql_helper(
self,
multi_use,
first=True,
count=0,
partition=None,
sql_count=0,
query_options=None,
timeout=google.api_core.gapic_v1.method.DEFAULT,
retry=google.api_core.gapic_v1.method.DEFAULT,
):
from google.protobuf.struct_pb2 import Struct
from google.cloud.spanner_v1.proto.result_set_pb2 import (
PartialResultSet,
ResultSetMetadata,
ResultSetStats,
)
from google.cloud.spanner_v1.proto.transaction_pb2 import (
TransactionSelector,
TransactionOptions,
)
from google.cloud.spanner_v1.proto.type_pb2 import Type, StructType
from google.cloud.spanner_v1.proto.type_pb2 import STRING, INT64
from google.cloud.spanner_v1._helpers import (
_make_value_pb,
_merge_query_options,
)
VALUES = [[u"bharney", u"rhubbyl", 31], [u"phred", u"phlyntstone", 32]]
VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES]
MODE = 2 # PROFILE
struct_type_pb = StructType(
fields=[
StructType.Field(name="first_name", type=Type(code=STRING)),
StructType.Field(name="last_name", type=Type(code=STRING)),
StructType.Field(name="age", type=Type(code=INT64)),
]
)
metadata_pb = ResultSetMetadata(row_type=struct_type_pb)
stats_pb = ResultSetStats(
query_stats=Struct(fields={"rows_returned": _make_value_pb(2)})
)
result_sets = [
PartialResultSet(values=VALUE_PBS[0], metadata=metadata_pb),
PartialResultSet(values=VALUE_PBS[1], stats=stats_pb),
]
iterator = _MockIterator(*result_sets)
database = _Database()
api = database.spanner_api = self._make_spanner_api()
api.execute_streaming_sql.return_value = iterator
session = _Session(database)
derived = self._makeDerived(session)
derived._multi_use = multi_use
derived._read_request_count = count
derived._execute_sql_count = sql_count
if not first:
derived._transaction_id = TXN_ID
result_set = derived.execute_sql(
SQL_QUERY_WITH_PARAM,
PARAMS,
PARAM_TYPES,
query_mode=MODE,
query_options=query_options,
partition=partition,
retry=retry,
timeout=timeout,
)
self.assertEqual(derived._read_request_count, count + 1)
if multi_use:
self.assertIs(result_set._source, derived)
else:
self.assertIsNone(result_set._source)
self.assertEqual(list(result_set), VALUES)
self.assertEqual(result_set.metadata, metadata_pb)
self.assertEqual(result_set.stats, stats_pb)
txn_options = TransactionOptions(
read_only=TransactionOptions.ReadOnly(strong=True)
)
if multi_use:
if first:
expected_transaction = TransactionSelector(begin=txn_options)
else:
expected_transaction = TransactionSelector(id=TXN_ID)
else:
expected_transaction = TransactionSelector(single_use=txn_options)
expected_params = Struct(
fields={key: _make_value_pb(value) for (key, value) in PARAMS.items()}
)
expected_query_options = database._instance._client._query_options
if query_options:
expected_query_options = _merge_query_options(
expected_query_options, query_options
)
api.execute_streaming_sql.assert_called_once_with(
self.SESSION_NAME,
SQL_QUERY_WITH_PARAM,
transaction=expected_transaction,
params=expected_params,
param_types=PARAM_TYPES,
query_mode=MODE,
query_options=expected_query_options,
partition_token=partition,
seqno=sql_count,
metadata=[("google-cloud-resource-prefix", database.name)],
timeout=timeout,
retry=retry,
)
self.assertEqual(derived._execute_sql_count, sql_count + 1)
self.assertSpanAttributes(
"CloudSpanner.ReadWriteTransaction",
status=StatusCanonicalCode.OK,
attributes=dict(BASE_ATTRIBUTES, **{"db.statement": SQL_QUERY_WITH_PARAM}),
)
def _execute_sql_helper(self,
multi_use,
first=True,
count=0,
partition=None,
sql_count=0):
from google.protobuf.struct_pb2 import Struct
from google.cloud.spanner_v1.proto.result_set_pb2 import (
PartialResultSet, ResultSetMetadata, ResultSetStats)
from google.cloud.spanner_v1.proto.transaction_pb2 import (
TransactionSelector, TransactionOptions)
from google.cloud.spanner_v1.proto.type_pb2 import Type, StructType
from google.cloud.spanner_v1.proto.type_pb2 import STRING, INT64
from google.cloud.spanner_v1._helpers import _make_value_pb
VALUES = [
[u'bharney', u'rhubbyl', 31],
[u'phred', u'phlyntstone', 32],
]
VALUE_PBS = [[_make_value_pb(item) for item in row] for row in VALUES]
MODE = 2 # PROFILE
struct_type_pb = StructType(fields=[
StructType.Field(name='first_name', type=Type(code=STRING)),
StructType.Field(name='last_name', type=Type(code=STRING)),
StructType.Field(name='age', type=Type(code=INT64)),
])
metadata_pb = ResultSetMetadata(row_type=struct_type_pb)
stats_pb = ResultSetStats(query_stats=Struct(
fields={
'rows_returned': _make_value_pb(2),
}))
result_sets = [
PartialResultSet(values=VALUE_PBS[0], metadata=metadata_pb),
PartialResultSet(values=VALUE_PBS[1], stats=stats_pb),
]
iterator = _MockIterator(*result_sets)
database = _Database()
api = database.spanner_api = self._make_spanner_api()
api.execute_streaming_sql.return_value = iterator
session = _Session(database)
derived = self._makeDerived(session)
derived._multi_use = multi_use
derived._read_request_count = count
derived._execute_sql_count = sql_count
if not first:
derived._transaction_id = TXN_ID
result_set = derived.execute_sql(SQL_QUERY_WITH_PARAM,
PARAMS,
PARAM_TYPES,
query_mode=MODE,
partition=partition)
self.assertEqual(derived._read_request_count, count + 1)
if multi_use:
self.assertIs(result_set._source, derived)
else:
self.assertIsNone(result_set._source)
self.assertEqual(list(result_set), VALUES)
self.assertEqual(result_set.metadata, metadata_pb)
self.assertEqual(result_set.stats, stats_pb)
txn_options = TransactionOptions(read_only=TransactionOptions.ReadOnly(
strong=True))
if multi_use:
if first:
expected_transaction = TransactionSelector(begin=txn_options)
else:
expected_transaction = TransactionSelector(id=TXN_ID)
else:
expected_transaction = TransactionSelector(single_use=txn_options)
expected_params = Struct(fields={
key: _make_value_pb(value)
for (key, value) in PARAMS.items()
})
api.execute_streaming_sql.assert_called_once_with(
self.SESSION_NAME,
SQL_QUERY_WITH_PARAM,
transaction=expected_transaction,
params=expected_params,
param_types=PARAM_TYPES,
query_mode=MODE,
partition_token=partition,
seqno=sql_count,
metadata=[('google-cloud-resource-prefix', database.name)],
)
self.assertEqual(derived._execute_sql_count, sql_count + 1)
