diff --git a/google/cloud/bigtable/row.py b/google/cloud/bigtable/row.py index 6c54ac0e5..e5aa0503a 100644 --- a/google/cloud/bigtable/row.py +++ b/google/cloud/bigtable/row.py @@ -14,16 +14,13 @@ """User-friendly container for Google Cloud Bigtable Row.""" - -import struct - from google.cloud._helpers import _datetime_from_microseconds # type: ignore from google.cloud._helpers import _microseconds_from_datetime # type: ignore from google.cloud._helpers import _to_bytes # type: ignore -from google.cloud.bigtable_v2.types import data as data_v2_pb2 +from google.cloud.bigtable.data import mutations +from google.cloud.bigtable.data import read_modify_write_rules as rmw_rules -_PACK_I64 = struct.Struct(">q").pack MAX_MUTATIONS = 100000 """The maximum number of mutations that a row can accumulate.""" @@ -157,26 +154,21 @@ def _set_cell(self, column_family_id, column, value, timestamp=None, state=None) :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - column = _to_bytes(column) - if isinstance(value, int): - value = _PACK_I64(value) - value = _to_bytes(value) if timestamp is None: - # Use -1 for current Bigtable server time. - timestamp_micros = -1 + # Use current Bigtable server time. + timestamp_micros = mutations._SERVER_SIDE_TIMESTAMP else: timestamp_micros = _microseconds_from_datetime(timestamp) # Truncate to millisecond granularity. timestamp_micros -= timestamp_micros % 1000 - mutation_val = data_v2_pb2.Mutation.SetCell( - family_name=column_family_id, - column_qualifier=column, + mutation = mutations.SetCell( + family=column_family_id, + qualifier=column, + new_value=value, timestamp_micros=timestamp_micros, - value=value, ) - mutation_pb = data_v2_pb2.Mutation(set_cell=mutation_val) - self._get_mutations(state).append(mutation_pb) + self._get_mutations(state).append(mutation) def _delete(self, state=None): """Helper for :meth:`delete` @@ -191,9 +183,7 @@ def _delete(self, state=None): :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - mutation_val = data_v2_pb2.Mutation.DeleteFromRow() - mutation_pb = data_v2_pb2.Mutation(delete_from_row=mutation_val) - self._get_mutations(state).append(mutation_pb) + self._get_mutations(state).append(mutations.DeleteAllFromRow()) def _delete_cells(self, column_family_id, columns, time_range=None, state=None): """Helper for :meth:`delete_cell` and :meth:`delete_cells`. @@ -222,27 +212,28 @@ def _delete_cells(self, column_family_id, columns, time_range=None, state=None): """ mutations_list = self._get_mutations(state) if columns is self.ALL_COLUMNS: - mutation_val = data_v2_pb2.Mutation.DeleteFromFamily( - family_name=column_family_id + self._get_mutations(state).append( + mutations.DeleteAllFromFamily(family_to_delete=column_family_id) ) - mutation_pb = data_v2_pb2.Mutation(delete_from_family=mutation_val) - mutations_list.append(mutation_pb) else: - delete_kwargs = {} + start_timestamp_micros = None + end_timestamp_micros = None if time_range is not None: - delete_kwargs["time_range"] = time_range._to_pb() + timestamps = time_range._to_dict() + start_timestamp_micros = timestamps.get("start_timestamp_micros") + end_timestamp_micros = timestamps.get("end_timestamp_micros") to_append = [] for column in columns: column = _to_bytes(column) - # time_range will never change if present, but the rest of - # delete_kwargs will - delete_kwargs.update( - family_name=column_family_id, column_qualifier=column + to_append.append( + mutations.DeleteRangeFromColumn( + family=column_family_id, + qualifier=column, + start_timestamp_micros=start_timestamp_micros, + end_timestamp_micros=end_timestamp_micros, + ) ) - mutation_val = data_v2_pb2.Mutation.DeleteFromColumn(**delete_kwargs) - mutation_pb = data_v2_pb2.Mutation(delete_from_column=mutation_val) - to_append.append(mutation_pb) # We don't add the mutations until all columns have been # processed without error. @@ -284,7 +275,7 @@ class DirectRow(_SetDeleteRow): def __init__(self, row_key, table=None): super(DirectRow, self).__init__(row_key, table) - self._pb_mutations = [] + self._mutations = [] def _get_mutations(self, state=None): # pylint: disable=unused-argument """Gets the list of mutations for a given state. @@ -299,7 +290,12 @@ def _get_mutations(self, state=None): # pylint: disable=unused-argument :rtype: list :returns: The list to add new mutations to (for the current state). """ - return self._pb_mutations + return self._mutations + + def _get_mutation_pbs(self): + """Gets the list of mutation protos.""" + + return [mut._to_pb() for mut in self._get_mutations()] def get_mutations_size(self): """Gets the total mutations size for current row @@ -313,7 +309,7 @@ def get_mutations_size(self): """ mutation_size = 0 - for mutation in self._get_mutations(): + for mutation in self._get_mutation_pbs(): mutation_size += mutation._pb.ByteSize() return mutation_size @@ -486,7 +482,7 @@ def clear(self): :end-before: [END bigtable_api_row_clear] :dedent: 4 """ - del self._pb_mutations[:] + del self._mutations[:] class ConditionalRow(_SetDeleteRow): @@ -597,17 +593,15 @@ def commit(self): % (MAX_MUTATIONS, num_true_mutations, num_false_mutations) ) - data_client = self._table._instance._client.table_data_client - resp = data_client.check_and_mutate_row( - table_name=self._table.name, + table = self._table._table_impl + resp = table.check_and_mutate_row( row_key=self._row_key, - predicate_filter=self._filter._to_pb(), - app_profile_id=self._table._app_profile_id, - true_mutations=true_mutations, - false_mutations=false_mutations, + predicate=self._filter, + true_case_mutations=true_mutations, + false_case_mutations=false_mutations, ) self.clear() - return resp.predicate_matched + return resp # pylint: disable=arguments-differ def set_cell(self, column_family_id, column, value, timestamp=None, state=True): @@ -797,7 +791,7 @@ class AppendRow(Row): def __init__(self, row_key, table): super(AppendRow, self).__init__(row_key, table) - self._rule_pb_list = [] + self._rule_list = [] def clear(self): """Removes all currently accumulated modifications on current row. @@ -809,7 +803,7 @@ def clear(self): :end-before: [END bigtable_api_row_clear] :dedent: 4 """ - del self._rule_pb_list[:] + del self._rule_list[:] def append_cell_value(self, column_family_id, column, value): """Appends a value to an existing cell. @@ -842,12 +836,11 @@ def append_cell_value(self, column_family_id, column, value): the targeted cell is unset, it will be treated as containing the empty string. """ - column = _to_bytes(column) - value = _to_bytes(value) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, column_qualifier=column, append_value=value + self._rule_list.append( + rmw_rules.AppendValueRule( + family=column_family_id, qualifier=column, append_value=value + ) ) - self._rule_pb_list.append(rule_pb) def increment_cell_value(self, column_family_id, column, int_value): """Increments a value in an existing cell. @@ -886,13 +879,11 @@ def increment_cell_value(self, column_family_id, column, int_value): big-endian signed integer), or the entire request will fail. """ - column = _to_bytes(column) - rule_pb = data_v2_pb2.ReadModifyWriteRule( - family_name=column_family_id, - column_qualifier=column, - increment_amount=int_value, + self._rule_list.append( + rmw_rules.IncrementRule( + family=column_family_id, qualifier=column, increment_amount=int_value + ) ) - self._rule_pb_list.append(rule_pb) def commit(self): """Makes a ``ReadModifyWriteRow`` API request. @@ -925,7 +916,7 @@ def commit(self): :raises: :class:`ValueError ` if the number of mutations exceeds the :data:`MAX_MUTATIONS`. """ - num_mutations = len(self._rule_pb_list) + num_mutations = len(self._rule_list) if num_mutations == 0: return {} if num_mutations > MAX_MUTATIONS: @@ -934,12 +925,10 @@ def commit(self): "allowable %d." % (num_mutations, MAX_MUTATIONS) ) - data_client = self._table._instance._client.table_data_client - row_response = data_client.read_modify_write_row( - table_name=self._table.name, + table = self._table._table_impl + row_response = table.read_modify_write_row( row_key=self._row_key, - rules=self._rule_pb_list, - app_profile_id=self._table._app_profile_id, + rules=self._rule_list, ) # Reset modifications after commit-ing request. @@ -983,47 +972,13 @@ def _parse_rmw_row_response(row_response): } """ result = {} - for column_family in row_response.row.families: - column_family_id, curr_family = _parse_family_pb(column_family) - result[column_family_id] = curr_family + for cell in row_response.cells: + result.setdefault(cell.family, {}).setdefault(cell.qualifier, []).append( + (cell.value, _datetime_from_microseconds(cell.timestamp_micros)) + ) return result -def _parse_family_pb(family_pb): - """Parses a Family protobuf into a dictionary. - - :type family_pb: :class:`._generated.data_pb2.Family` - :param family_pb: A protobuf - - :rtype: tuple - :returns: A string and dictionary. The string is the name of the - column family and the dictionary has column names (within the - family) as keys and cell lists as values. Each cell is - represented with a two-tuple with the value (in bytes) and the - timestamp for the cell. For example: - - .. code:: python - - { - b'col-name1': [ - (b'cell-val', datetime.datetime(...)), - (b'cell-val-newer', datetime.datetime(...)), - ], - b'col-name2': [ - (b'altcol-cell-val', datetime.datetime(...)), - ], - } - """ - result = {} - for column in family_pb.columns: - result[column.qualifier] = cells = [] - for cell in column.cells: - val_pair = (cell.value, _datetime_from_microseconds(cell.timestamp_micros)) - cells.append(val_pair) - - return family_pb.name, result - - class PartialRowData(object): """Representation of partial row in a Google Cloud Bigtable Table. diff --git a/google/cloud/bigtable/table.py b/google/cloud/bigtable/table.py index 9d2897daa..39e02e4c3 100644 --- a/google/cloud/bigtable/table.py +++ b/google/cloud/bigtable/table.py @@ -1368,7 +1368,7 @@ def _compile_mutation_entries(table_name, rows): for row in rows: _check_row_table_name(table_name, row) _check_row_type(row) - mutations = row._get_mutations() + mutations = row._get_mutation_pbs() entries.append(entry_klass(row_key=row.row_key, mutations=mutations)) mutations_count += len(mutations) diff --git a/tests/system/v2_client/test_data_api.py b/tests/system/v2_client/test_data_api.py index f41aa8377..0f3716472 100644 --- a/tests/system/v2_client/test_data_api.py +++ b/tests/system/v2_client/test_data_api.py @@ -14,7 +14,6 @@ import datetime import operator -import struct import pytest @@ -50,6 +49,8 @@ INITIAL_ROW_SPLITS = [b"row_split_1", b"row_split_2", b"row_split_3"] JOY_EMOJI = "\N{FACE WITH TEARS OF JOY}" +GAP_MARGIN_OF_ERROR = 0.05 + PASS_ALL_FILTER = row_filters.PassAllFilter(True) BLOCK_ALL_FILTER = row_filters.BlockAllFilter(True) @@ -266,7 +267,7 @@ def test_error_handler(exc): retry._initial * retry._multiplier**times_triggered, retry._maximum, ) - assert gap <= max_gap + assert gap <= max_gap + GAP_MARGIN_OF_ERROR times_triggered += 1 curr_time = next_time @@ -1041,22 +1042,11 @@ def test_table_direct_row_input_errors(data_table, rows_to_delete): with pytest.raises(TypeError): row.delete_cell(COLUMN_FAMILY_ID1, INT_COL_NAME) - # Unicode for column name and value does not get converted to bytes because - # internally we use to_bytes in ascii mode. - with pytest.raises(UnicodeEncodeError): - row.set_cell(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - row.delete_cell(COLUMN_FAMILY_ID1, JOY_EMOJI) - - # Various non int64s, we use struct to pack a Python int to bytes. - with pytest.raises(struct.error): + # Various non int64s + with pytest.raises(ValueError): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - with pytest.raises(struct.error): + with pytest.raises(ValueError): row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL2) # Since floats aren't ints, they aren't converted to bytes via struct.pack, @@ -1101,22 +1091,11 @@ def test_table_conditional_row_input_errors(data_table, rows_to_delete): with pytest.raises(TypeError): true_row.delete_cell(COLUMN_FAMILY_ID1, INT_COL_NAME) - # Unicode for column name and value does not get converted to bytes because - # internally we use to_bytes in ascii mode. - with pytest.raises(UnicodeEncodeError): - true_row.set_cell(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - true_row.delete_cell(COLUMN_FAMILY_ID1, JOY_EMOJI) - - # Various non int64s, we use struct to pack a Python int to bytes. - with pytest.raises(struct.error): + # Various non int64s + with pytest.raises(ValueError): true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - with pytest.raises(struct.error): + with pytest.raises(ValueError): true_row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL2) # Since floats aren't ints, they aren't converted to bytes via struct.pack, diff --git a/tests/unit/v2_client/test_row.py b/tests/unit/v2_client/test_row.py index b8a1917a9..a8d27732b 100644 --- a/tests/unit/v2_client/test_row.py +++ b/tests/unit/v2_client/test_row.py @@ -19,6 +19,9 @@ from ._testing import _make_credentials +_INSTANCE_ID = "test-instance" + + def _make_client(*args, **kwargs): from google.cloud.bigtable.client import Client @@ -66,7 +69,7 @@ def test_direct_row_constructor(): row = _make_direct_row(row_key, table) assert row._row_key == row_key assert row._table is table - assert row._pb_mutations == [] + assert row._mutations == [] def test_direct_row_constructor_with_unicode(): @@ -89,10 +92,28 @@ def test_direct_row__get_mutations(): row_key = b"row_key" row = _make_direct_row(row_key, None) - row._pb_mutations = mutations = object() + row._mutations = mutations = object() assert mutations is row._get_mutations(None) +def test_direct_row__get_mutation_pbs(): + from google.cloud.bigtable.data.mutations import SetCell, _SERVER_SIDE_TIMESTAMP + + row_key = b"row_key" + row = _make_direct_row(row_key, None) + + mutation = SetCell( + family="column_family_id", + qualifier=b"column", + new_value=b"value", + timestamp_micros=_SERVER_SIDE_TIMESTAMP, + ) + + row._mutations = [mutation] + + assert row._get_mutation_pbs() == [mutation._to_pb()] + + def test_direct_row_get_mutations_size(): row_key = b"row_key" row = _make_direct_row(row_key, None) @@ -108,7 +129,7 @@ def test_direct_row_get_mutations_size(): row.set_cell(column_family_id2, column2, value) total_mutations_size = 0 - for mutation in row._get_mutations(): + for mutation in row._get_mutation_pbs(): total_mutations_size += mutation._pb.ByteSize() assert row.get_mutations_size() == total_mutations_size @@ -123,26 +144,27 @@ def _set_cell_helper( ): import struct + from google.cloud.bigtable.data.mutations import SetCell + row_key = b"row_key" column_family_id = "column_family_id" if column is None: column = b"column" table = object() row = _make_direct_row(row_key, table) - assert row._pb_mutations == [] + assert row._mutations == [] row.set_cell(column_family_id, column, value, timestamp=timestamp) if isinstance(value, int): value = struct.pack(">q", value) - expected_pb = _MutationPB( - set_cell=_MutationSetCellPB( - family_name=column_family_id, - column_qualifier=column_bytes or column, - timestamp_micros=timestamp_micros, - value=value, - ) + expected_mutation = SetCell( + family=column_family_id, + qualifier=column_bytes or column, + new_value=value, + timestamp_micros=timestamp_micros, ) - assert row._pb_mutations == [expected_pb] + + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_set_cell(): @@ -183,13 +205,15 @@ def test_direct_row_set_cell_with_non_null_timestamp(): def test_direct_row_delete(): + from google.cloud.bigtable.data.mutations import DeleteAllFromRow + row_key = b"row_key" row = _make_direct_row(row_key, object()) - assert row._pb_mutations == [] + assert row._mutations == [] row.delete() - expected_pb = _MutationPB(delete_from_row=_MutationDeleteFromRowPB()) - assert row._pb_mutations == [expected_pb] + expected_mutation = DeleteAllFromRow() + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cell(): @@ -213,14 +237,14 @@ def _delete_cells(self, *args, **kwargs): mock_row = MockRow(row_key, table) # Make sure no values are set before calling the method. - assert mock_row._pb_mutations == [] + assert mock_row._mutations == [] assert mock_row._args == [] assert mock_row._kwargs == [] # Actually make the request against the mock class. time_range = object() mock_row.delete_cell(column_family_id, column, time_range=time_range) - assert mock_row._pb_mutations == [] + assert mock_row._mutations == [] assert mock_row._args == [(column_family_id, [column])] assert mock_row._kwargs == [{"state": None, "time_range": time_range}] @@ -238,19 +262,18 @@ def test_direct_row_delete_cells_non_iterable(): def test_direct_row_delete_cells_all_columns(): from google.cloud.bigtable.row import DirectRow + from google.cloud.bigtable.data.mutations import DeleteAllFromFamily row_key = b"row_key" column_family_id = "column_family_id" table = object() row = _make_direct_row(row_key, table) - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, DirectRow.ALL_COLUMNS) - expected_pb = _MutationPB( - delete_from_family=_MutationDeleteFromFamilyPB(family_name=column_family_id) - ) - assert row._pb_mutations == [expected_pb] + expected_mutation = DeleteAllFromFamily(family_to_delete=column_family_id) + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cells_no_columns(): @@ -260,12 +283,14 @@ def test_direct_row_delete_cells_no_columns(): row = _make_direct_row(row_key, table) columns = [] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns) - assert row._pb_mutations == [] + assert row._mutations == [] def _delete_cells_helper(time_range=None): + from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn + row_key = b"row_key" column = b"column" column_family_id = "column_family_id" @@ -273,17 +298,19 @@ def _delete_cells_helper(time_range=None): row = _make_direct_row(row_key, table) columns = [column] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns, time_range=time_range) - expected_pb = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column - ) - ) + expected_mutation = DeleteRangeFromColumn(family=column_family_id, qualifier=column) if time_range is not None: - expected_pb.delete_from_column.time_range._pb.CopyFrom(time_range._to_pb()._pb) - assert row._pb_mutations == [expected_pb] + time_range_pb = time_range._to_pb() + if time_range_pb.start_timestamp_micros: + expected_mutation.start_timestamp_micros = ( + time_range_pb.start_timestamp_micros + ) + if time_range_pb.end_timestamp_micros: + expected_mutation.end_timestamp_micros = time_range_pb.end_timestamp_micros + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cells_no_time_range(): @@ -311,13 +338,15 @@ def test_direct_row_delete_cells_with_bad_column(): row = _make_direct_row(row_key, table) columns = [column, object()] - assert row._pb_mutations == [] + assert row._mutations == [] with pytest.raises(TypeError): row.delete_cells(column_family_id, columns) - assert row._pb_mutations == [] + assert row._mutations == [] def test_direct_row_delete_cells_with_string_columns(): + from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn + row_key = b"row_key" column_family_id = "column_family_id" column1 = "column1" @@ -328,20 +357,16 @@ def test_direct_row_delete_cells_with_string_columns(): row = _make_direct_row(row_key, table) columns = [column1, column2] - assert row._pb_mutations == [] + assert row._mutations == [] row.delete_cells(column_family_id, columns) - expected_pb1 = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column1_bytes - ) + expected_mutation1 = DeleteRangeFromColumn( + family=column_family_id, qualifier=column1_bytes ) - expected_pb2 = _MutationPB( - delete_from_column=_MutationDeleteFromColumnPB( - family_name=column_family_id, column_qualifier=column2_bytes - ) + expected_mutation2 = DeleteRangeFromColumn( + family=column_family_id, qualifier=column2_bytes ) - assert row._pb_mutations == [expected_pb1, expected_pb2] + _assert_mutations_equal(row._mutations, [expected_mutation1, expected_mutation2]) def test_direct_row_commit(): @@ -517,50 +542,54 @@ def test_append_row_constructor(): row = _make_append_row(row_key, table) assert row._row_key == row_key assert row._table is table - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_clear(): row_key = b"row_key" table = object() row = _make_append_row(row_key, table) - row._rule_pb_list = [1, 2, 3] + row._rule_list = [1, 2, 3] row.clear() - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_append_cell_value(): + from google.cloud.bigtable.data.read_modify_write_rules import AppendValueRule + table = object() row_key = b"row_key" row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] column = b"column" column_family_id = "column_family_id" value = b"bytes-val" row.append_cell_value(column_family_id, column, value) - expected_pb = _ReadModifyWriteRulePB( - family_name=column_family_id, column_qualifier=column, append_value=value + expected_pb = AppendValueRule( + family=column_family_id, qualifier=column, append_value=value ) - assert row._rule_pb_list == [expected_pb] + _assert_mutations_equal(row._rule_list, [expected_pb]) def test_append_row_increment_cell_value(): + from google.cloud.bigtable.data.read_modify_write_rules import IncrementRule + table = object() row_key = b"row_key" row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] column = b"column" column_family_id = "column_family_id" int_value = 281330 row.increment_cell_value(column_family_id, column, int_value) - expected_pb = _ReadModifyWriteRulePB( - family_name=column_family_id, - column_qualifier=column, + expected_pb = IncrementRule( + family=column_family_id, + qualifier=column, increment_amount=int_value, ) - assert row._rule_pb_list == [expected_pb] + _assert_mutations_equal(row._rule_list, [expected_pb]) def test_append_row_commit(): @@ -605,7 +634,7 @@ def mock_parse_rmw_row_response(row_response): call_args = api.read_modify_write_row.call_args_list[0] assert app_profile_id == call_args.app_profile_id[0] assert result == expected_result - assert row._rule_pb_list == [] + assert row._rule_list == [] def test_append_row_commit_no_rules(): @@ -618,7 +647,7 @@ def test_append_row_commit_no_rules(): client = _make_client(project=project_id, credentials=credentials, admin=True) table = _Table(None, client=client) row = _make_append_row(row_key, table) - assert row._rule_pb_list == [] + assert row._rule_list == [] # Patch the stub used by the API method. stub = _FakeStub() @@ -637,8 +666,8 @@ def test_append_row_commit_too_many_mutations(): row_key = b"row_key" table = object() row = _make_append_row(row_key, table) - row._rule_pb_list = [1, 2, 3] - num_mutations = len(row._rule_pb_list) + row._rule_list = [1, 2, 3] + num_mutations = len(row._rule_list) with _Monkey(MUT, MAX_MUTATIONS=num_mutations - 1): with pytest.raises(ValueError): row.commit() @@ -648,6 +677,8 @@ def test__parse_rmw_row_response(): from google.cloud._helpers import _datetime_from_microseconds from google.cloud.bigtable.row import _parse_rmw_row_response + from google.cloud.bigtable.data.row import Row + col_fam1 = "col-fam-id" col_fam2 = "col-fam-id2" col_name1 = b"col-name1" @@ -696,59 +727,16 @@ def test__parse_rmw_row_response(): ), ] ) - sample_input = _ReadModifyWriteRowResponsePB(row=response_row) + sample_input = Row._from_pb(response_row) assert expected_output == _parse_rmw_row_response(sample_input) -def test__parse_family_pb(): - from google.cloud._helpers import _datetime_from_microseconds - from google.cloud.bigtable.row import _parse_family_pb - - col_fam1 = "col-fam-id" - col_name1 = b"col-name1" - col_name2 = b"col-name2" - cell_val1 = b"cell-val" - cell_val2 = b"cell-val-newer" - cell_val3 = b"altcol-cell-val" - - microseconds = 5554441037 - timestamp = _datetime_from_microseconds(microseconds) - expected_dict = { - col_name1: [(cell_val1, timestamp), (cell_val2, timestamp)], - col_name2: [(cell_val3, timestamp)], - } - expected_output = (col_fam1, expected_dict) - sample_input = _FamilyPB( - name=col_fam1, - columns=[ - _ColumnPB( - qualifier=col_name1, - cells=[ - _CellPB(value=cell_val1, timestamp_micros=microseconds), - _CellPB(value=cell_val2, timestamp_micros=microseconds), - ], - ), - _ColumnPB( - qualifier=col_name2, - cells=[_CellPB(value=cell_val3, timestamp_micros=microseconds)], - ), - ], - ) - assert expected_output == _parse_family_pb(sample_input) - - def _CheckAndMutateRowResponsePB(*args, **kw): from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 return messages_v2_pb2.CheckAndMutateRowResponse(*args, **kw) -def _ReadModifyWriteRowResponsePB(*args, **kw): - from google.cloud.bigtable_v2.types import bigtable as messages_v2_pb2 - - return messages_v2_pb2.ReadModifyWriteRowResponse(*args, **kw) - - def _CellPB(*args, **kw): from google.cloud.bigtable_v2.types import data as data_v2_pb2 @@ -767,46 +755,18 @@ def _FamilyPB(*args, **kw): return data_v2_pb2.Family(*args, **kw) -def _MutationPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation(*args, **kw) - - -def _MutationSetCellPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.SetCell(*args, **kw) - - -def _MutationDeleteFromColumnPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromColumn(*args, **kw) - - -def _MutationDeleteFromFamilyPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromFamily(*args, **kw) - - -def _MutationDeleteFromRowPB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 - - return data_v2_pb2.Mutation.DeleteFromRow(*args, **kw) - - def _RowPB(*args, **kw): from google.cloud.bigtable_v2.types import data as data_v2_pb2 return data_v2_pb2.Row(*args, **kw) -def _ReadModifyWriteRulePB(*args, **kw): - from google.cloud.bigtable_v2.types import data as data_v2_pb2 +def _assert_mutations_equal(mutations_1, mutations_2): + assert len(mutations_1) == len(mutations_2) - return data_v2_pb2.ReadModifyWriteRule(*args, **kw) + for i in range(0, len(mutations_1)): + assert type(mutations_1[i]) is type(mutations_2[i]) + assert mutations_1[i]._to_pb() == mutations_2[i]._to_pb() class _Instance(object): @@ -822,6 +782,12 @@ def __init__(self, name, client=None, app_profile_id=None): self.client = client self.mutated_rows = [] + self._table_impl = self._instance._client._veneer_data_client.get_table( + _INSTANCE_ID, + self.name, + app_profile_id=self._app_profile_id, + ) + def mutate_rows(self, rows): from google.rpc import status_pb2 diff --git a/tests/unit/v2_client/test_table.py b/tests/unit/v2_client/test_table.py index 5601f6b5e..d1cc518f7 100644 --- a/tests/unit/v2_client/test_table.py +++ b/tests/unit/v2_client/test_table.py @@ -1736,10 +1736,9 @@ def _do_mutate_retryable_rows_helper( expected_entries = [] for row, prior_status in zip(rows, worker.responses_statuses): if prior_status is None or prior_status.code in RETRYABLES: - mutations = row._get_mutations().copy() # row clears on success entry = data_messages_v2_pb2.MutateRowsRequest.Entry( row_key=row.row_key, - mutations=mutations, + mutations=row._get_mutation_pbs().copy(), # row clears on success ) expected_entries.append(entry)