feat: Reworked MutateRows to use the data client

Author: gkevinzheng

google/cloud/bigtable/data/_async/client.py

@@ -1556,6 +1556,33 @@ async def mutate_row(
             exception_factory=_retry_exception_factory,
         )
 
+    @CrossSync.convert
+    def _get_mutate_rows_operation(
+        self,
+        mutation_entries: list[RowMutationEntry],
+        *,
+        operation_timeout: float | TABLE_DEFAULT,
+        attempt_timeout: float | None | TABLE_DEFAULT,
+        retryable_errors: Sequence[type[Exception]] | TABLE_DEFAULT,
+    ) -> CrossSync._MutateRowsOperation:
+        """
+        Gets the bulk mutate rows operation object for the given mutation entries.
+        """
+        operation_timeout, attempt_timeout = _get_timeouts(
+            operation_timeout, attempt_timeout, self
+        )
+        retryable_excs = _get_retryable_errors(retryable_errors, self)
+
+        operation = CrossSync._MutateRowsOperation(
+            self.client._gapic_client,
+            self,
+            mutation_entries,
+            operation_timeout,
+            attempt_timeout,
+            retryable_exceptions=retryable_excs,
+        )
+        return operation
+
     @CrossSync.convert
     async def bulk_mutate_rows(
         self,
@@ -1597,18 +1624,11 @@ async def bulk_mutate_rows(
                 Contains details about any failed entries in .exceptions
             ValueError: if invalid arguments are provided
         """
-        operation_timeout, attempt_timeout = _get_timeouts(
-            operation_timeout, attempt_timeout, self
-        )
-        retryable_excs = _get_retryable_errors(retryable_errors, self)
-
-        operation = CrossSync._MutateRowsOperation(
-            self.client._gapic_client,
-            self,
+        operation = self._get_mutate_rows_operation(
             mutation_entries,
-            operation_timeout,
-            attempt_timeout,
-            retryable_exceptions=retryable_excs,
+            operation_timeout=operation_timeout,
+            attempt_timeout=attempt_timeout,
+            retryable_errors=retryable_errors,
         )
         await operation.start()
 

google/cloud/bigtable/data/_sync_autogen/client.py

@@ -1299,6 +1299,29 @@ def mutate_row(
             exception_factory=_retry_exception_factory,
         )
 
+    def _get_mutate_rows_operation(
+        self,
+        mutation_entries: list[RowMutationEntry],
+        *,
+        operation_timeout: float | TABLE_DEFAULT,
+        attempt_timeout: float | None | TABLE_DEFAULT,
+        retryable_errors: Sequence[type[Exception]] | TABLE_DEFAULT,
+    ) -> CrossSync._Sync_Impl._MutateRowsOperation:
+        """Gets the bulk mutate rows operation object for the given mutation entries."""
+        (operation_timeout, attempt_timeout) = _get_timeouts(
+            operation_timeout, attempt_timeout, self
+        )
+        retryable_excs = _get_retryable_errors(retryable_errors, self)
+        operation = CrossSync._Sync_Impl._MutateRowsOperation(
+            self.client._gapic_client,
+            self,
+            mutation_entries,
+            operation_timeout,
+            attempt_timeout,
+            retryable_exceptions=retryable_excs,
+        )
+        return operation
+
     def bulk_mutate_rows(
         self,
         mutation_entries: list[RowMutationEntry],
@@ -1337,17 +1360,11 @@ def bulk_mutate_rows(
             MutationsExceptionGroup: if one or more mutations fails
                 Contains details about any failed entries in .exceptions
             ValueError: if invalid arguments are provided"""
-        (operation_timeout, attempt_timeout) = _get_timeouts(
-            operation_timeout, attempt_timeout, self
-        )
-        retryable_excs = _get_retryable_errors(retryable_errors, self)
-        operation = CrossSync._Sync_Impl._MutateRowsOperation(
-            self.client._gapic_client,
-            self,
+        operation = self._get_mutate_rows_operation(
             mutation_entries,
-            operation_timeout,
-            attempt_timeout,
-            retryable_exceptions=retryable_excs,
+            operation_timeout=operation_timeout,
+            attempt_timeout=attempt_timeout,
+            retryable_errors=retryable_errors,
         )
         operation.start()
 

google/cloud/bigtable/row.py

@@ -153,15 +153,16 @@ def _set_cell(self, column_family_id, column, value, timestamp=None, state=None)
                       integer (8 bytes).
 
         :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (Optional) The timestamp of the operation.
+        :param timestamp: (Optional) The timestamp of the operation. If a
+                         timestamp is not provided, the current system time
+                         will be used.
 
         :type state: bool
         :param state: (Optional) The state that is passed along to
                       :meth:`_get_mutations`.
         """
-        if timestamp is None:
-            # Use current Bigtable server time.
-            timestamp_micros = mutations._SERVER_SIDE_TIMESTAMP
+        if timestamp is None or timestamp == mutations._SERVER_SIDE_TIMESTAMP:
+            timestamp_micros = timestamp
         else:
             timestamp_micros = _microseconds_from_datetime(timestamp)
             # Truncate to millisecond granularity.
@@ -351,7 +352,9 @@ def set_cell(self, column_family_id, column, value, timestamp=None):
                       integer (8 bytes).
 
         :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (Optional) The timestamp of the operation.
+        :param timestamp: (Optional) The timestamp of the operation. If a
+                         timestamp is not provided, the current system time
+                         will be used.
         """
         self._set_cell(column_family_id, column, value, timestamp=timestamp, state=None)
 
@@ -651,7 +654,9 @@ def set_cell(self, column_family_id, column, value, timestamp=None, state=True):
                       integer (8 bytes).
 
         :type timestamp: :class:`datetime.datetime`
-        :param timestamp: (Optional) The timestamp of the operation.
+        :param timestamp: (Optional) The timestamp of the operation. If a
+                         timestamp is not provided, the current system time
+                         will be used.
 
         :type state: bool
         :param state: (Optional) The state that the mutation should be