def upsert(
transaction: spanner_transaction.Transaction,
table_name: str,
columns: Iterable[str],
values: Iterable[Iterable[Any]],
) -> None:
"""Inserts or updates rows in the given table based on the provided values.
All non-nullable columns must be specified, similarly to the insert method.
The presence or absence of data in the table will not cause an exception
to be thrown, unlike insert or update.
Args:
transaction: The Spanner transaction to execute the request on
table_name: The Spanner table being modified
columns: Which columns to write on the Spanner table
values: A list of rows to write to the table. The order of the values in
each sublist must match the order of the columns specified in the
`columns` parameter.
"""
_logger.debug("Upsert table=%s columns=%s values=%s", table_name, columns,
values)
transaction.insert_or_update(table=table_name,
columns=columns,
values=values)
def delete(transaction: spanner_transaction.Transaction, table_name: str,
keyset: spanner.KeySet) -> None:
"""Deletes rows from the given table based on the provided KeySet.
Args:
transaction: The Spanner transaction to execute the request on
table_name: The Spanner table being modified
keyset: Contains a list of primary keys that indicates which rows to delete
from the Spanner table
"""
_logger.debug('Delete table=%s keys=%s', table_name, keyset.keys)
transaction.delete(table=table_name, keyset=keyset)
def test_run_in_transaction_w_commit_error(self):
from google.gax.errors import GaxError
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],
]
gax_api = _SpannerApi(_commit_error=True, )
database = _Database(self.DATABASE_NAME)
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = 'DEADBEEF'
begun_txn = session._transaction = Transaction(session)
begun_txn._transaction_id = b'FACEDACE'
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
with self.assertRaises(GaxError):
session.run_in_transaction(unit_of_work)
self.assertIsNone(session._transaction)
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIs(txn, begun_txn)
self.assertEqual(txn.committed, None)
self.assertEqual(args, ())
self.assertEqual(kw, {})
def sql_query(
transaction: spanner_transaction.Transaction, query: str,
parameters: Dict[str, Any],
parameter_types: Dict[str, type_pb2.Type]) -> List[Iterable[Any]]:
"""Executes a given SQL query against the Spanner database.
This isn't technically read-only, but it's necessary to implement the read-
only features of the ORM
Args:
transaction: The Spanner transaction to execute the request on
query: The SQL query to run
parameters: A mapping from the names of the parameters used in the SQL query
to the value to be substituted in for that parameter
parameter_types: A mapping from the names of the parameters used in the SQL
query to the type of the value being substituted in for that parameter
Returns:
A list of lists. Each sublist is a result row from the SQL query. For
SELECT queries, the order of values in the sublist matches the order
of the columns requested from the SELECT clause of the query.
"""
_logger.debug('Executing SQL:\n%s\n%s\n%s', query, parameters,
parameter_types)
stream_results = transaction.execute_sql(query,
params=parameters,
param_types=parameter_types)
return list(stream_results)
def find(
transaction: spanner_transaction.Transaction,
table_name: str,
columns: Iterable[str],
keyset: spanner.KeySet,
) -> List[Iterable[Any]]:
"""Retrieves rows from the given table based on the provided KeySet.
Args:
transaction: The Spanner transaction to execute the request on
table_name: The Spanner table being queried
columns: Which columns to retrieve from the Spanner table
keyset: Contains a list of primary keys that indicates which rows to
retrieve from the Spanner table
Returns:
A list of lists. Each sublist is the set of `columns` requested from
a row in the Spanner table whose primary key matches one of the
primary keys in the `keyset`. The order of the values in the sublist
matches the order of the columns from the `columns` parameter.
"""
_logger.debug("Find table=%s columns=%s keys=%s", table_name, columns,
keyset.keys)
stream_results = transaction.read(table=table_name,
columns=columns,
keyset=keyset)
return list(stream_results)
def insert(transaction: spanner_transaction.Transaction, table_name: str,
columns: Iterable[str], values: Iterable[Iterable[Any]]) -> None:
"""Adds rows to the given table based on the provided values.
All non-nullable columns must be specified. Note that if a row is specified
for which the primary key already exists in the table, an exception will
be thrown and the insert will be aborted.
Args:
transaction: The Spanner transaction to execute the request on
table_name: The Spanner table being modified
columns: Which columns to write on the Spanner table
values: A list of rows to write to the table. The order of the values in
each sublist must match the order of the columns specified in the
`columns` parameter.
"""
_logger.debug('Insert table=%s columns=%s values=%s', table_name, columns,
values)
transaction.insert(table=table_name, columns=columns, values=values)
def update(
transaction: spanner_transaction.Transaction,
table_name: str,
columns: Iterable[str],
values: Iterable[Iterable[Any]],
) -> None:
"""Updates rows in the given table based on the provided values.
Note that if a row is specified for which the primary key does not
exist in the table, an exception will be thrown and the update
will be aborted.
Args:
transaction: The Spanner transaction to execute the request on
table_name: The Spanner table being modified
columns: Which columns to write on the Spanner table
values: A list of rows to write to the table. The order of the values in
each sublist must match the order of the columns specified in the
`columns` parameter.
"""
_logger.debug("Update table=%s columns=%s values=%s", table_name, columns,
values)
transaction.update(table=table_name, columns=columns, values=values)
def transaction(self):
"""Create a transaction to perform a set of reads with shared staleness.
:rtype: :class:`~google.cloud.spanner_v1.transaction.Transaction`
:returns: a transaction bound to this session
:raises ValueError: if the session has not yet been created.
"""
if self._session_id is None:
raise ValueError("Session has not been created.")
if self._transaction is not None:
self._transaction._rolled_back = True
del self._transaction
txn = self._transaction = Transaction(self)
return txn
def test_run_in_transaction_w_commit_error(self):
from google.api_core.exceptions import Unknown
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"
gax_api = self._make_spanner_api()
gax_api.commit.side_effect = Unknown("error")
database = self._make_database()
database.spanner_api = gax_api
session = self._make_one(database)
session._session_id = self.SESSION_ID
begun_txn = session._transaction = Transaction(session)
begun_txn._transaction_id = TRANSACTION_ID
assert session._transaction._transaction_id
called_with = []
def unit_of_work(txn, *args, **kw):
called_with.append((txn, args, kw))
txn.insert(TABLE_NAME, COLUMNS, VALUES)
with self.assertRaises(Unknown):
session.run_in_transaction(unit_of_work)
self.assertIsNone(session._transaction)
self.assertEqual(len(called_with), 1)
txn, args, kw = called_with[0]
self.assertIs(txn, begun_txn)
self.assertEqual(txn.committed, None)
self.assertEqual(args, ())
self.assertEqual(kw, {})
gax_api.begin_transaction.assert_not_called()
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.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 _execute_sql_in_transaction(transaction: Transaction,
queries: List[str]):
for sql in queries:
transaction.execute_update(sql)
