[DRAFT] feat: add new sync data client

google/cloud/bigtable/data/__init__.py

@@ -20,6 +20,11 @@
 
 from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync
 
+from google.cloud.bigtable.data._sync.client import BigtableDataClient
+from google.cloud.bigtable.data._sync.client import Table
+
+from google.cloud.bigtable.data._sync.mutations_batcher import MutationsBatcher
+
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 from google.cloud.bigtable.data.read_rows_query import RowRange
 from google.cloud.bigtable.data.row import Row
@@ -48,12 +53,15 @@
 __version__: str = package_version.__version__
 
 __all__ = (
+    "BigtableDataClient",
+    "Table",
+    "MutationsBatcher",
     "BigtableDataClientAsync",
     "TableAsync",
+    "MutationsBatcherAsync",
     "RowKeySamples",
     "ReadRowsQuery",
     "RowRange",
-    "MutationsBatcherAsync",
     "Mutation",
     "RowMutationEntry",
     "SetCell",

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

@@ -101,7 +101,9 @@ def __init__(
             bt_exceptions._MutateRowsIncomplete,
         )
         sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60)
-        self._operation = retries.retry_target_async(
+        # Note: _operation could be a raw coroutine, but using a lambda
+        # wrapper helps unify with sync code
+        self._operation = lambda: retries.retry_target_async(
             self._run_attempt,
             self.is_retryable,
             sleep_generator,
@@ -125,7 +127,7 @@ async def start(self):
         """
         try:
             # trigger mutate_rows
-            await self._operation
+            await self._operation()
         except Exception as exc:
             # exceptions raised by retryable are added to the list of exceptions for all unfinalized mutations
             incomplete_indices = self.remaining_indices.copy()

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

@@ -17,7 +17,6 @@
 
 from typing import (
     TYPE_CHECKING,
-    AsyncGenerator,
     AsyncIterable,
     Awaitable,
     Sequence,
@@ -32,9 +31,7 @@
 from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 from google.cloud.bigtable.data.exceptions import InvalidChunk
 from google.cloud.bigtable.data.exceptions import _RowSetComplete
-from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
-from google.cloud.bigtable.data._helpers import _make_metadata
-from google.cloud.bigtable.data._helpers import _retry_exception_factory
+from google.cloud.bigtable.data import _helpers
 
 from google.api_core import retry as retries
 from google.api_core.retry import exponential_sleep_generator
@@ -80,7 +77,7 @@ def __init__(
         attempt_timeout: float,
         retryable_exceptions: Sequence[type[Exception]] = (),
     ):
-        self.attempt_timeout_gen = _attempt_timeout_generator(
+        self.attempt_timeout_gen = _helpers._attempt_timeout_generator(
             attempt_timeout, operation_timeout
         )
         self.operation_timeout = operation_timeout
@@ -94,14 +91,14 @@ def __init__(
             self.request = query._to_pb(table)
         self.table = table
         self._predicate = retries.if_exception_type(*retryable_exceptions)
-        self._metadata = _make_metadata(
+        self._metadata = _helpers._make_metadata(
             table.table_name,
             table.app_profile_id,
         )
         self._last_yielded_row_key: bytes | None = None
         self._remaining_count: int | None = self.request.rows_limit or None
 
-    def start_operation(self) -> AsyncGenerator[Row, None]:
+    def start_operation(self) -> AsyncIterable[Row]:
         """
         Start the read_rows operation, retrying on retryable errors.
         """
@@ -110,10 +107,10 @@ def start_operation(self) -> AsyncGenerator[Row, None]:
             self._predicate,
             exponential_sleep_generator(0.01, 60, multiplier=2),
             self.operation_timeout,
-            exception_factory=_retry_exception_factory,
+            exception_factory=_helpers._retry_exception_factory,
         )
 
-    def _read_rows_attempt(self) -> AsyncGenerator[Row, None]:
+    def _read_rows_attempt(self) -> AsyncIterable[Row]:
         """
         Attempt a single read_rows rpc call.
         This function is intended to be wrapped by retry logic,
@@ -148,7 +145,7 @@ def _read_rows_attempt(self) -> AsyncGenerator[Row, None]:
 
     async def chunk_stream(
         self, stream: Awaitable[AsyncIterable[ReadRowsResponsePB]]
-    ) -> AsyncGenerator[ReadRowsResponsePB.CellChunk, None]:
+    ) -> AsyncIterable[ReadRowsResponsePB.CellChunk]:
         """
         process chunks out of raw read_rows stream
         """
@@ -193,9 +190,7 @@ async def chunk_stream(
                     current_key = None
 
     @staticmethod
-    async def merge_rows(
-        chunks: AsyncGenerator[ReadRowsResponsePB.CellChunk, None] | None
-    ):
+    async def merge_rows(chunks: AsyncIterable[ReadRowsResponsePB.CellChunk] | None):
         """
         Merge chunks into rows
         """