fix(spanner): ensure session pool is bound before database use

Author: sinhasubham

google/cloud/spanner_v1/_async/database.py

@@ -14,10 +14,8 @@
 
 """User-friendly container for Cloud Spanner Database."""
 __CROSS_SYNC_OUTPUT__ = "google.cloud.spanner_v1.database"
-import asyncio
 import copy
 import functools
-import inspect
 import logging
 import re
 import threading
@@ -223,14 +221,8 @@ def __init__(
             pool = BurstyPool(database_role=database_role)
 
         self._pool = pool
-        res = pool.bind(self)
-        try:
-            loop = asyncio.get_running_loop()
-            if loop.is_running() and inspect.isawaitable(res):
-                loop.create_task(res)
-        except RuntimeError:
-            # No running loop, bind should have been sync or will be failed later
-            pass
+        # Note: self._pool.bind(self) should be called via Instance.database()
+        # factory method to ensure proper async initialization.
         is_experimental_host = (
             self._instance.experimental_host is not None if self._instance else False
         )

google/cloud/spanner_v1/_async/instance.py

@@ -433,7 +433,8 @@ async def delete(self):
 
         await api.delete_instance(name=self.name, metadata=metadata)
 
-    def database(
+    @CrossSync.convert
+    async def database(
         self,
         database_id,
         ddl_statements=(),
@@ -498,7 +499,7 @@ def database(
         """
 
         if not enable_interceptors_in_tests:
-            return Database(
+            db = Database(
                 database_id,
                 self,
                 ddl_statements=ddl_statements,
@@ -511,7 +512,7 @@ def database(
                 proto_descriptors=proto_descriptors,
             )
         else:
-            return TestDatabase(
+            db = TestDatabase(
                 database_id,
                 self,
                 ddl_statements=ddl_statements,
@@ -523,6 +524,9 @@ def database(
                 enable_drop_protection=enable_drop_protection,
             )
 
+        await db._pool.bind(db)
+        return db
+
     @CrossSync.convert
     async def list_databases(self, page_size=None):
         """List databases for the instance.

google/cloud/spanner_v1/batch.py

@@ -20,8 +20,9 @@
 import functools
 import time
 from typing import List, Optional
+
 from google.api_core.exceptions import InternalServerError
-from google.cloud.spanner_v1._helpers import _retry, _retry_on_aborted_exception
+
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
     _check_rst_stream_error,
@@ -31,6 +32,8 @@
     _merge_Transaction_Options,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
+    _retry,
+    _retry_on_aborted_exception,
     _SessionWrapper,
     _validate_client_context,
 )

google/cloud/spanner_v1/client.py

@@ -32,34 +32,36 @@
 import threading
 from typing import Optional
 import warnings
+
 import google.api_core.client_options
 from google.api_core.gapic_v1 import client_info
 from google.auth.credentials import AnonymousCredentials
 import grpc
+
 from google.cloud.client import ClientWithProject
 from google.cloud.spanner_admin_database_v1 import (
     DatabaseAdminClient as DatabaseAdminClient,
 )
 from google.cloud.spanner_admin_database_v1.services.database_admin.transports.grpc import (
     DatabaseAdminGrpcTransport,
 )
+from google.cloud.spanner_admin_instance_v1 import (
+    ListInstanceConfigsRequest,
+    ListInstancesRequest,
+)
 from google.cloud.spanner_admin_instance_v1 import (
     InstanceAdminClient as InstanceAdminClient,
 )
 from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.grpc import (
     InstanceAdminGrpcTransport,
 )
-from google.cloud.spanner_admin_instance_v1 import (
-    ListInstanceConfigsRequest,
-    ListInstancesRequest,
-)
-from google.cloud.spanner_v1.instance import Instance
 from google.cloud.spanner_v1._helpers import (
     _merge_query_options,
     _metadata_with_prefix,
     _validate_client_context,
 )
 from google.cloud.spanner_v1.gapic_version import __version__
+from google.cloud.spanner_v1.instance import Instance
 from google.cloud.spanner_v1.metrics.constants import METRIC_EXPORT_INTERVAL_MS
 from google.cloud.spanner_v1.metrics.metrics_exporter import (
     CloudMonitoringMetricsExporter,

google/cloud/spanner_v1/database.py

@@ -17,21 +17,21 @@
 
 """User-friendly container for Cloud Spanner Database."""
 
-import asyncio
 import copy
 import functools
-import inspect
 import logging
 import re
 import threading
 from typing import Optional
+
 from google.api_core import gapic_v1
 from google.api_core.exceptions import Aborted
 from google.api_core.retry import Retry
 import google.auth.credentials
 from google.iam.v1 import iam_policy_pb2, options_pb2
 from google.protobuf.field_mask_pb2 import FieldMask
 import grpc
+
 from google.cloud.aio._cross_sync import CrossSync
 from google.cloud.exceptions import NotFound
 from google.cloud.spanner_admin_database_v1 import (
@@ -44,15 +44,6 @@
 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.batch import Batch, MutationGroups
-from google.cloud.spanner_v1.database_sessions_manager import (
-    DatabaseSessionsManager,
-    TransactionType,
-)
-from google.cloud.spanner_v1.pool import BurstyPool
-from google.cloud.spanner_v1.session import Session
-from google.cloud.spanner_v1.snapshot import Snapshot, _restart_on_unavailable
-from google.cloud.spanner_v1.streamed import StreamedResultSet
 from google.cloud.spanner_v1._helpers import (
     _augment_errors_with_request_id,
     _merge_query_options,
@@ -61,11 +52,30 @@
     _metadata_with_request_id,
     _metadata_with_request_id_and_req_id,
 )
+from google.cloud.spanner_v1._opentelemetry_tracing import (
+    add_span_event,
+    get_current_span,
+    trace_call,
+)
+from google.cloud.spanner_v1.batch import Batch, MutationGroups
+from google.cloud.spanner_v1.database_sessions_manager import (
+    DatabaseSessionsManager,
+    TransactionType,
+)
 from google.cloud.spanner_v1.keyset import KeySet
 from google.cloud.spanner_v1.merged_result_set import MergedResultSet
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+from google.cloud.spanner_v1.pool import BurstyPool
 from google.cloud.spanner_v1.services.spanner.client import (
     SpannerClient as SpannerClient,
 )
+from google.cloud.spanner_v1.services.spanner.transports.grpc import (
+    SpannerGrpcTransport,
+)
+from google.cloud.spanner_v1.session import Session
+from google.cloud.spanner_v1.snapshot import Snapshot, _restart_on_unavailable
+from google.cloud.spanner_v1.streamed import StreamedResultSet
+from google.cloud.spanner_v1.table import Table
 from google.cloud.spanner_v1.transaction import (
     BatchTransactionId,
     DefaultTransactionOptions,
@@ -76,16 +86,6 @@
     TransactionSelector,
 )
 from google.cloud.spanner_v1.types.type import Type, TypeCode
-from google.cloud.spanner_v1.services.spanner.transports.grpc import (
-    SpannerGrpcTransport,
-)
-from google.cloud.spanner_v1._opentelemetry_tracing import (
-    add_span_event,
-    get_current_span,
-    trace_call,
-)
-from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
-from google.cloud.spanner_v1.table import Table
 
 SPANNER_DATA_SCOPE = "https://www.googleapis.com/auth/spanner.data"
 _DATABASE_NAME_RE = re.compile(
@@ -204,13 +204,6 @@ def __init__(
         if pool is None:
             pool = BurstyPool(database_role=database_role)
         self._pool = pool
-        res = pool.bind(self)
-        try:
-            loop = asyncio.get_running_loop()
-            if loop.is_running() and inspect.isawaitable(res):
-                loop.create_task(res)
-        except RuntimeError:
-            pass
         is_experimental_host = (
             self._instance.experimental_host is not None if self._instance else False
         )

google/cloud/spanner_v1/database_sessions_manager.py

@@ -24,12 +24,13 @@
 from threading import Thread
 from typing import Optional
 from weakref import ref
+
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1.session import Session
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     add_span_event,
     get_current_span,
 )
+from google.cloud.spanner_v1.session import Session
 
 
 class TransactionType(Enum):

google/cloud/spanner_v1/instance.py

@@ -19,10 +19,12 @@
 
 import re
 import typing
+
 from google.api_core.exceptions import InvalidArgument
 import google.api_core.operation
 from google.protobuf.empty_pb2 import Empty
 from google.protobuf.field_mask_pb2 import FieldMask
+
 from google.cloud.aio._cross_sync import CrossSync
 from google.cloud.exceptions import NotFound
 from google.cloud.spanner_admin_database_v1 import (
@@ -34,10 +36,10 @@
 )
 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.database import Database
-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
+from google.cloud.spanner_v1.database import Database
+from google.cloud.spanner_v1.testing.database_test import TestDatabase
 
 _INSTANCE_NAME_RE = re.compile(
     "^projects/(?P<project>[^/]+)/instances/(?P<instance_id>[a-z][-a-z0-9]*)$"
@@ -456,7 +458,7 @@ def database(
         :rtype: :class:`~google.cloud.spanner_v1.database.Database`
         :returns: a database owned by this instance."""
         if not enable_interceptors_in_tests:
-            return Database(
+            db = Database(
                 database_id,
                 self,
                 ddl_statements=ddl_statements,
@@ -469,7 +471,7 @@ def database(
                 proto_descriptors=proto_descriptors,
             )
         else:
-            return TestDatabase(
+            db = TestDatabase(
                 database_id,
                 self,
                 ddl_statements=ddl_statements,
@@ -480,6 +482,8 @@ def database(
                 database_role=database_role,
                 enable_drop_protection=enable_drop_protection,
             )
+        db._pool.bind(db)
+        return db
 
     def list_databases(self, page_size=None):
         """List databases for the instance.

google/cloud/spanner_v1/pool.py

@@ -21,9 +21,9 @@
 import datetime
 import time
 from warnings import warn
+
 from google.cloud.aio._cross_sync import CrossSync
 from google.cloud.exceptions import NotFound
-from google.cloud.spanner_v1.session import Session
 from google.cloud.spanner_v1._helpers import (
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
@@ -34,6 +34,7 @@
     trace_call,
 )
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+from google.cloud.spanner_v1.session import Session
 from google.cloud.spanner_v1.types.spanner import BatchCreateSessionsRequest
 from google.cloud.spanner_v1.types.spanner import Session as SessionProto
 

google/cloud/spanner_v1/session.py

@@ -21,14 +21,13 @@
 from functools import total_ordering
 import time
 from typing import MutableMapping, Optional
+
 from google.api_core.exceptions import Aborted, GoogleAPICallError, NotFound
 from google.api_core.gapic_v1 import method
+
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1._helpers import _delay_until_retry
-from google.cloud.spanner_v1.batch import Batch
-from google.cloud.spanner_v1.snapshot import Snapshot
-from google.cloud.spanner_v1.transaction import Transaction
 from google.cloud.spanner_v1._helpers import (
+    _delay_until_retry,
     _get_retry_delay,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
@@ -38,7 +37,10 @@
     get_current_span,
     trace_call,
 )
+from google.cloud.spanner_v1.batch import Batch
 from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+from google.cloud.spanner_v1.snapshot import Snapshot
+from google.cloud.spanner_v1.transaction import Transaction
 from google.cloud.spanner_v1.types.spanner import (
     CreateSessionRequest,
     ExecuteSqlRequest,

google/cloud/spanner_v1/snapshot.py

@@ -19,6 +19,7 @@
 
 import functools
 from typing import List, Optional, Union
+
 from google.api_core import gapic_v1
 from google.api_core.exceptions import (
     Aborted,
@@ -27,9 +28,8 @@
     ServiceUnavailable,
 )
 from google.protobuf.struct_pb2 import Struct
+
 from google.cloud.aio._cross_sync import CrossSync
-from google.cloud.spanner_v1._helpers import _retry
-from google.cloud.spanner_v1.streamed import StreamedResultSet
 from google.cloud.spanner_v1._helpers import (
     AtomicCounter,
     _augment_error_with_request_id,
@@ -40,11 +40,13 @@
     _merge_request_options,
     _metadata_with_leader_aware_routing,
     _metadata_with_prefix,
+    _retry,
     _SessionWrapper,
     _validate_client_context,
 )
 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.streamed import StreamedResultSet
 from google.cloud.spanner_v1.types import MultiplexedSessionPrecommitToken
 from google.cloud.spanner_v1.types.mutation import Mutation
 from google.cloud.spanner_v1.types.result_set import PartialResultSet, ResultSet