[DRAFT] chore: use cross sync for async methods

google/cloud/bigtable/data/__init__.py

@@ -50,10 +50,10 @@
 __all__ = (
     "BigtableDataClientAsync",
     "TableAsync",
+    "MutationsBatcherAsync",
     "RowKeySamples",
     "ReadRowsQuery",
     "RowRange",
-    "MutationsBatcherAsync",
     "Mutation",
     "RowMutationEntry",
     "SetCell",

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

@@ -15,38 +15,36 @@
 from __future__ import annotations
 
 from typing import Sequence, TYPE_CHECKING
-from dataclasses import dataclass
 import functools
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry as retries
-import google.cloud.bigtable_v2.types.bigtable as types_pb
 import google.cloud.bigtable.data.exceptions as bt_exceptions
 from google.cloud.bigtable.data._helpers import _make_metadata
 from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 from google.cloud.bigtable.data._helpers import _retry_exception_factory
 
 # mutate_rows requests are limited to this number of mutations
 from google.cloud.bigtable.data.mutations import _MUTATE_ROWS_REQUEST_MUTATION_LIMIT
+from google.cloud.bigtable.data.mutations import _EntryWithProto
+
+from google.cloud.bigtable.data._sync.cross_sync import CrossSync
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable_v2.services.bigtable.async_client import (
-        BigtableAsyncClient,
-    )
     from google.cloud.bigtable.data.mutations import RowMutationEntry
-    from google.cloud.bigtable.data._async.client import TableAsync
-
 
-@dataclass
-class _EntryWithProto:
-    """
-    A dataclass to hold a RowMutationEntry and its corresponding proto representation.
-    """
+    if CrossSync.is_async:
+        from google.cloud.bigtable_v2.services.bigtable.async_client import (
+            BigtableAsyncClient,
+        )
 
-    entry: RowMutationEntry
-    proto: types_pb.MutateRowsRequest.Entry
+        CrossSync.add_mapping("GapicClient", BigtableAsyncClient)
 
 
+@CrossSync.export_sync(
+    path="google.cloud.bigtable.data._sync._mutate_rows._MutateRowsOperation",
+    add_mapping_for_name="_MutateRowsOperation",
+)
 class _MutateRowsOperationAsync:
     """
     MutateRowsOperation manages the logic of sending a set of row mutations,
@@ -66,10 +64,11 @@ class _MutateRowsOperationAsync:
             If not specified, the request will run until operation_timeout is reached.
     """
 
+    @CrossSync.convert
     def __init__(
         self,
-        gapic_client: "BigtableAsyncClient",
-        table: "TableAsync",
+        gapic_client: "CrossSync.GapicClient",
+        table: "CrossSync.Table",
         mutation_entries: list["RowMutationEntry"],
         operation_timeout: float,
         attempt_timeout: float | None,
@@ -100,7 +99,7 @@ def __init__(
             bt_exceptions._MutateRowsIncomplete,
         )
         sleep_generator = retries.exponential_sleep_generator(0.01, 2, 60)
-        self._operation = retries.retry_target_async(
+        self._operation = lambda: CrossSync.retry_target(
             self._run_attempt,
             self.is_retryable,
             sleep_generator,
@@ -115,6 +114,7 @@ def __init__(
         self.remaining_indices = list(range(len(self.mutations)))
         self.errors: dict[int, list[Exception]] = {}
 
+    @CrossSync.convert
     async def start(self):
         """
         Start the operation, and run until completion
@@ -124,7 +124,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()
@@ -151,6 +151,7 @@ async def start(self):
                     all_errors, len(self.mutations)
                 )
 
+    @CrossSync.convert
     async def _run_attempt(self):
         """
         Run a single attempt of the mutate_rows rpc.

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

@@ -15,13 +15,7 @@
 
 from __future__ import annotations
 
-from typing import (
-    TYPE_CHECKING,
-    AsyncGenerator,
-    AsyncIterable,
-    Awaitable,
-    Sequence,
-)
+from typing import Sequence
 
 from google.cloud.bigtable_v2.types import ReadRowsRequest as ReadRowsRequestPB
 from google.cloud.bigtable_v2.types import ReadRowsResponse as ReadRowsResponsePB
@@ -32,22 +26,21 @@
 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.exceptions import _ResetRow
 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.api_core import retry as retries
 from google.api_core.retry import exponential_sleep_generator
 
-if TYPE_CHECKING:
-    from google.cloud.bigtable.data._async.client import TableAsync
-
-
-class _ResetRow(Exception):
-    def __init__(self, chunk):
-        self.chunk = chunk
+from google.cloud.bigtable.data._sync.cross_sync import CrossSync
 
 
+@CrossSync.export_sync(
+    path="google.cloud.bigtable.data._sync._read_rows._ReadRowsOperation",
+    add_mapping_for_name="_ReadRowsOperation",
+)
 class _ReadRowsOperationAsync:
     """
     ReadRowsOperation handles the logic of merging chunks from a ReadRowsResponse stream
@@ -82,7 +75,7 @@ class _ReadRowsOperationAsync:
     def __init__(
         self,
         query: ReadRowsQuery,
-        table: "TableAsync",
+        table: "CrossSync.Table",
         operation_timeout: float,
         attempt_timeout: float,
         retryable_exceptions: Sequence[type[Exception]] = (),
@@ -108,22 +101,22 @@ def __init__(
         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) -> CrossSync.Iterable[Row]:
         """
         Start the read_rows operation, retrying on retryable errors.
 
         Yields:
             Row: The next row in the stream
         """
-        return retries.retry_target_stream_async(
+        return CrossSync.retry_target_stream(
             self._read_rows_attempt,
             self._predicate,
             exponential_sleep_generator(0.01, 60, multiplier=2),
             self.operation_timeout,
             exception_factory=_retry_exception_factory,
         )
 
-    def _read_rows_attempt(self) -> AsyncGenerator[Row, None]:
+    def _read_rows_attempt(self) -> CrossSync.Iterable[Row]:
         """
         Attempt a single read_rows rpc call.
         This function is intended to be wrapped by retry logic,
@@ -159,9 +152,10 @@ def _read_rows_attempt(self) -> AsyncGenerator[Row, None]:
         chunked_stream = self.chunk_stream(gapic_stream)
         return self.merge_rows(chunked_stream)
 
+    @CrossSync.convert
     async def chunk_stream(
-        self, stream: Awaitable[AsyncIterable[ReadRowsResponsePB]]
-    ) -> AsyncGenerator[ReadRowsResponsePB.CellChunk, None]:
+        self, stream: CrossSync.Awaitable[CrossSync.Iterable[ReadRowsResponsePB]]
+    ) -> CrossSync.Iterable[ReadRowsResponsePB.CellChunk]:
         """
         process chunks out of raw read_rows stream
 
@@ -211,9 +205,12 @@ async def chunk_stream(
                     current_key = None
 
     @staticmethod
+    @CrossSync.convert(
+        replace_symbols={"__aiter__": "__iter__", "__anext__": "__next__"}
+    )
     async def merge_rows(
-        chunks: AsyncGenerator[ReadRowsResponsePB.CellChunk, None] | None
-    ) -> AsyncGenerator[Row, None]:
+        chunks: CrossSync.Iterable[ReadRowsResponsePB.CellChunk] | None,
+    ) -> CrossSync.Iterable[Row]:
         """
         Merge chunks into rows
 
@@ -229,7 +226,7 @@ async def merge_rows(
         while True:
             try:
                 c = await it.__anext__()
-            except StopAsyncIteration:
+            except CrossSync.StopIteration:
                 # stream complete
                 return
             row_key = c.row_key
@@ -322,7 +319,7 @@ async def merge_rows(
                 ):
                     raise InvalidChunk("reset row with data")
                 continue
-            except StopAsyncIteration:
+            except CrossSync.StopIteration:
                 raise InvalidChunk("premature end of stream")
 
     @staticmethod