def _partition_read_helper(
self, multi_use, w_txn,
size=None, max_partitions=None, index=None):
from google.cloud.spanner_v1.keyset import KeySet
from google.cloud.spanner_v1.types import Partition
from google.cloud.spanner_v1.types import PartitionOptions
from google.cloud.spanner_v1.types import PartitionResponse
from google.cloud.spanner_v1.types import Transaction
from google.cloud.spanner_v1.proto.transaction_pb2 import (
TransactionSelector)
keyset = KeySet(all_=True)
new_txn_id = b'ABECAB91'
token_1 = b'FACE0FFF'
token_2 = b'BADE8CAF'
response = PartitionResponse(
partitions=[
Partition(partition_token=token_1),
Partition(partition_token=token_2),
],
transaction=Transaction(id=new_txn_id),
)
database = _Database()
api = database.spanner_api = _FauxSpannerAPI(
_partition_read_response=response)
session = _Session(database)
derived = self._makeDerived(session)
derived._multi_use = multi_use
if w_txn:
derived._transaction_id = TXN_ID
tokens = list(derived.partition_read(
TABLE_NAME, COLUMNS, keyset,
index=index,
partition_size_bytes=size,
max_partitions=max_partitions,
))
self.assertEqual(tokens, [token_1, token_2])
(r_session, table, key_set, transaction, r_index, columns,
partition_options, metadata) = api._partition_read_with
self.assertEqual(r_session, self.SESSION_NAME)
self.assertEqual(table, TABLE_NAME)
self.assertEqual(key_set, keyset._to_pb())
self.assertIsInstance(transaction, TransactionSelector)
self.assertEqual(transaction.id, TXN_ID)
self.assertFalse(transaction.HasField('begin'))
self.assertEqual(r_index, index)
self.assertEqual(columns, COLUMNS)
self.assertEqual(
partition_options,
PartitionOptions(
partition_size_bytes=size, max_partitions=max_partitions))
self.assertEqual(
metadata, [('google-cloud-resource-prefix', database.name)])
def process_read_batch(self, batch):
"""Process a single, partitioned read.
:type batch: mapping
:param batch:
one of the mappings returned from an earlier call to
:meth:`generate_read_batches`.
:rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
:returns: a result set instance which can be used to consume rows.
"""
kwargs = copy.deepcopy(batch["read"])
keyset_dict = kwargs.pop("keyset")
kwargs["keyset"] = KeySet._from_dict(keyset_dict)
return self._get_snapshot().read(partition=batch["partition"], **kwargs)
def _read_helper(self, multi_use, first=True, 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)
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
TXN_ID = b'DEADBEEF'
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
TOKEN = b'DEADBEEF'
database = _Database()
api = database.spanner_api = _FauxSpannerAPI(
_streaming_read_response=_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
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)
(r_session, table, columns, key_set, transaction, index,
limit, resume_token, metadata) = api._streaming_read_with
self.assertEqual(r_session, self.SESSION_NAME)
self.assertEqual(table, TABLE_NAME)
self.assertEqual(columns, COLUMNS)
self.assertEqual(key_set, KEYSET._to_pb())
self.assertIsInstance(transaction, TransactionSelector)
if multi_use:
if first:
self.assertTrue(transaction.begin.read_only.strong)
else:
self.assertEqual(transaction.id, TXN_ID)
else:
self.assertTrue(transaction.single_use.read_only.strong)
self.assertEqual(index, INDEX)
self.assertEqual(limit, LIMIT)
self.assertEqual(resume_token, b'')
self.assertEqual(
metadata, [('google-cloud-resource-prefix', database.name)])
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)
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 = _FauxSpannerAPI(
_streaming_read_response=_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)
(r_session, table, columns, key_set, transaction, index, limit,
resume_token, r_partition, metadata) = api._streaming_read_with
self.assertEqual(r_session, self.SESSION_NAME)
self.assertEqual(table, TABLE_NAME)
self.assertEqual(columns, COLUMNS)
self.assertEqual(key_set, keyset._to_pb())
self.assertIsInstance(transaction, TransactionSelector)
if multi_use:
if first:
self.assertTrue(transaction.begin.read_only.strong)
else:
self.assertEqual(transaction.id, TXN_ID)
else:
self.assertTrue(transaction.single_use.read_only.strong)
self.assertEqual(index, INDEX)
if partition is not None:
self.assertEqual(limit, 0)
self.assertEqual(r_partition, partition)
else:
self.assertEqual(limit, LIMIT)
self.assertIsNone(r_partition)
self.assertEqual(resume_token, b'')
self.assertEqual(
metadata, [('google-cloud-resource-prefix', database.name)])