def test_read_w_ranges(self):
ROW_COUNT = 4000
START = 1000
END = 2000
session, committed = self._set_up_table(ROW_COUNT)
snapshot = session.snapshot(read_timestamp=committed, multi_use=True)
all_data_rows = list(self._row_data(ROW_COUNT))
closed_closed = KeyRange(start_closed=[START], end_closed=[END])
keyset = KeySet(ranges=(closed_closed, ))
rows = list(snapshot.read(self.TABLE, self.COLUMNS, keyset))
expected = all_data_rows[START:END + 1]
self._check_row_data(rows, expected)
closed_open = KeyRange(start_closed=[START], end_open=[END])
keyset = KeySet(ranges=(closed_open, ))
rows = list(snapshot.read(self.TABLE, self.COLUMNS, keyset))
expected = all_data_rows[START:END]
self._check_row_data(rows, expected)
open_open = KeyRange(start_open=[START], end_open=[END])
keyset = KeySet(ranges=(open_open, ))
rows = list(snapshot.read(self.TABLE, self.COLUMNS, keyset))
expected = all_data_rows[START + 1:END]
self._check_row_data(rows, expected)
open_closed = KeyRange(start_open=[START], end_closed=[END])
keyset = KeySet(ranges=(open_closed, ))
rows = list(snapshot.read(self.TABLE, self.COLUMNS, keyset))
expected = all_data_rows[START + 1:END + 1]
self._check_row_data(rows, expected)
def table(cls, table, columns, index="", keyset=None):
"""
A convenient method to construct ReadOperation from table.
Args:
table: name of the table from which to fetch data.
columns: names of columns to be retrieved.
index: (optional) name of index to use, rather than the table's primary
key.
keyset: (optional) `KeySet` keys / ranges identifying rows to be
retrieved.
"""
keyset = keyset or KeySet(all_=True)
if not isinstance(keyset, KeySet):
raise ValueError(
"keyset must be an instance of class "
"google.cloud.spanner.KeySet")
return cls(
is_sql=False,
is_table=True,
read_operation="process_read_batch",
kwargs={
'table': table,
'columns': columns,
'index': index,
'keyset': keyset
})
def test_transaction_read_and_insert_or_update_then_commit(self):
from google.cloud.spanner import KeySet
keyset = KeySet(all_=True)
retry = RetryInstanceState(_has_all_ddl)
retry(self._db.reload)()
session = self._db.session()
session.create()
self.to_delete.append(session)
with session.batch() as batch:
batch.delete(self.TABLE, keyset)
with session.transaction() as transaction:
rows = list(transaction.read(self.TABLE, self.COLUMNS, keyset))
self.assertEqual(rows, [])
transaction.insert_or_update(
self.TABLE, self.COLUMNS, self.ROW_DATA)
# Inserted rows can't be read until after commit.
rows = list(transaction.read(self.TABLE, self.COLUMNS, keyset))
self.assertEqual(rows, [])
rows = list(session.read(self.TABLE, self.COLUMNS, keyset))
self._check_row_data(rows)
def _set_up_table(self, row_count):
from google.cloud.spanner import KeySet
def _row_data(max_index):
for index in range(max_index):
yield [
index,
'First%09d' % (index, ),
'Last09%d' % (index),
'*****@*****.**' % (index, )
]
keyset = KeySet(all_=True)
retry = RetryInstanceState(_has_all_ddl)
retry(self._db.reload)()
session = self._db.session()
session.create()
self.to_delete.append(session)
with session.transaction() as transaction:
transaction.delete(self.TABLE, keyset)
transaction.insert(self.TABLE, self.COLUMNS, _row_data(row_count))
return session, keyset, transaction.committed
def _read_w_concurrent_update(self, transaction, pkey):
keyset = KeySet(keys=[(pkey, )])
rows = list(transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset))
self.assertEqual(len(rows), 1)
pkey, value = rows[0]
transaction.update(COUNTERS_TABLE, COUNTERS_COLUMNS,
[[pkey, value + 1]])
def test_read_w_single_key(self):
ROW_COUNT = 40
session, committed = self._set_up_table(ROW_COUNT)
snapshot = session.snapshot(read_timestamp=committed)
rows = list(
snapshot.read(self.TABLE, self.COLUMNS, KeySet(keys=[(0, )])))
all_data_rows = list(self._row_data(ROW_COUNT))
expected = [all_data_rows[0]]
self._check_row_data(rows, expected)
def test_read_w_multiple_keys(self):
ROW_COUNT = 40
indices = [0, 5, 17]
session, committed = self._set_up_table(ROW_COUNT)
snapshot = session.snapshot(read_timestamp=committed)
rows = list(
snapshot.read(self.TABLE, self.COLUMNS,
KeySet(keys=[(index, ) for index in indices])))
all_data_rows = list(self._row_data(ROW_COUNT))
expected = [row for row in all_data_rows if row[0] in indices]
self._check_row_data(rows, expected)
def _handle_abort_unit_of_work(self, transaction):
keyset_1 = KeySet(keys=[(self.KEY1, )])
rows_1 = list(
transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset_1))
assert len(rows_1) == 1
row_1 = rows_1[0]
value_1 = row_1[1]
self.handler_running.set()
self.provoker_done.wait()
keyset_2 = KeySet(keys=[(self.KEY2, )])
rows_2 = list(
transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset_2))
assert len(rows_2) == 1
row_2 = rows_2[0]
value_2 = row_2[1]
transaction.update(COUNTERS_TABLE, COUNTERS_COLUMNS,
[[self.KEY2, value_1 + value_2]])
def _provoke_abort_unit_of_work(self, transaction):
keyset = KeySet(keys=[(self.KEY1, )])
rows = list(transaction.read(COUNTERS_TABLE, COUNTERS_COLUMNS, keyset))
assert len(rows) == 1
row = rows[0]
value = row[1]
self.provoker_started.set()
self.handler_running.wait()
transaction.update(COUNTERS_TABLE, COUNTERS_COLUMNS,
[[self.KEY1, value + 1]])
def test_batch_insert_then_read(self):
from google.cloud.spanner import KeySet
keyset = KeySet(all_=True)
retry = RetryInstanceState(_has_all_ddl)
retry(self._db.reload)()
session = self._db.session()
session.create()
self.to_delete.append(session)
batch = session.batch()
batch.delete(self.TABLE, keyset)
batch.insert(self.TABLE, self.COLUMNS, self.ROW_DATA)
batch.commit()
snapshot = session.snapshot(read_timestamp=batch.committed)
rows = list(snapshot.read(self.TABLE, self.COLUMNS, keyset))
self._check_row_data(rows)
def _transaction_concurrency_helper(self, unit_of_work, pkey):
INITIAL_VALUE = 123
NUM_THREADS = 3 # conforms to equivalent Java systest.
retry = RetryInstanceState(_has_all_ddl)
retry(self._db.reload)()
session = self._db.session()
session.create()
self.to_delete.append(session)
with session.batch() as batch:
batch.insert_or_update(
COUNTERS_TABLE, COUNTERS_COLUMNS, [[pkey, INITIAL_VALUE]])
# We don't want to run the threads' transactions in the current
# session, which would fail.
txn_sessions = []
for _ in range(NUM_THREADS):
txn_session = self._db.session()
txn_sessions.append(txn_session)
txn_session.create()
self.to_delete.append(txn_session)
threads = [
threading.Thread(
target=txn_session.run_in_transaction,
args=(unit_of_work, pkey))
for txn_session in txn_sessions]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
keyset = KeySet(keys=[(pkey,)])
rows = list(session.read(
COUNTERS_TABLE, COUNTERS_COLUMNS, keyset))
self.assertEqual(len(rows), 1)
_, value = rows[0]
self.assertEqual(value, INITIAL_VALUE + len(threads))
class _TestData(object):
TABLE = 'contacts'
COLUMNS = ('contact_id', 'first_name', 'last_name', 'email')
ROW_DATA = (
(1, u'Phred', u'Phlyntstone', u'*****@*****.**'),
(2, u'Bharney', u'Rhubble', u'*****@*****.**'),
(3, u'Wylma', u'Phlyntstone', u'*****@*****.**'),
)
ALL = KeySet(all_=True)
SQL = 'SELECT * FROM contacts ORDER BY contact_id'
def _assert_timestamp(self, value, nano_value):
self.assertIsInstance(value, datetime.datetime)
self.assertIsNone(value.tzinfo)
self.assertIs(nano_value.tzinfo, UTC)
self.assertEqual(value.year, nano_value.year)
self.assertEqual(value.month, nano_value.month)
self.assertEqual(value.day, nano_value.day)
self.assertEqual(value.hour, nano_value.hour)
self.assertEqual(value.minute, nano_value.minute)
self.assertEqual(value.second, nano_value.second)
self.assertEqual(value.microsecond, nano_value.microsecond)
if isinstance(value, TimestampWithNanoseconds):
self.assertEqual(value.nanosecond, nano_value.nanosecond)
else:
self.assertEqual(value.microsecond * 1000, nano_value.nanosecond)
def _check_row_data(self, row_data, expected=None):
if expected is None:
expected = self.ROW_DATA
self.assertEqual(len(row_data), len(expected))
for found, expected in zip(row_data, expected):
self.assertEqual(len(found), len(expected))
for found_cell, expected_cell in zip(found, expected):
if isinstance(found_cell, TimestampWithNanoseconds):
self._assert_timestamp(expected_cell, found_cell)
elif isinstance(found_cell, float) and math.isnan(found_cell):
self.assertTrue(math.isnan(expected_cell))
else:
self.assertEqual(found_cell, expected_cell)
def delete_data(cls, table_name, id_list, transaction=None):
"""
Delete data row
:type table_name: str
:param table_name: database table name
:type id_list: list
:param id_list: id tuple list
eg. [('1'), ('2')]
:type transaction: Transaction
:param transaction:
"""
start_time = time()
key_set = KeySet(keys=id_list, all_=False)
if transaction is None:
db_instance = Connection.get_instance()
with db_instance.batch() as batch:
batch.delete(table_name, key_set)
else:
transaction.delete(table_name, key_set)
logging.debug('Query completion Time: %s', (time() - start_time))