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 = self._read_timestamp elif self._min_read_timestamp: key = "min_read_timestamp" value = self._min_read_timestamp elif self._max_staleness: key = "max_staleness" value = self._max_staleness elif self._exact_staleness: key = "exact_staleness" value = 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_begin_ok_exact_strong(self): from google.cloud.spanner_v1 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=session.name, options=expected_txn_options, metadata=[("google-cloud-resource-prefix", database.name)], ) self.assertSpanAttributes( "CloudSpanner.BeginTransaction", status=StatusCanonicalCode.OK, attributes=BASE_ATTRIBUTES, )
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()) with trace_call("CloudSpanner.BeginTransaction", self._session): response = api.begin_transaction(session=self._session.name, options=txn_options, metadata=metadata) self._transaction_id = response.id return self._transaction_id
def commit(self, return_commit_stats=False): """Commit mutations to the database. :type return_commit_stats: bool :param return_commit_stats: If true, the response will return commit stats which can be accessed though commit_stats. :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)} request = CommitRequest( session=self._session.name, mutations=self._mutations, single_use_transaction=txn_options, return_commit_stats=return_commit_stats, ) with trace_call("CloudSpanner.Commit", self._session, trace_attributes): response = api.commit( request=request, metadata=metadata, ) self.committed = response.commit_timestamp self.commit_stats = response.commit_stats return self.committed
def test_run_in_transaction_callback_raises_non_gax_error(self): from google.cloud.spanner_v1 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) gax_api.rollback.assert_called_once_with( session=self.SESSION_NAME, transaction_id=TRANSACTION_ID, metadata=[("google-cloud-resource-prefix", database.name)], )
def _make_txn_selector(self): from google.cloud.spanner_v1 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 commit(self, return_commit_stats=False, request_options=None): """Commit mutations to the database. :type return_commit_stats: bool :param return_commit_stats: If true, the response will return commit stats which can be accessed though commit_stats. :type request_options: :class:`google.cloud.spanner_v1.types.RequestOptions` :param request_options: (Optional) Common options for this request. If a dict is provided, it must be of the same form as the protobuf message :class:`~google.cloud.spanner_v1.types.RequestOptions`. :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)} if request_options is None: request_options = RequestOptions() elif type(request_options) == dict: request_options = RequestOptions(request_options) request_options.transaction_tag = self.transaction_tag # Request tags are not supported for commit requests. request_options.request_tag = None request = CommitRequest( session=self._session.name, mutations=self._mutations, single_use_transaction=txn_options, return_commit_stats=return_commit_stats, request_options=request_options, ) with trace_call("CloudSpanner.Commit", self._session, trace_attributes): response = api.commit( request=request, metadata=metadata, ) self.committed = response.commit_timestamp self.commit_stats = response.commit_stats 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( session=self._session.name, mutations=self._mutations, single_use_transaction=txn_options, metadata=metadata, ) self.committed = response.commit_timestamp return self.committed
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.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.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: from google.cloud.spanner_v1.transaction import Transaction if param_types is None: raise ValueError("Specify 'param_types' when passing 'params'.") params_pb = Transaction._make_params_pb(params, param_types) else: params_pb = {} 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=session.name, options=txn_options, metadata=metadata ) txn_selector = TransactionSelector(id=txn.id) request = ExecuteSqlRequest( session=session.name, sql=dml, transaction=txn_selector, params=params_pb, param_types=param_types, query_options=query_options, ) restart = functools.partial( api.execute_streaming_sql, request=request, 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 return _retry_on_aborted(execute_pdml, DEFAULT_RETRY_BACKOFF)()
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 import CommitResponse from google.cloud.spanner_v1 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) ] * 2, ) self.assertEqual( gax_api.commit.call_args_list, [ mock.call( session=self.SESSION_NAME, mutations=txn._mutations, transaction_id=TRANSACTION_ID, metadata=[("google-cloud-resource-prefix", database.name)], ) ] * 2, )
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 import CommitResponse from google.cloud.spanner_v1 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) ] * 2, ) self.assertEqual( gax_api.commit.call_args_list, [ mock.call( session=self.SESSION_NAME, mutations=txn._mutations, transaction_id=TRANSACTION_ID, metadata=[("google-cloud-resource-prefix", database.name)], ) ] * 2, )
def test_run_in_transaction_w_args_w_kwargs_wo_abort(self): import datetime from google.cloud.spanner_v1 import CommitResponse from google.cloud.spanner_v1 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) gax_api.commit.assert_called_once_with( session=self.SESSION_NAME, mutations=txn._mutations, transaction_id=TRANSACTION_ID, 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 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) ] * 3, ) self.assertEqual( gax_api.commit.call_args_list, [ mock.call( session=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 import CommitResponse from google.cloud.spanner_v1 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( session=self.SESSION_NAME, options=expected_options, metadata=[("google-cloud-resource-prefix", database.name)], ) gax_api.commit.assert_called_once_with( session=self.SESSION_NAME, mutations=txn._mutations, transaction_id=TRANSACTION_ID, metadata=[("google-cloud-resource-prefix", database.name)], )
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 import ( PartialResultSet, ResultSetMetadata, ResultSetStats, ) from google.cloud.spanner_v1 import ( TransactionSelector, TransactionOptions, ) from google.cloud.spanner_v1 import ExecuteSqlRequest from google.cloud.spanner_v1 import Type, StructType from google.cloud.spanner_v1 import TypeCode from google.cloud.spanner_v1._helpers import ( _make_value_pb, _merge_query_options, ) VALUES = [[u"bharney", u"rhubbyl", 31], [u"phred", u"phlyntstone", 32]] MODE = 2 # PROFILE struct_type_pb = StructType( fields=[ StructType.Field(name="first_name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="last_name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.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(metadata=metadata_pb), PartialResultSet(stats=stats_pb), ] for i in range(len(result_sets)): result_sets[i].values.extend(VALUES[i]) 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 ) expected_request = ExecuteSqlRequest( session=self.SESSION_NAME, sql=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, ) api.execute_streaming_sql.assert_called_once_with( request=expected_request, 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 _read_helper(self, multi_use, first=True, count=0, partition=None): from google.protobuf.struct_pb2 import Struct from google.cloud.spanner_v1 import ( PartialResultSet, ResultSetMetadata, ResultSetStats, ) from google.cloud.spanner_v1 import ( TransactionSelector, TransactionOptions, ) from google.cloud.spanner_v1 import ReadRequest from google.cloud.spanner_v1 import Type, StructType from google.cloud.spanner_v1 import TypeCode 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]] struct_type_pb = StructType( fields=[ StructType.Field(name="name", type_=Type(code=TypeCode.STRING)), StructType.Field(name="age", type_=Type(code=TypeCode.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(metadata=metadata_pb), PartialResultSet(stats=stats_pb), ] for i in range(len(result_sets)): result_sets[i].values.extend(VALUES[i]) 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 expected_request = ReadRequest( session=self.SESSION_NAME, table=TABLE_NAME, columns=COLUMNS, key_set=keyset._to_pb(), transaction=expected_transaction, index=INDEX, limit=expected_limit, partition_token=partition, ) api.streaming_read.assert_called_once_with( request=expected_request, metadata=[("google-cloud-resource-prefix", database.name)], ) self.assertSpanAttributes( "CloudSpanner.ReadOnlyTransaction", attributes=dict( BASE_ATTRIBUTES, table_id=TABLE_NAME, columns=tuple(COLUMNS) ), )