fixed imports

google/cloud/bigtable/data/__init__.py

@@ -18,21 +18,21 @@
 
 from google.cloud.bigtable import gapic_version as package_version
 
-from google.cloud.bigtable.client import BigtableDataClient
-from google.cloud.bigtable.client import Table
-
-from google.cloud.bigtable.read_rows_query import ReadRowsQuery
-from google.cloud.bigtable.read_rows_query import RowRange
-from google.cloud.bigtable.row import Row
-from google.cloud.bigtable.row import Cell
-
-from google.cloud.bigtable.mutations_batcher import MutationsBatcher
-from google.cloud.bigtable.mutations import Mutation
-from google.cloud.bigtable.mutations import RowMutationEntry
-from google.cloud.bigtable.mutations import SetCell
-from google.cloud.bigtable.mutations import DeleteRangeFromColumn
-from google.cloud.bigtable.mutations import DeleteAllFromFamily
-from google.cloud.bigtable.mutations import DeleteAllFromRow
+from google.cloud.bigtable.data._async.client import BigtableDataClientAsync
+from google.cloud.bigtable.data._async.client import TableAsync
+
+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
+from google.cloud.bigtable.data.row import Cell
+
+from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync
+from google.cloud.bigtable.data.mutations import Mutation
+from google.cloud.bigtable.data.mutations import RowMutationEntry
+from google.cloud.bigtable.data.mutations import SetCell
+from google.cloud.bigtable.data.mutations import DeleteRangeFromColumn
+from google.cloud.bigtable.data.mutations import DeleteAllFromFamily
+from google.cloud.bigtable.data.mutations import DeleteAllFromRow
 
 # Type alias for the output of sample_keys
 RowKeySamples = List[Tuple[bytes, int]]
@@ -42,12 +42,12 @@
 __version__: str = package_version.__version__
 
 __all__ = (
-    "BigtableDataClient",
-    "Table",
+    "BigtableDataClientAsync",
+    "TableAsync",
     "RowKeySamples",
     "ReadRowsQuery",
     "RowRange",
-    "MutationsBatcher",
+    "MutationsBatcherAsync",
     "Mutation",
     "RowMutationEntry",
     "SetCell",

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

@@ -19,20 +19,20 @@
 
 from google.api_core import exceptions as core_exceptions
 from google.api_core import retry_async as retries
-import google.cloud.bigtable.exceptions as bt_exceptions
-from google.cloud.bigtable._helpers import _make_metadata
-from google.cloud.bigtable._helpers import _convert_retry_deadline
-from google.cloud.bigtable._helpers import _attempt_timeout_generator
+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 _convert_retry_deadline
+from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
 
 # mutate_rows requests are limited to this number of mutations
-from google.cloud.bigtable.mutations import MUTATE_ROWS_REQUEST_MUTATION_LIMIT
+from google.cloud.bigtable.data.mutations import MUTATE_ROWS_REQUEST_MUTATION_LIMIT
 
 if TYPE_CHECKING:
     from google.cloud.bigtable_v2.services.bigtable.async_client import (
         BigtableAsyncClient,
     )
-    from google.cloud.bigtable.client import Table
-    from google.cloud.bigtable.mutations import RowMutationEntry
+    from google.cloud.bigtable.data.client import Table
+    from google.cloud.bigtable.data.mutations import RowMutationEntry
 
 
 class _MutateRowsOperationAsync:

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

@@ -33,16 +33,16 @@
 
 from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse
 from google.cloud.bigtable_v2.services.bigtable.async_client import BigtableAsyncClient
-from google.cloud.bigtable.row import Row, _LastScannedRow
-from google.cloud.bigtable.exceptions import InvalidChunk
-from google.cloud.bigtable.exceptions import _RowSetComplete
-from google.cloud.bigtable.exceptions import IdleTimeout
-from google.cloud.bigtable._read_rows_state_machine import _StateMachine
+from google.cloud.bigtable.data.row import Row, _LastScannedRow
+from google.cloud.bigtable.data.exceptions import InvalidChunk
+from google.cloud.bigtable.data.exceptions import _RowSetComplete
+from google.cloud.bigtable.data.exceptions import IdleTimeout
+from google.cloud.bigtable.data._read_rows_state_machine import _StateMachine
 from google.api_core import retry_async as retries
 from google.api_core import exceptions as core_exceptions
-from google.cloud.bigtable._helpers import _make_metadata
-from google.cloud.bigtable._helpers import _attempt_timeout_generator
-from google.cloud.bigtable._helpers import _convert_retry_deadline
+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 _convert_retry_deadline
 
 
 class _ReadRowsOperationAsync(AsyncIterable[Row]):

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

@@ -43,34 +43,34 @@
 from google.api_core.exceptions import GoogleAPICallError
 from google.api_core import retry_async as retries
 from google.api_core import exceptions as core_exceptions
-from google.cloud.bigtable._read_rows import _ReadRowsOperationAsync
+from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync
+from google.cloud.bigtable.data._async._read_rows import ReadRowsIteratorAsync
 
 import google.auth.credentials
 import google.auth._default
 from google.api_core import client_options as client_options_lib
-from google.cloud.bigtable.row import Row
-from google.cloud.bigtable.read_rows_query import ReadRowsQuery
-from google.cloud.bigtable.iterators import ReadRowsIteratorAsync
-from google.cloud.bigtable.exceptions import FailedQueryShardError
-from google.cloud.bigtable.exceptions import ShardedReadRowsExceptionGroup
-
-from google.cloud.bigtable.mutations import Mutation, RowMutationEntry
-from google.cloud.bigtable._mutate_rows import _MutateRowsOperationAsync
-from google.cloud.bigtable._helpers import _make_metadata
-from google.cloud.bigtable._helpers import _convert_retry_deadline
-from google.cloud.bigtable.mutations_batcher import MutationsBatcher
-from google.cloud.bigtable.mutations_batcher import _MB_SIZE
-from google.cloud.bigtable._helpers import _attempt_timeout_generator
-
-from google.cloud.bigtable.read_modify_write_rules import ReadModifyWriteRule
-from google.cloud.bigtable.row_filters import RowFilter
-from google.cloud.bigtable.row_filters import StripValueTransformerFilter
-from google.cloud.bigtable.row_filters import CellsRowLimitFilter
-from google.cloud.bigtable.row_filters import RowFilterChain
+from google.cloud.bigtable.data.row import Row
+from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
+from google.cloud.bigtable.data.exceptions import FailedQueryShardError
+from google.cloud.bigtable.data.exceptions import ShardedReadRowsExceptionGroup
+
+from google.cloud.bigtable.data.mutations import Mutation, RowMutationEntry
+from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync
+from google.cloud.bigtable.data._helpers import _make_metadata
+from google.cloud.bigtable.data._helpers import _convert_retry_deadline
+from google.cloud.bigtable.data._async.mutations_batcher import MutationsBatcherAsync
+from google.cloud.bigtable.data._async.mutations_batcher import _MB_SIZE
+from google.cloud.bigtable.data._helpers import _attempt_timeout_generator
+
+from google.cloud.bigtable.data.read_modify_write_rules import ReadModifyWriteRule
+from google.cloud.bigtable.data.row_filters import RowFilter
+from google.cloud.bigtable.data.row_filters import StripValueTransformerFilter
+from google.cloud.bigtable.data.row_filters import CellsRowLimitFilter
+from google.cloud.bigtable.data.row_filters import RowFilterChain
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable import RowKeySamples
-    from google.cloud.bigtable import ShardedQuery
+    from google.cloud.bigtable.data.data import RowKeySamples
+    from google.cloud.bigtable.data.data import ShardedQuery
 
 # used by read_rows_sharded to limit how many requests are attempted in parallel
 CONCURRENCY_LIMIT = 10
@@ -556,7 +556,7 @@ async def read_row(
         See read_rows_stream
 
         Raises:
-            - google.cloud.bigtable.exceptions.RowNotFound: if the row does not exist
+            - google.cloud.bigtable.data.exceptions.RowNotFound: if the row does not exist
         Returns:
             - the individual row requested, or None if it does not exist
         """
@@ -762,7 +762,7 @@ def mutations_batcher(
         flow_control_max_bytes: int = 100 * _MB_SIZE,
         batch_operation_timeout: float | None = None,
         batch_per_request_timeout: float | None = None,
-    ) -> MutationsBatcher:
+    ) -> MutationsBatcherAsync:
         """
         Returns a new mutations batcher instance.
 
@@ -782,9 +782,9 @@ def mutations_batcher(
           - batch_per_request_timeout: timeout for each individual request, in seconds. If None,
               table default_per_request_timeout will be used
         Returns:
-            - a MutationsBatcher context manager that can batch requests
+            - a MutationsBatcherAsync context manager that can batch requests
         """
-        return MutationsBatcher(
+        return MutationsBatcherAsync(
             self,
             flush_interval=flush_interval,
             flush_limit_mutation_count=flush_limit_mutation_count,
@@ -930,7 +930,7 @@ async def bulk_mutate_rows(
         if per_request_timeout is not None and per_request_timeout > operation_timeout:
             raise ValueError("per_request_timeout must be less than operation_timeout")
 
-        operation = _MutateRowsOperation(
+        operation = _MutateRowsOperationAsync(
             self.client._gapic_client,
             self,
             mutation_entries,

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

@@ -20,16 +20,13 @@
 import warnings
 from collections import deque
 
-from google.cloud.bigtable.mutations import RowMutationEntry
-from google.cloud.bigtable.exceptions import MutationsExceptionGroup
-from google.cloud.bigtable.exceptions import FailedMutationEntryError
+from google.cloud.bigtable.data.mutations import RowMutationEntry
+from google.cloud.bigtable.data.exceptions import MutationsExceptionGroup
+from google.cloud.bigtable.data.exceptions import FailedMutationEntryError
 
-from google.cloud.bigtable._mutate_rows import _MutateRowsOperation
-from google.cloud.bigtable._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT
-from google.cloud.bigtable.mutations import Mutation
-
-if TYPE_CHECKING:
-    from google.cloud.bigtable.client import Table  # pragma: no cover
+from google.cloud.bigtable.data._async._mutate_rows import _MutateRowsOperationAsync
+from google.cloud.bigtable.data._async._mutate_rows import MUTATE_ROWS_REQUEST_MUTATION_LIMIT
+from google.cloud.bigtable.data.mutations import Mutation
 
 # used to make more readable default values
 _MB_SIZE = 1024 * 1024
@@ -179,7 +176,7 @@ class MutationsBatcherAsync:
 
     def __init__(
         self,
-        table: "Table",
+        table: TableAsync,
         *,
         flush_interval: float | None = 5,
         flush_limit_mutation_count: int | None = 1000,
@@ -353,7 +350,7 @@ async def _execute_mutate_rows(
         if self._table.app_profile_id:
             request["app_profile_id"] = self._table.app_profile_id
         try:
-            operation = _MutateRowsOperation(
+            operation = _MutateRowsOperationAsync(
                 self._table.client._gapic_client,
                 self._table,
                 batch,

google/cloud/bigtable/data/_helpers.py

@@ -18,7 +18,7 @@
 import time
 
 from google.api_core import exceptions as core_exceptions
-from google.cloud.bigtable.exceptions import RetryExceptionGroup
+from google.cloud.bigtable.data.exceptions import RetryExceptionGroup
 
 """
 Helper functions used in various places in the library.

google/cloud/bigtable/data/_read_rows_state_machine.py

@@ -17,8 +17,8 @@
 from typing import Type
 
 from google.cloud.bigtable_v2.types.bigtable import ReadRowsResponse
-from google.cloud.bigtable.row import Row, Cell, _LastScannedRow
-from google.cloud.bigtable.exceptions import InvalidChunk
+from google.cloud.bigtable.data.row import Row, Cell, _LastScannedRow
+from google.cloud.bigtable.data.exceptions import InvalidChunk
 
 """
 This module provides classes for the read_rows state machine:

google/cloud/bigtable/data/exceptions.py

@@ -19,13 +19,13 @@
 from typing import Any, TYPE_CHECKING
 
 from google.api_core import exceptions as core_exceptions
-from google.cloud.bigtable.row import Row
+from google.cloud.bigtable.data.row import Row
 
 is_311_plus = sys.version_info >= (3, 11)
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable.mutations import RowMutationEntry
-    from google.cloud.bigtable.read_rows_query import ReadRowsQuery
+    from google.cloud.bigtable.data.mutations import RowMutationEntry
+    from google.cloud.bigtable.data.read_rows_query import ReadRowsQuery
 
 
 class IdleTimeout(core_exceptions.DeadlineExceeded):

google/cloud/bigtable/data/mutations.py

@@ -19,7 +19,7 @@
 from abc import ABC, abstractmethod
 from sys import getsizeof
 
-from google.cloud.bigtable.read_modify_write_rules import MAX_INCREMENT_VALUE
+from google.cloud.bigtable.data.read_modify_write_rules import MAX_INCREMENT_VALUE
 
 # special value for SetCell mutation timestamps. If set, server will assign a timestamp
 SERVER_SIDE_TIMESTAMP = -1

google/cloud/bigtable/data/read_rows_query.py

@@ -18,11 +18,11 @@
 from bisect import bisect_right
 from collections import defaultdict
 from dataclasses import dataclass
-from google.cloud.bigtable.row_filters import RowFilter
+from google.cloud.bigtable.data.row_filters import RowFilter
 
 if TYPE_CHECKING:
-    from google.cloud.bigtable import RowKeySamples
-    from google.cloud.bigtable import ShardedQuery
+    from google.cloud.bigtable.data import RowKeySamples
+    from google.cloud.bigtable.data import ShardedQuery
 
 
 @dataclass

tests/unit/data/_async/test_client.py

@@ -1294,12 +1294,12 @@ async def test_read_rows_per_request_timeout(
 
     @pytest.mark.asyncio
     async def test_read_rows_idle_timeout(self):
-        from google.cloud.bigtable.data._async.client import ReadRowsIterator
+        from google.cloud.bigtable.data._async.client import ReadRowsIteratorAsync
         from google.cloud.bigtable_v2.services.bigtable.async_client import (
             BigtableAsyncClient,
         )
         from google.cloud.bigtable.data.exceptions import IdleTimeout
-        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperation
+        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync
 
         chunks = [
             self._make_chunk(row_key=b"test_1"),
@@ -1310,15 +1310,15 @@ async def test_read_rows_idle_timeout(self):
                 chunks
             )
             with mock.patch.object(
-                ReadRowsIterator, "_start_idle_timer"
+                ReadRowsIteratorAsync, "_start_idle_timer"
             ) as start_idle_timer:
                 client = self._make_client()
                 table = client.get_table("instance", "table")
                 query = ReadRowsQuery()
                 gen = await table.read_rows_stream(query)
             # should start idle timer on creation
             start_idle_timer.assert_called_once()
-        with mock.patch.object(_ReadRowsOperation, "aclose", AsyncMock()) as aclose:
+        with mock.patch.object(_ReadRowsOperationAsync, "aclose", AsyncMock()) as aclose:
             # start idle timer with our own value
             await gen._start_idle_timer(0.1)
             # should timeout after being abandoned
@@ -1397,13 +1397,13 @@ async def test_read_rows_revise_request(self):
         """
         Ensure that _revise_request is called between retries
         """
-        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperation
+        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync
         from google.cloud.bigtable.data.exceptions import InvalidChunk
 
         with mock.patch.object(
-            _ReadRowsOperation, "_revise_request_rowset"
+            _ReadRowsOperationAsync, "_revise_request_rowset"
         ) as revise_rowset:
-            with mock.patch.object(_ReadRowsOperation, "aclose"):
+            with mock.patch.object(_ReadRowsOperationAsync, "aclose"):
                 revise_rowset.return_value = "modified"
                 async with self._make_table() as table:
                     read_rows = table.client._gapic_client.read_rows
@@ -1431,11 +1431,11 @@ async def test_read_rows_default_timeouts(self):
         """
         Ensure that the default timeouts are set on the read rows operation when not overridden
         """
-        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperation
+        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync
 
         operation_timeout = 8
         per_request_timeout = 4
-        with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op:
+        with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op:
             mock_op.side_effect = RuntimeError("mock error")
             async with self._make_table(
                 default_operation_timeout=operation_timeout,
@@ -1454,11 +1454,11 @@ async def test_read_rows_default_timeout_override(self):
         """
         When timeouts are passed, they overwrite default values
         """
-        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperation
+        from google.cloud.bigtable.data._async._read_rows import _ReadRowsOperationAsync
 
         operation_timeout = 8
         per_request_timeout = 4
-        with mock.patch.object(_ReadRowsOperation, "__init__") as mock_op:
+        with mock.patch.object(_ReadRowsOperationAsync, "__init__") as mock_op:
             mock_op.side_effect = RuntimeError("mock error")
             async with self._make_table(
                 default_operation_timeout=99, default_per_request_timeout=97