perf(spanner): multiple micro-optimizations for Spanner

packages/google-cloud-spanner/.cross_sync/generate.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from __future__ import annotations
-from typing import Sequence
+
 import ast
+from typing import Sequence
+
 """
 Entrypoint for initiating an async -> sync conversion using CrossSync
 
@@ -35,12 +37,13 @@ def extract_header_comments(file_path) -> str:
                 header.append(line)
             else:
                 break
-    header.append("\n# This file is automatically generated by CrossSync. Do not edit manually.\n\n")
+    header.append(
+        "\n# This file is automatically generated by CrossSync. Do not edit manually.\n\n"
+    )
     return "".join(header)
 
 
 class CrossSyncOutputFile:
-
     def __init__(self, output_path: str, ast_tree, header: str | None = None):
         self.output_path = output_path
         self.tree = ast_tree
@@ -56,15 +59,19 @@ def render(self, with_formatter=True, save_to_disk: bool = True) -> str:
         """
         full_str = self.header + ast.unparse(self.tree)
         if with_formatter:
-            import black  # type: ignore
-            import autoflake  # type: ignore
-
-            full_str = black.format_str(
-                autoflake.fix_code(full_str, remove_all_unused_imports=True),
-                mode=black.FileMode(),
-            )
+            try:
+                import autoflake  # type: ignore
+                import black  # type: ignore
+
+                full_str = black.format_str(
+                    autoflake.fix_code(full_str, remove_all_unused_imports=True),
+                    mode=black.FileMode(),
+                )
+            except ImportError:
+                pass
         if save_to_disk:
             import os
+
             os.makedirs(os.path.dirname(self.output_path), exist_ok=True)
             with open(self.output_path, "w") as f:
                 f.write(full_str)
@@ -73,10 +80,15 @@ def render(self, with_formatter=True, save_to_disk: bool = True) -> str:
 
 def convert_files_in_dir(directory: str) -> set[CrossSyncOutputFile]:
     import glob
+    import os
+
     from transformers import CrossSyncFileProcessor
 
-    # find all python files in the directory
-    files = glob.glob(directory + "/**/*.py", recursive=True)
+    # find all python files in the directory or use single file
+    if os.path.isfile(directory):
+        files = [directory]
+    else:
+        files = glob.glob(directory + "/**/*.py", recursive=True)
     # keep track of the output files pointed to by the annotated classes
     artifacts: set[CrossSyncOutputFile] = set()
     file_transformer = CrossSyncFileProcessor()

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/_helpers.py

@@ -1,5 +1,6 @@
 import asyncio
 import inspect
+import os
 import time
 
 from google.api_core.exceptions import Aborted
@@ -8,7 +9,8 @@
 async def _delay_until_retry(exc, deadline, attempts, default_retry_delay=None):
     from google.cloud.spanner_v1._helpers import _get_retry_delay
 
-    cause = exc.errors[0] if hasattr(exc, "errors") and exc.errors else exc
+    errors = getattr(exc, "errors", None)
+    cause = errors[0] if errors else exc
     now = time.time()
     if now >= deadline:
         raise exc
@@ -153,3 +155,44 @@ def _create_experimental_host_transport(
         client_key,
         interceptors=interceptors,
     )
+
+
+_PENDING_DRAIN_TASKS = set()
+if hasattr(os, "register_at_fork"):
+    os.register_at_fork(after_in_child=_PENDING_DRAIN_TASKS.clear)
+
+
+def _drain_stream(iterator):
+    """Drain an async stream iterator to EOF in the background.
+
+    Called when PartialResultSet.last is True to allow the caller to return immediately
+    while consuming trailing gRPC metadata so the stream terminates cleanly with status OK.
+    """
+    if iterator is None:
+        return
+
+    async def _drain():
+        try:
+            async for _ in iterator:
+                pass
+        except asyncio.CancelledError:
+            if hasattr(iterator, "cancel"):
+                try:
+                    iterator.cancel()
+                except Exception:
+                    pass
+            raise
+        except Exception:
+            pass
+
+    try:
+        task = asyncio.create_task(_drain())
+        _PENDING_DRAIN_TASKS.add(task)
+        task.add_done_callback(_PENDING_DRAIN_TASKS.discard)
+    except RuntimeError:
+        # Event loop may be closed or not running.
+        if hasattr(iterator, "cancel"):
+            try:
+                iterator.cancel()
+            except Exception:
+                pass

packages/google-cloud-spanner/google/cloud/spanner_v1/_async/database_sessions_manager.py

@@ -123,8 +123,8 @@ async def _get_multiplexed_session(self) -> Session:
         """Returns a multiplexed session from the database session manager.
 
         If the multiplexed session is not defined, creates a new multiplexed
-        session and starts a maintenance thread to periodically delete and
-        recreate it so that it remains valid. Otherwise, simply returns the
+        session and starts a maintenance thread to periodically rotate
+        it so that it remains valid. Otherwise, simply returns the
         current multiplexed session.
 
         :rtype: :class:`~google.cloud.spanner_v1.session.Session`
@@ -167,8 +167,8 @@ def _build_maintenance_thread(
         self, session: Optional[Session] = None
     ) -> CrossSync.Task:
         """Builds and returns a multiplexed session maintenance thread for
-        the database session manager. This thread will periodically delete
-        and recreate the multiplexed session to ensure that it is always valid.
+        the database session manager. This thread will periodically rotate
+        the multiplexed session to ensure that it is always valid.
 
         :type session: :class:`~google.cloud.spanner_v1.session.Session`
         :param session: (Optional) The multiplexed session to maintain.
@@ -209,25 +209,18 @@ async def _rotate_multiplexed_session(self) -> bool:
             return False
 
         async with self._multiplexed_session_lock:
-            old_session = self._multiplexed_session
             self._multiplexed_session = new_session
 
-        if old_session is not None:
-            try:
-                await CrossSync.run_if_async(old_session.delete)
-            except Exception:
-                pass
-
         return True
 
     @staticmethod
     @CrossSync.convert
     async def _maintain_multiplexed_session(session_manager_ref) -> None:
         """Maintains the multiplexed session for the database session manager.
 
-        This method will delete and recreate the referenced database session manager's
+        This method will periodically rotate the referenced database session manager's
         multiplexed session to ensure that it is always valid. The method will run until
-        the database session manager is deleted or the multiplexed session is deleted.
+        the database session manager is garbage collected or the session manager is closed.
 
         :type session_manager_ref: :class:`_weakref.ReferenceType`
         :param session_manager_ref: A weak reference to the database session manager."""
@@ -292,7 +285,4 @@ async def close(self) -> None:
                     pass
             else:
                 self._multiplexed_session_thread.join()
-        if self._multiplexed_session is not None:
-            session_to_delete = self._multiplexed_session
-            self._multiplexed_session = None
-            await session_to_delete.delete()
+        self._multiplexed_session = None