diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py index d769f2d24644..3aacac0d0f7a 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row.py @@ -14,17 +14,14 @@ """User-friendly container for Google Cloud Bigtable Row.""" -import struct - from google.cloud._helpers import ( _datetime_from_microseconds, # type: ignore _microseconds_from_datetime, # type: ignore _to_bytes, # type: ignore ) -from google.cloud.bigtable_v2.types import data as data_v2_pb2 - -_PACK_I64 = struct.Struct(">q").pack +from google.cloud.bigtable.data import mutations +from google.cloud.bigtable.data import read_modify_write_rules as rmw_rules MAX_MUTATIONS = 100000 """The maximum number of mutations that a row can accumulate.""" @@ -158,26 +155,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` @@ -192,9 +184,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`. @@ -221,33 +211,30 @@ def _delete_cells(self, column_family_id, columns, time_range=None, state=None): :param state: (Optional) The state that is passed along to :meth:`_get_mutations`. """ - 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 = {} - if time_range is not None: - delete_kwargs["time_range"] = time_range._to_pb() + timestamps = time_range._to_dict() if time_range else {} + 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. - mutations_list.extend(to_append) + self._get_mutations(state).extend(to_append) class DirectRow(_SetDeleteRow): @@ -285,7 +272,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. @@ -300,7 +287,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 @@ -314,7 +306,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 @@ -487,7 +479,7 @@ def clear(self): :end-before: [END bigtable_api_row_clear] :dedent: 4 """ - del self._pb_mutations[:] + del self._mutations[:] class ConditionalRow(_SetDeleteRow): @@ -598,17 +590,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): @@ -798,7 +788,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. @@ -810,7 +800,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. @@ -843,12 +833,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. @@ -887,13 +876,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. @@ -926,7 +913,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: @@ -935,12 +922,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. @@ -984,47 +969,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: + column_family = result.setdefault(cell.family, {}) + column = column_family.setdefault(cell.qualifier, []) + column.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/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py b/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py index 9b289ce1c6b5..449543bc132c 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/row_filters.py @@ -17,33 +17,39 @@ import struct from google.cloud.bigtable.data.row_filters import ( # noqa: F401 - RowFilter, - SinkFilter, - _BoolFilter, - PassAllFilter, + ApplyLabelFilter, BlockAllFilter, - _RegexFilter, + CellsColumnLimitFilter, + CellsRowLimitFilter, + CellsRowOffsetFilter, + ColumnQualifierRegexFilter, + FamilyNameRegexFilter, + PassAllFilter, + RowFilter, + RowFilterChain, + RowFilterUnion, RowKeyRegexFilter, RowSampleFilter, - FamilyNameRegexFilter, - ColumnQualifierRegexFilter, + SinkFilter, + StripValueTransformerFilter, TimestampRange, - TimestampRangeFilter as BaseTimestampRangeFilter, - ColumnRangeFilter as BaseColumnRangeFilter, - ValueRegexFilter, - ValueRangeFilter, ValueBitmaskFilter, + ValueRangeFilter, + ValueRegexFilter, + _BoolFilter, _CellCountFilter, - CellsRowOffsetFilter, - CellsRowLimitFilter, - CellsColumnLimitFilter, - StripValueTransformerFilter, - ApplyLabelFilter, _FilterCombination, - RowFilterChain, - RowFilterUnion, + _RegexFilter, +) +from google.cloud.bigtable.data.row_filters import ( + ColumnRangeFilter as BaseColumnRangeFilter, +) +from google.cloud.bigtable.data.row_filters import ( ConditionalRowFilter as BaseConditionalRowFilter, ) +from google.cloud.bigtable.data.row_filters import ( + TimestampRangeFilter as BaseTimestampRangeFilter, +) _PACK_I64 = struct.Struct(">q").pack diff --git a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py index 4c2f2a703933..590ce08aeec8 100644 --- a/packages/google-cloud-bigtable/google/cloud/bigtable/table.py +++ b/packages/google-cloud-bigtable/google/cloud/bigtable/table.py @@ -1374,7 +1374,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/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py index cb90d62de934..64907c706e55 100644 --- a/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py +++ b/packages/google-cloud-bigtable/tests/system/v2_client/test_data_api.py @@ -13,7 +13,6 @@ # limitations under the License. import operator -import struct from datetime import datetime, timedelta, timezone import pytest @@ -46,6 +45,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) @@ -1047,22 +1048,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, @@ -1107,22 +1097,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, @@ -1171,39 +1150,17 @@ def test_table_append_row_input_errors(data_table, rows_to_delete): rows_to_delete.append(data_table.direct_row(ROW_KEY)) # Column names should be convertible to bytes (str or bytes) - with pytest.raises(TypeError): + with pytest.raises(AttributeError): row.append_cell_value(COLUMN_FAMILY_ID1, INT_COL_NAME, CELL_VAL1) - with pytest.raises(TypeError): + with pytest.raises(AttributeError): row.increment_cell_value(COLUMN_FAMILY_ID1, INT_COL_NAME, 1) - # Unicode for column name and value - with pytest.raises(UnicodeEncodeError): - row.append_cell_value(COLUMN_FAMILY_ID1, JOY_EMOJI, CELL_VAL1) - - with pytest.raises(UnicodeEncodeError): - row.append_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, JOY_EMOJI) - - with pytest.raises(UnicodeEncodeError): - row.increment_cell_value(COLUMN_FAMILY_ID1, JOY_EMOJI, 1) - - # Non-integer cell values for increment_cell_value with pytest.raises(ValueError): row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, OVERFLOW_INT_CELL_VAL) - # increment_cell_value does not do input validation on the int_value, instead using - # proto-plus to do validation. - row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, FLOAT_CELL_VAL) - row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME2, FLOAT_CELL_VAL2) - row.commit() - - row_data = data_table.read_row(ROW_KEY) - assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME1][0].value == int( - FLOAT_CELL_VAL - ).to_bytes(8, byteorder="big", signed=True) - assert row_data.cells[COLUMN_FAMILY_ID1][COL_NAME2][0].value == int( - FLOAT_CELL_VAL2 - ).to_bytes(8, byteorder="big", signed=True) + with pytest.raises(TypeError): + row.increment_cell_value(COLUMN_FAMILY_ID1, COL_NAME1, FLOAT_CELL_VAL) # Can't have more than MAX_MUTATIONS mutations, but only enforced after # a row.commit diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py index 5544327b3a97..37205eb7f508 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row.py @@ -18,6 +18,8 @@ from ._testing import _make_credentials +_INSTANCE_ID = "test-instance" + def _make_client(*args, **kwargs): from google.cloud.bigtable.client import Client @@ -66,7 +68,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 +91,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 _SERVER_SIDE_TIMESTAMP, SetCell + + 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 +128,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 +143,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(): @@ -184,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(): @@ -214,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,6 +261,7 @@ def test_direct_row_delete_cells_non_iterable(): def test_direct_row_delete_cells_all_columns(): + from google.cloud.bigtable.data.mutations import DeleteAllFromFamily from google.cloud.bigtable.row import DirectRow row_key = b"row_key" @@ -245,13 +269,11 @@ def test_direct_row_delete_cells_all_columns(): 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(): @@ -261,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" @@ -274,17 +298,17 @@ 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] + timestamps = time_range._to_dict() + expected_mutation.start_timestamp_micros = timestamps.get( + "start_timestamp_micros" + ) + expected_mutation.end_timestamp_micros = timestamps.get("end_timestamp_micros") + _assert_mutations_equal(row._mutations, [expected_mutation]) def test_direct_row_delete_cells_no_time_range(): @@ -314,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" @@ -331,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(): @@ -521,50 +543,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(): @@ -610,7 +636,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(): @@ -623,7 +649,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() @@ -643,8 +669,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() @@ -653,6 +679,7 @@ def test_append_row_commit_too_many_mutations(): def test__parse_rmw_row_response(): from google.cloud._helpers import _datetime_from_microseconds + from google.cloud.bigtable.data.row import Row from google.cloud.bigtable.row import _parse_rmw_row_response col_fam1 = "col-fam-id" @@ -703,60 +730,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 @@ -775,46 +758,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): @@ -830,6 +785,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/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_filters.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_filters.py index 276fbf9af9f2..6f6de3c00899 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_filters.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_row_filters.py @@ -18,10 +18,16 @@ from google.cloud.bigtable import row_filters from google.cloud.bigtable.row_filters import ( _BoolFilter as _BaseBoolFilter, - _RegexFilter as _BaseRegexFilter, +) +from google.cloud.bigtable.row_filters import ( _CellCountFilter as _BaseCellCountFilter, +) +from google.cloud.bigtable.row_filters import ( _FilterCombination as _BaseFilterCombination, ) +from google.cloud.bigtable.row_filters import ( + _RegexFilter as _BaseRegexFilter, +) def test_bool_filter_constructor(): @@ -1228,6 +1234,7 @@ def test_all_row_filters_to_pb_backwards_compatibility(filter_instance): def test_timestamp_range_to_pb_backwards_compatibility(): from datetime import datetime, timezone + from google.cloud.bigtable.row_filters import TimestampRange start = datetime(2023, 1, 1, tzinfo=timezone.utc) diff --git a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py index 55cd5eb9a134..25cc35f7ef7d 100644 --- a/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py +++ b/packages/google-cloud-bigtable/tests/unit/v2_client/test_table.py @@ -1743,10 +1743,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)