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 = self._make_spanner_api()
api.partition_read.return_value = 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])
expected_txn_selector = TransactionSelector(id=TXN_ID)
expected_partition_options = PartitionOptions(
partition_size_bytes=size, max_partitions=max_partitions)
api.partition_read.assert_called_once_with(
session=self.SESSION_NAME,
table=TABLE_NAME,
columns=COLUMNS,
key_set=keyset._to_pb(),
transaction=expected_txn_selector,
index=index,
partition_options=expected_partition_options,
metadata=[("google-cloud-resource-prefix", database.name)],
)
def _partition_query_helper(self,
multi_use,
w_txn,
size=None,
max_partitions=None):
from google.protobuf.struct_pb2 import Struct
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)
from google.cloud.spanner_v1._helpers import _make_value_pb
txn_id = b'DEADBEEF'
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_query_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_query(
SQL_QUERY_WITH_PARAM,
PARAMS,
PARAM_TYPES,
partition_size_bytes=size,
max_partitions=max_partitions,
))
self.assertEqual(tokens, [token_1, token_2])
(r_session, sql, transaction, params, param_types, partition_options,
metadata) = api._partition_query_with
self.assertEqual(r_session, self.SESSION_NAME)
self.assertEqual(sql, SQL_QUERY_WITH_PARAM)
self.assertIsInstance(transaction, TransactionSelector)
self.assertEqual(transaction.id, txn_id)
self.assertFalse(transaction.HasField('begin'))
expected_params = Struct(fields={
key: _make_value_pb(value)
for (key, value) in PARAMS.items()
})
self.assertEqual(params, expected_params)
self.assertEqual(param_types, PARAM_TYPES)
self.assertEqual(
partition_options,
PartitionOptions(partition_size_bytes=size,
max_partitions=max_partitions))
self.assertEqual(metadata,
[('google-cloud-resource-prefix', database.name)])
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)
txn_id = b'DEADBEEF'
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 _partition_query_helper(self,
multi_use,
w_txn,
size=None,
max_partitions=None):
from google.protobuf.struct_pb2 import Struct
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)
from google.cloud.spanner_v1._helpers import _make_value_pb
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 = self._make_spanner_api()
api.partition_query.return_value = 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_query(
SQL_QUERY_WITH_PARAM,
PARAMS,
PARAM_TYPES,
partition_size_bytes=size,
max_partitions=max_partitions,
))
self.assertEqual(tokens, [token_1, token_2])
expected_params = Struct(fields={
key: _make_value_pb(value)
for (key, value) in PARAMS.items()
})
expected_txn_selector = TransactionSelector(id=TXN_ID)
expected_partition_options = PartitionOptions(
partition_size_bytes=size, max_partitions=max_partitions)
api.partition_query.assert_called_once_with(
session=self.SESSION_NAME,
sql=SQL_QUERY_WITH_PARAM,
transaction=expected_txn_selector,
params=expected_params,
param_types=PARAM_TYPES,
partition_options=expected_partition_options,
metadata=[('google-cloud-resource-prefix', database.name)],
)