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(
self,
sql,
params=None,
param_types=None,
partition_size_bytes=None,
max_partitions=None,
):
"""Perform a ``ParitionQuery`` API request.
:type sql: str
:param sql: SQL query statement
:type params: dict, {str -> column value}
:param params: values for parameter replacement. Keys must match
the names used in ``sql``.
: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 partition_size_bytes: int
:param partition_size_bytes:
(Optional) desired size for each partition generated. The service
uses this as a hint, the actual partition size may differ.
:type max_partitions: int
:param max_partitions:
(Optional) desired maximum number of partitions generated. The
service uses this as a hint, the actual number of partitions may
differ.
:rtype: iterable of bytes
:returns: a sequence of partition tokens
:raises ValueError:
for single-use snapshots, or if a transaction ID is
already associtated with the snapshot.
"""
if not self._multi_use:
raise ValueError("Cannot use single-use snapshot.")
if self._transaction_id is None:
raise ValueError("Transaction not started.")
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
database = self._session._database
api = database.spanner_api
metadata = _metadata_with_prefix(database.name)
transaction = self._make_txn_selector()
partition_options = PartitionOptions(
partition_size_bytes=partition_size_bytes,
max_partitions=max_partitions)
trace_attributes = {"db.statement": sql}
with trace_call(
"CloudSpanner.PartitionReadWriteTransaction",
self._session,
trace_attributes,
):
response = api.partition_query(
session=self._session.name,
sql=sql,
transaction=transaction,
params=params_pb,
param_types=param_types,
partition_options=partition_options,
metadata=metadata,
)
return [partition.partition_token for partition in response.partitions]
def partition_read(
self,
table,
columns,
keyset,
index="",
partition_size_bytes=None,
max_partitions=None,
):
"""Perform a ``ParitionRead`` API request for rows in a table.
:type table: str
:param table: name of the table from which to fetch data
:type columns: list of str
:param columns: names of columns to be retrieved
:type keyset: :class:`~google.cloud.spanner_v1.keyset.KeySet`
:param keyset: keys / ranges identifying rows to be retrieved
:type index: str
:param index: (Optional) name of index to use, rather than the
table's primary key
:type partition_size_bytes: int
:param partition_size_bytes:
(Optional) desired size for each partition generated. The service
uses this as a hint, the actual partition size may differ.
:type max_partitions: int
:param max_partitions:
(Optional) desired maximum number of partitions generated. The
service uses this as a hint, the actual number of partitions may
differ.
:rtype: iterable of bytes
:returns: a sequence of partition tokens
:raises ValueError:
for single-use snapshots, or if a transaction ID is
already associtated with the snapshot.
"""
if not self._multi_use:
raise ValueError("Cannot use single-use snapshot.")
if self._transaction_id is None:
raise ValueError("Transaction not started.")
database = self._session._database
api = database.spanner_api
metadata = _metadata_with_prefix(database.name)
transaction = self._make_txn_selector()
partition_options = PartitionOptions(
partition_size_bytes=partition_size_bytes,
max_partitions=max_partitions)
trace_attributes = {"table_id": table, "columns": columns}
with trace_call("CloudSpanner.PartitionReadOnlyTransaction",
self._session, trace_attributes):
response = api.partition_read(
session=self._session.name,
table=table,
columns=columns,
key_set=keyset._to_pb(),
transaction=transaction,
index=index,
partition_options=partition_options,
metadata=metadata,
)
return [partition.partition_token for partition in response.partitions]
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)],
)