test: stabilize async test suite

Author: sinhasubham

google/cloud/spanner_v1/__init__.py

@@ -20,6 +20,31 @@
 
 __version__: str = package_version.__version__
 
+from google.cloud.spanner_v1 import param_types
+from google.cloud.spanner_v1._async.client import Client as AsyncClient
+from google.cloud.spanner_v1._async.pool import BurstyPool as AsyncBurstyPool
+from google.cloud.spanner_v1._async.pool import PingingPool as AsyncPingingPool
+from google.cloud.spanner_v1._async.pool import (
+    AbstractSessionPool as AsyncAbstractSessionPool,
+)
+from google.cloud.spanner_v1._async.pool import FixedSizePool as AsyncFixedSizePool
+from google.cloud.spanner_v1._async.pool import (
+    TransactionPingingPool as AsyncTransactionPingingPool,
+)
+from google.cloud.spanner_v1.client import Client
+from google.cloud.spanner_v1.keyset import KeyRange, KeySet
+from google.cloud.spanner_v1.pool import (
+    AbstractSessionPool,
+    BurstyPool,
+    FixedSizePool,
+    PingingPool,
+    TransactionPingingPool,
+)
+
+from .data_types import Interval, JsonObject
+from .exceptions import wrap_with_request_id
+from .services.spanner import SpannerAsyncClient, SpannerClient
+from .transaction import BatchTransactionId, DefaultTransactionOptions
 from .types import RequestOptions
 from .types.commit_response import CommitResponse
 from .types.keys import KeyRange as KeyRangePB
@@ -61,32 +86,6 @@
 from .types.transaction import Transaction, TransactionOptions, TransactionSelector
 from .types.type import StructType, Type, TypeAnnotationCode, TypeCode
 
-from .data_types import Interval, JsonObject
-from .exceptions import wrap_with_request_id
-from .services.spanner import SpannerAsyncClient, SpannerClient
-from .transaction import BatchTransactionId, DefaultTransactionOptions
-
-from google.cloud.spanner_v1 import param_types
-from google.cloud.spanner_v1.client import Client
-from google.cloud.spanner_v1.keyset import KeyRange, KeySet
-from google.cloud.spanner_v1.pool import (
-    AbstractSessionPool,
-    BurstyPool,
-    FixedSizePool,
-    PingingPool,
-    TransactionPingingPool,
-)
-from google.cloud.spanner_v1._async.client import Client as AsyncClient
-from google.cloud.spanner_v1._async.pool import BurstyPool as AsyncBurstyPool
-from google.cloud.spanner_v1._async.pool import PingingPool as AsyncPingingPool
-from google.cloud.spanner_v1._async.pool import (
-    AbstractSessionPool as AsyncAbstractSessionPool,
-)
-from google.cloud.spanner_v1._async.pool import FixedSizePool as AsyncFixedSizePool
-from google.cloud.spanner_v1._async.pool import (
-    TransactionPingingPool as AsyncTransactionPingingPool,
-)
-
 COMMIT_TIMESTAMP = "spanner.commit_timestamp()"
 """Placeholder be used to store commit timestamp of a transaction in a column.
 This value can only be used for timestamp columns that have set the option

google/cloud/spanner_v1/_async/_helpers.py

@@ -1,4 +1,5 @@
 import asyncio
+import inspect
 import time
 
 from google.api_core.exceptions import Aborted
@@ -7,9 +8,15 @@
 async def _delay_until_retry(exc, deadline, attempts, default_retry_delay=None):
     from google.cloud.spanner_v1._helpers import _get_retry_delay
 
-    delay = _get_retry_delay(exc, attempts, default_retry_delay)
-    if time.time() + delay > deadline:
+    cause = exc.errors[0] if hasattr(exc, "errors") and exc.errors else exc
+    now = time.time()
+    if now >= deadline:
         raise exc
+
+    delay = _get_retry_delay(cause, attempts, default_retry_delay)
+    if now + delay > deadline:
+        raise exc
+
     await asyncio.sleep(delay)
 
 
@@ -39,9 +46,14 @@ async def _retry(
     retries = 0
     while True:
         try:
-            return await func()
+            res = func()
+            if asyncio.iscoroutine(res) or inspect.isawaitable(res):
+                return await res
+            return res
         except Exception as e:
-            if allowed_exceptions and type(e) in allowed_exceptions:
+            if allowed_exceptions is not None:
+                if type(e) not in allowed_exceptions:
+                    raise e
                 _check_err = allowed_exceptions.get(type(e))
                 if callable(_check_err) and not _check_err(e):
                     raise e
@@ -80,8 +92,8 @@ def _create_experimental_host_transport(
     Raises:
         ValueError: If TLS/mTLS configuration is invalid.
     """
-    import grpc.aio
     from google.auth.credentials import AnonymousCredentials
+    import grpc.aio
 
     channel = None
     if use_plain_text:

google/cloud/spanner_v1/_async/batch.py

@@ -29,14 +29,15 @@
     RequestOptions,
     TransactionOptions,
 )
-from google.cloud.spanner_v1._async._helpers import _retry_on_aborted_exception
+from google.cloud.spanner_v1._async._helpers import _retry, _retry_on_aborted_exception
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
     _merge_client_context,
     _merge_request_options,
     _validate_client_context,
     _check_rst_stream_error,
     _make_list_value_pbs,
+
     _merge_Transaction_Options,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,

google/cloud/spanner_v1/_async/client.py

@@ -66,11 +66,9 @@
     ListInstanceConfigsRequest,
     ListInstancesRequest,
 )
-from google.cloud.spanner_v1 import (
-    DefaultTransactionOptions,
-    ExecuteSqlRequest,
-    __version__,
-)
+from google.cloud.spanner_v1.gapic_version import __version__
+from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
 from google.cloud.spanner_v1._helpers import (
     _create_experimental_host_transport,
     _validate_client_context,
@@ -402,12 +400,12 @@ def instance_admin_api(self):
                     transport=transport,
                 )
             elif self._experimental_host:
-                from google.cloud.spanner_v1._helpers import (
-                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
-                )
                 from google.cloud.spanner_v1._async._helpers import (
                     _create_experimental_host_transport as _create_experimental_host_transport_async,
                 )
+                from google.cloud.spanner_v1._helpers import (
+                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
+                )
 
                 if CrossSync.is_async:
                     transport = _create_experimental_host_transport_async(
@@ -456,12 +454,12 @@ def database_admin_api(self):
                     transport=transport,
                 )
             elif self._experimental_host:
-                from google.cloud.spanner_v1._helpers import (
-                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
-                )
                 from google.cloud.spanner_v1._async._helpers import (
                     _create_experimental_host_transport as _create_experimental_host_transport_async,
                 )
+                from google.cloud.spanner_v1._helpers import (
+                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
+                )
 
                 if CrossSync.is_async:
                     transport = _create_experimental_host_transport_async(

google/cloud/spanner_v1/_async/database.py

@@ -43,15 +43,13 @@
 from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
 from google.cloud.spanner_admin_database_v1.types import DatabaseDialect
-from google.cloud.spanner_v1 import (
-    DefaultTransactionOptions,
-    ExecuteSqlRequest,
-    RequestOptions,
-    TransactionOptions,
-    TransactionSelector,
-    Type,
-    TypeCode,
-)
+from google.cloud.spanner_v1.transaction import DefaultTransactionOptions
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
+from google.cloud.spanner_v1.types.spanner import RequestOptions
+from google.cloud.spanner_v1.types.transaction import TransactionOptions
+from google.cloud.spanner_v1.types.transaction import TransactionSelector
+from google.cloud.spanner_v1.types.type import Type
+from google.cloud.spanner_v1.types.type import TypeCode
 from google.cloud.spanner_v1._async.batch import Batch, MutationGroups
 from google.cloud.spanner_v1._async.database_sessions_manager import (
     DatabaseSessionsManager,
@@ -482,12 +480,12 @@ def spanner_api(self):
                 )
                 return self._spanner_api
             if self._experimental_host is not None:
-                from google.cloud.spanner_v1._helpers import (
-                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
-                )
                 from google.cloud.spanner_v1._async._helpers import (
                     _create_experimental_host_transport as _create_experimental_host_transport_async,
                 )
+                from google.cloud.spanner_v1._helpers import (
+                    _create_experimental_host_transport as _create_experimental_host_transport_sync,
+                )
 
                 if CrossSync.is_async:
                     transport = _create_experimental_host_transport_async(

google/cloud/spanner_v1/_async/pool.py

@@ -21,8 +21,8 @@
 
 from google.cloud.aio._cross_sync import CrossSync
 from google.cloud.exceptions import NotFound
-from google.cloud.spanner_v1 import BatchCreateSessionsRequest
-from google.cloud.spanner_v1 import Session as SessionProto
+from google.cloud.spanner_v1.types.spanner import BatchCreateSessionsRequest
+from google.cloud.spanner_v1.types.spanner import Session as SessionProto
 from google.cloud.spanner_v1._async.session import Session
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_leader_aware_routing,
@@ -35,7 +35,9 @@
 )
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 
-_NOW = datetime.datetime.utcnow  # unit tests may replace
+
+def _NOW():
+    return datetime.datetime.now(datetime.timezone.utc)
 
 
 @CrossSync.convert_class

google/cloud/spanner_v1/_async/session.py

@@ -14,7 +14,7 @@
 
 """Wrapper for Cloud Spanner Session objects."""
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.session"
-from datetime import datetime
+from datetime import datetime, timezone
 from functools import total_ordering
 import time
 from typing import MutableMapping, Optional
@@ -23,12 +23,13 @@
 from google.api_core.gapic_v1 import method
 
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1 import CreateSessionRequest, ExecuteSqlRequest
+from google.cloud.spanner_v1.types.spanner import CreateSessionRequest
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
+from google.cloud.spanner_v1._async._helpers import _delay_until_retry
 from google.cloud.spanner_v1._async.batch import Batch
 from google.cloud.spanner_v1._async.snapshot import Snapshot
 from google.cloud.spanner_v1._async.transaction import Transaction
 from google.cloud.spanner_v1._helpers import (
-    _delay_until_retry,
     _get_retry_delay,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
@@ -77,7 +78,7 @@ def __init__(self, database, labels=None, database_role=None, is_multiplexed=Fal
         self._labels: MutableMapping[str, str] = labels
         self._database_role: Optional[str] = database_role
         self._is_multiplexed: bool = is_multiplexed
-        self._last_use_time: datetime = datetime.utcnow()
+        self._last_use_time: datetime = datetime.now(timezone.utc)
 
     @property
     def _resource_info(self):
@@ -393,7 +394,7 @@ async def read(self, table, columns, keyset, index="", limit=0, column_info=None
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
-        return self.snapshot().read(
+        return await self.snapshot().read(
             table, columns, keyset, index, limit, column_info=column_info
         )
 
@@ -460,7 +461,7 @@ async def execute_sql(
         :rtype: :class:`~google.cloud.spanner_v1.streamed.StreamedResultSet`
         :returns: a result set instance which can be used to consume rows.
         """
-        return self.snapshot().execute_sql(
+        return await self.snapshot().execute_sql(
             sql,
             params,
             param_types,
@@ -599,7 +600,7 @@ async def run_in_transaction(self, func, *args, **kw):
                         "Transaction was aborted in user operation, retrying",
                         attributes,
                     )
-                    _delay_until_retry(
+                    await _delay_until_retry(
                         exc,
                         deadline,
                         attempts,
@@ -645,7 +646,7 @@ async def run_in_transaction(self, func, *args, **kw):
                         "Transaction was aborted during commit, retrying",
                         attributes,
                     )
-                    _delay_until_retry(
+                    await _delay_until_retry(
                         exc,
                         deadline,
                         attempts,

google/cloud/spanner_v1/_async/snapshot.py

@@ -27,21 +27,19 @@
 from google.protobuf.struct_pb2 import Struct
 
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1 import (
-    BeginTransactionRequest,
-    ExecuteSqlRequest,
-    Mutation,
-    PartialResultSet,
-    PartitionOptions,
-    PartitionQueryRequest,
-    PartitionReadRequest,
-    ReadRequest,
-    RequestOptions,
-    ResultSet,
-    Transaction,
-    TransactionOptions,
-    TransactionSelector,
-)
+from google.cloud.spanner_v1.types.mutation import Mutation
+from google.cloud.spanner_v1.types.result_set import PartialResultSet
+from google.cloud.spanner_v1.types.result_set import ResultSet
+from google.cloud.spanner_v1.types.spanner import BeginTransactionRequest
+from google.cloud.spanner_v1.types.spanner import ExecuteSqlRequest
+from google.cloud.spanner_v1.types.spanner import PartitionOptions
+from google.cloud.spanner_v1.types.spanner import PartitionQueryRequest
+from google.cloud.spanner_v1.types.spanner import PartitionReadRequest
+from google.cloud.spanner_v1.types.spanner import ReadRequest
+from google.cloud.spanner_v1.types.spanner import RequestOptions
+from google.cloud.spanner_v1.types.transaction import Transaction
+from google.cloud.spanner_v1.types.transaction import TransactionOptions
+from google.cloud.spanner_v1.types.transaction import TransactionSelector
 from google.cloud.spanner_v1._async._helpers import _retry
 from google.cloud.spanner_v1._async.streamed import StreamedResultSet
 from google.cloud.spanner_v1._helpers import (

google/cloud/spanner_v1/_async/streamed.py

@@ -18,7 +18,9 @@
 
 from google.cloud import exceptions
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1 import PartialResultSet, ResultSetMetadata, TypeCode
+from google.cloud.spanner_v1.types.result_set import PartialResultSet
+from google.cloud.spanner_v1.types.result_set import ResultSetMetadata
+from google.cloud.spanner_v1.types.type import TypeCode
 from google.cloud.spanner_v1._helpers import _get_type_decoder, _parse_nullable