chore: resolve rebase conflicts from main and stabilize async tests

Author: sinhasubham

google/cloud/spanner_dbapi/batch_dml_executor.py

@@ -97,6 +97,7 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
         many_result_set.add_iter(res)
         cursor._row_count = sum([max(val, 0) for val in res])
     else:
+        retry_count = 0
         while True:
             try:
                 transaction = connection.transaction_checkout()
@@ -108,6 +109,9 @@ def run_batch_dml(cursor: "Cursor", statements: List[Statement]):
                     if not transaction._transaction_id:
                         # This should normally not happen,
                         # but we safeguard against it just to be sure.
+                        if retry_count > 0:
+                            raise OperationalError(status.message)
+                        retry_count += 1
                         transaction._reset_and_begin()
                         continue
                     raise OperationalError(status.message)

google/cloud/spanner_dbapi/cursor.py

@@ -439,6 +439,18 @@ def executemany(self, operation, seq_of_params):
 
             self._result_set = many_result_set
             self._itr = many_result_set
+        except (AlreadyExists, FailedPrecondition, OutOfRange) as e:
+            exception = IntegrityError(getattr(e, "details", e))
+            exception.__cause__ = e
+            raise exception
+        except InvalidArgument as e:
+            exception = ProgrammingError(getattr(e, "details", e))
+            exception.__cause__ = e
+            raise exception
+        except InternalServerError as e:
+            exception = OperationalError(getattr(e, "details", e))
+            exception.__cause__ = e
+            raise exception
         except Exception as e:
             exception = e
             raise

google/cloud/spanner_v1/_async/batch.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """Context manager for Cloud Spanner batched writes."""
+
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.batch"
 import functools
 import time
@@ -30,7 +31,6 @@
     _validate_client_context,
     _check_rst_stream_error,
     _make_list_value_pbs,
-
     _merge_Transaction_Options,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
@@ -441,8 +441,6 @@ def wrapped_method():
                 )
                 return batch_write_method()
 
-            from google.cloud.spanner_v1._async._helpers import _retry
-
             response = await _retry(
                 wrapped_method,
                 allowed_exceptions={

google/cloud/spanner_v1/_async/client.py

@@ -23,6 +23,7 @@
 * a :class:`~google.cloud.spanner_v1.instance.Instance` owns a
   :class:`~google.cloud.spanner_v1.database.Database`
 """
+
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.client"
 import logging
 import os
@@ -69,7 +70,6 @@
 from google.cloud.spanner_v1._async.instance import Instance
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
-    _create_experimental_host_transport,
     _merge_query_options,
     _metadata_with_prefix,
     _validate_client_context,

google/cloud/spanner_v1/_async/database.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 """User-friendly container for Cloud Spanner Database."""
+
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.database"
 import copy
 import functools
@@ -230,11 +231,7 @@ def __init__(
         self._experimental_host = (
             self._instance.experimental_host if self._instance else None
         )
-        is_experimental_host = self._experimental_host is not None
-
-        self._sessions_manager = DatabaseSessionsManager(
-            self, pool 
-        )
+        self._sessions_manager = DatabaseSessionsManager(self, pool)
 
     @property
     def _resource_info(self):
@@ -1522,7 +1519,9 @@ async def __aenter__(self):
             event_attributes={"id": self._session.session_id},
         )
 
-        batch = self._batch = Batch(session=self._session, client_context=self._client_context)
+        batch = self._batch = Batch(
+            session=self._session, client_context=self._client_context
+        )
         if self._request_options.transaction_tag:
             batch.transaction_tag = self._request_options.transaction_tag
 
@@ -1588,7 +1587,9 @@ async def __aenter__(self):
             transaction_type
         )
 
-        return MutationGroups(session=self._session, client_context=self._client_context)
+        return MutationGroups(
+            session=self._session, client_context=self._client_context
+        )
 
     @CrossSync.convert(sync_name="__exit__")
     async def __aexit__(self, exc_type, exc_val, exc_tb):
@@ -1711,7 +1712,9 @@ def from_dict(cls, database, mapping):
         instance._session_id = session._session_id = mapping["session_id"]
         instance._client_context = mapping.get("client_context")
 
-        snapshot = instance._snapshot = session.snapshot(client_context=instance._client_context)
+        snapshot = instance._snapshot = session.snapshot(
+            client_context=instance._client_context
+        )
         instance._transaction_id = snapshot._transaction_id = mapping["transaction_id"]
 
         return instance

google/cloud/spanner_v1/_async/instance.py

@@ -34,7 +34,12 @@
 from google.cloud.spanner_admin_database_v1.types import backup, spanner_database_admin
 from google.cloud.spanner_admin_instance_v1 import Instance as InstancePB
 from google.cloud.spanner_v1._async.database import Database
-from google.cloud.spanner_v1.testing.database_test import TestDatabase
+
+if CrossSync.is_async:
+    from google.cloud.spanner_v1._async.testing.database_test import TestDatabase
+else:
+    from google.cloud.spanner_v1.testing.database_test import TestDatabase
+
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
 from google.cloud.spanner_v1.backup import Backup
 
@@ -524,7 +529,9 @@ async def database(
                 enable_drop_protection=enable_drop_protection,
             )
 
-        await db._pool.bind(db)
+        res = db._pool.bind(db)
+        if res is not None:
+            await res
         return db
 
     @CrossSync.convert

google/cloud/spanner_v1/_async/transaction.py

@@ -13,13 +13,15 @@
 # limitations under the License.
 
 """Spanner read-write transaction support."""
+
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.transaction"
 
 from dataclasses import dataclass, field
 import functools
 from typing import Any, Optional
 
 from google.cloud.spanner_v1._helpers import (
+    AtomicCounter,
     _make_value_pb,
     _merge_query_options,
     _metadata_with_prefix,
@@ -35,18 +37,8 @@
 from google.protobuf.struct_pb2 import Struct
 
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1._async._helpers import _retry
 from google.cloud.spanner_v1._async.batch import _BatchBase
 from google.cloud.spanner_v1._async.snapshot import _SnapshotBase
-from google.cloud.spanner_v1._helpers import (
-    AtomicCounter,
-    _check_rst_stream_error,
-    _make_value_pb,
-    _merge_query_options,
-    _merge_Transaction_Options,
-    _metadata_with_leader_aware_routing,
-    _metadata_with_prefix,
-)
 from google.cloud.spanner_v1._opentelemetry_tracing import add_span_event, trace_call
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
 from google.cloud.spanner_v1.types.commit_response import CommitResponse

google/cloud/spanner_v1/_helpers.py

@@ -32,8 +32,6 @@
 from google.rpc.error_details_pb2 import RetryInfo
 
 from google.cloud._helpers import _date_from_iso8601_date
-from google.cloud.spanner_v1.types import ExecuteSqlRequest
-from google.cloud.spanner_v1.types import TransactionOptions
 from google.cloud.spanner_v1.types import ClientContext
 from google.cloud.spanner_v1.types import RequestOptions
 from google.cloud.spanner_v1.data_types import JsonObject, Interval

google/cloud/spanner_v1/batch.py

@@ -20,7 +20,7 @@
 import time
 from typing import List, Optional
 from google.api_core.exceptions import InternalServerError
-from google.cloud.spanner_v1._helpers import _retry_on_aborted_exception
+from google.cloud.spanner_v1._helpers import _retry, _retry_on_aborted_exception
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
     _merge_client_context,
@@ -382,8 +382,6 @@ def wrapped_method():
                 )
                 return batch_write_method()
 
-            from google.cloud.spanner_v1._helpers import _retry
-
             response = _retry(
                 wrapped_method,
                 allowed_exceptions={InternalServerError: _check_rst_stream_error},

google/cloud/spanner_v1/database.py

@@ -204,7 +204,6 @@ def __init__(
         self._experimental_host = (
             self._instance.experimental_host if self._instance else None
         )
-        is_experimental_host = self._experimental_host is not None
         self._sessions_manager = DatabaseSessionsManager(self, pool)
 
     @property