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))