test: add mockserver tests for asyncIO operations

Author: sinhasubham

google/cloud/aio/_cross_sync/cross_sync.py

@@ -91,6 +91,8 @@ class CrossSync(metaclass=MappingMeta):
     Task: TypeAlias = asyncio.Task
     Event: TypeAlias = asyncio.Event
     Semaphore: TypeAlias = asyncio.Semaphore
+    LifoQueue: TypeAlias = asyncio.LifoQueue
+    PriorityQueue: TypeAlias = asyncio.PriorityQueue
     StopIteration: TypeAlias = StopAsyncIteration
     # provide aliases for common async type annotations
     Awaitable: TypeAlias = typing.Awaitable
@@ -160,6 +162,23 @@ async def run_if_async(func, *args, **kwargs):
             return await res
         return res
 
+    @staticmethod
+    async def queue_get(queue, block=True, timeout=None):
+        if not block:
+            return queue.get_nowait()
+        if timeout is not None:
+            return await asyncio.wait_for(queue.get(), timeout=timeout)
+        return await queue.get()
+
+    @staticmethod
+    async def queue_put(queue, item, block=True, timeout=None):
+        if not block:
+            return queue.put_nowait(item)
+        if timeout is not None:
+             await asyncio.wait_for(queue.put(item), timeout=timeout)
+        else:
+             await queue.put(item)
+
     @staticmethod
     async def gather_partials(
         partial_list: Sequence[Callable[[], Awaitable[T]]],
@@ -288,6 +307,8 @@ class _Sync_Impl(metaclass=MappingMeta):
         Task: TypeAlias = concurrent.futures.Future
         Event: TypeAlias = threading.Event
         Semaphore: TypeAlias = threading.Semaphore
+        LifoQueue: TypeAlias = queue.LifoQueue
+        PriorityQueue: TypeAlias = queue.PriorityQueue
         StopIteration: TypeAlias = StopIteration
         # type annotations
         Awaitable: TypeAlias = Union[T]
@@ -304,6 +325,14 @@ def run_if_async(func, *args, **kwargs):
             """
             return func(*args, **kwargs)
 
+        @staticmethod
+        def queue_get(queue, block=True, timeout=None):
+            return queue.get(block=block, timeout=timeout)
+
+        @staticmethod
+        def queue_put(queue, item, block=True, timeout=None):
+             queue.put(item, block=block, timeout=timeout)
+
         @classmethod
         def Mock(cls, *args, **kwargs):
             from unittest.mock import Mock

google/cloud/spanner_v1/_async/client.py

@@ -41,21 +41,31 @@
 
 
 from google.cloud.spanner_admin_database_v1 import DatabaseAdminAsyncClient as DatabaseAdminClient
-from google.cloud.spanner_admin_database_v1.services.database_admin.transports.grpc import (
-    DatabaseAdminGrpcTransport,
-)
+if CrossSync.is_async:
+    from google.cloud.spanner_admin_database_v1.services.database_admin.transports.grpc_asyncio import (
+        DatabaseAdminGrpcAsyncIOTransport as DatabaseAdminGrpcTransport,
+    )
+else:
+    from google.cloud.spanner_admin_database_v1.services.database_admin.transports.grpc import (
+        DatabaseAdminGrpcTransport,
+    )
 from google.cloud.spanner_admin_instance_v1 import InstanceAdminAsyncClient as InstanceAdminClient
-from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.grpc import (
-    InstanceAdminGrpcTransport,
-)
+if CrossSync.is_async:
+    from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.grpc_asyncio import (
+        InstanceAdminGrpcAsyncIOTransport as InstanceAdminGrpcTransport,
+    )
+else:
+    from google.cloud.spanner_admin_instance_v1.services.instance_admin.transports.grpc import (
+        InstanceAdminGrpcTransport,
+    )
 from google.cloud.spanner_admin_instance_v1 import ListInstanceConfigsRequest
 from google.cloud.spanner_admin_instance_v1 import ListInstancesRequest
 from google.cloud.spanner_v1 import __version__
 from google.cloud.spanner_v1 import ExecuteSqlRequest
 from google.cloud.spanner_v1 import DefaultTransactionOptions
 from google.cloud.spanner_v1._helpers import _merge_query_options
 from google.cloud.spanner_v1._helpers import _metadata_with_prefix
-from google.cloud.spanner_v1.instance import Instance
+from google.cloud.spanner_v1._async.instance import Instance
 from google.cloud.spanner_v1.metrics.constants import (
     METRIC_EXPORT_INTERVAL_MS,
 )

google/cloud/spanner_v1/_async/database.py

@@ -22,6 +22,8 @@
 from typing import Optional
 
 import grpc
+import asyncio
+import inspect
 import logging
 import re
 import threading
@@ -64,7 +66,7 @@
 from google.cloud.spanner_v1._async.batch import MutationGroups
 from google.cloud.spanner_v1.keyset import KeySet
 from google.cloud.spanner_v1.merged_result_set import MergedResultSet
-from google.cloud.spanner_v1.pool import BurstyPool
+from google.cloud.spanner_v1._async.pool import BurstyPool
 from google.cloud.spanner_v1._async.session import Session
 from google.cloud.spanner_v1._async.database_sessions_manager import (
     DatabaseSessionsManager,
@@ -73,9 +75,14 @@
 from google.cloud.spanner_v1._async.snapshot import _restart_on_unavailable
 from google.cloud.spanner_v1._async.snapshot import Snapshot
 from google.cloud.spanner_v1._async.streamed import StreamedResultSet
-from google.cloud.spanner_v1.services.spanner.transports.grpc import (
-    SpannerGrpcTransport,
-)
+if CrossSync.is_async:
+    from google.cloud.spanner_v1.services.spanner.transports.grpc_asyncio import (
+        SpannerGrpcAsyncIOTransport as SpannerGrpcTransport,
+    )
+else:
+    from google.cloud.spanner_v1.services.spanner.transports.grpc import (
+        SpannerGrpcTransport,
+    )
 from google.cloud.spanner_v1.table import Table
 from google.cloud.spanner_v1._opentelemetry_tracing import (
     add_span_event,
@@ -205,7 +212,14 @@ def __init__(
             pool = BurstyPool(database_role=database_role)
 
         self._pool = pool
-        pool.bind(self)
+        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
         is_experimental_host = self._instance.experimental_host is not None
 
         self._sessions_manager = DatabaseSessionsManager(
@@ -448,17 +462,21 @@ def spanner_api(self):
             client_info = self._instance._client._client_info
             client_options = self._instance._client._client_options
             if self._instance.emulator_host is not None:
-                transport = SpannerGrpcTransport(
-                    channel=grpc.insecure_channel(self._instance.emulator_host)
-                )
+                if CrossSync.is_async:
+                    channel = grpc.aio.insecure_channel(self._instance.emulator_host)
+                else:
+                    channel = grpc.insecure_channel(self._instance.emulator_host)
+                transport = SpannerGrpcTransport(channel=channel)
                 self._spanner_api = SpannerClient(
                     client_info=client_info, transport=transport
                 )
                 return self._spanner_api
             if self._instance.experimental_host is not None:
-                transport = SpannerGrpcTransport(
-                    channel=grpc.insecure_channel(self._instance.experimental_host)
-                )
+                if CrossSync.is_async:
+                    channel = grpc.aio.insecure_channel(self._instance.experimental_host)
+                else:
+                    channel = grpc.insecure_channel(self._instance.experimental_host)
+                transport = SpannerGrpcTransport(channel=channel)
                 self._spanner_api = SpannerClient(
                     client_info=client_info,
                     transport=transport,