fix lint

rahul2393 · rahul2393 · commit e8f813c0fc8d · 2025-05-24T12:02:52.000+05:30
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -199,7 +199,6 @@ def __init__(
         pool.bind(self)
         self._session_manager = DatabaseSessionsManager(database=self, pool=pool)
 
-
     @classmethod
     def from_pb(cls, database_pb, instance, pool=None):
         """Creates an instance of this class from a protobuf.
@@ -785,8 +784,8 @@ def execute_pdml():
                     # re-raising the error.
                     except NotImplementedError as exc:
                         if (
-                                "Transaction type partitioned_dml not supported with multiplexed sessions"
-                                in str(exc)
+                            "Transaction type partitioned_dml not supported with multiplexed sessions"
+                            in str(exc)
                         ):
                             self.session_options.disable_multiplexed(
                                 self.logger, TransactionType.PARTITIONED
@@ -1248,8 +1247,20 @@ def observability_options(self):
 
 class SessionCheckout(object):
     """Context manager for using a session from a database.
+
+    This is the recommended way to obtain sessions for database operations.
+    It automatically integrates with the database session manager to support
+    multiplexed sessions when enabled via environment variables.
+
+    For read-only operations, consider using ``database.snapshot()`` instead.
+
     :type database: :class:`~google.cloud.spanner_v1.database.Database`
     :param database: database to use the session from
+
+    :type transaction_type: :class:`~google.cloud.spanner_v1.session_options.TransactionType`
+    :param transaction_type: type of transaction this session will be used for.
+        Defaults to READ_WRITE. For read-only operations, use READ_ONLY to
+        enable multiplexed session support.
     """
 
     _session = None  # Not checked out until '__enter__'.
@@ -1261,6 +1272,7 @@ def __init__(
     ):
         # Move import here to avoid circular import
         from google.cloud.spanner_v1.database import Database
+
         if not isinstance(database, Database):
             raise TypeError(
                 "{class_name} must receive an instance of {expected_class_name}. Received: {actual_class_name}".format(
@@ -1290,6 +1302,7 @@ def __enter__(self):
     def __exit__(self, *ignored):
         self._database._session_manager.put_session(self._session)
 
+
 class BatchCheckout(object):
     """Context manager for using a batch from a database.
 
@@ -1339,8 +1352,12 @@ def __init__(
 
     def __enter__(self):
         """Begin ``with`` block."""
+        from google.cloud.spanner_v1.database import TransactionType
+
         current_span = get_current_span()
-        session = self._session = self._database._pool.get()
+        session = self._session = self._database._session_manager.get_session(
+            TransactionType.READ_WRITE
+        )
         add_span_event(current_span, "Using session", {"id": session.session_id})
         batch = self._batch = Batch(session)
         if self._request_options.transaction_tag:
@@ -1365,7 +1382,7 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     "CommitStats: {}".format(self._batch.commit_stats),
                     extra={"commit_stats": self._batch.commit_stats},
                 )
-            self._database._pool.put(self._session)
+            self._database._session_manager.put_session(self._session)
             current_span = get_current_span()
             add_span_event(
                 current_span,
@@ -1393,7 +1410,11 @@ def __init__(self, database):
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        from google.cloud.spanner_v1.database import TransactionType
+
+        session = self._session = self._database._session_manager.get_session(
+            TransactionType.READ_WRITE
+        )
         return MutationGroups(session)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1402,9 +1423,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # If NotFound exception occurs inside the with block
             # then we validate if the session still exists.
             if not self._session.exists():
-                self._session = self._database._pool._new_session()
+                self._session = self._database._session_manager._pool._new_session()
                 self._session.create()
-        self._database._pool.put(self._session)
+        self._database._session_manager.put_session(self._session)
 
 
 class SnapshotCheckout(object):
@@ -1432,7 +1453,11 @@ def __init__(self, database, **kw):
 
     def __enter__(self):
         """Begin ``with`` block."""
-        session = self._session = self._database._pool.get()
+        from google.cloud.spanner_v1.database import TransactionType
+
+        session = self._session = self._database._session_manager.get_session(
+            TransactionType.READ_ONLY
+        )
         return Snapshot(session, **self._kw)
 
     def __exit__(self, exc_type, exc_val, exc_tb):
@@ -1441,9 +1466,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
             # If NotFound exception occurs inside the with block
             # then we validate if the session still exists.
             if not self._session.exists():
-                self._session = self._database._pool._new_session()
+                self._session = self._database._session_manager._pool._new_session()
                 self._session.create()
-        self._database._pool.put(self._session)
+        self._database._session_manager.put_session(self._session)
 
 
 class BatchSnapshot(object):
@@ -1488,8 +1513,12 @@ def from_dict(cls, database, mapping):
 
         :rtype: :class:`BatchSnapshot`
         """
+        from google.cloud.spanner_v1.database import TransactionType
+
         instance = cls(database)
-        session = instance._session = database.session()
+        session = instance._session = database._session_manager.get_session(
+            TransactionType.READ_ONLY
+        )
         session._session_id = mapping["session_id"]
         snapshot = instance._snapshot = session.snapshot()
         snapshot._transaction_id = mapping["transaction_id"]
@@ -1522,8 +1551,12 @@ def _get_session(self):
            Caller is responsible for cleaning up the session after
            all partitions have been processed.
         """
+        from google.cloud.spanner_v1.database import TransactionType
+
         if self._session is None:
-            session = self._session = self._database.session()
+            session = self._session = self._database._session_manager.get_session(
+                TransactionType.READ_ONLY
+            )
             if self._session_id is None:
                 session.create()
             else:
@@ -1978,4 +2011,3 @@ def _retry_on_aborted(func, retry_config):
     """
     retry = retry_config.with_predicate(if_exception_type(Aborted))
     return retry(func)
-
diff --git a/google/cloud/spanner_v1/database_sessions_manager.py b/google/cloud/spanner_v1/database_sessions_manager.py
@@ -256,4 +256,4 @@ def _maintain_multiplexed_session(session_manager_ref) -> None:
                     session_manager._disable_multiplexed_sessions()
                     return
 
-            session_created_time = time.time()
+            session_created_time = time.time()
diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py
@@ -137,15 +137,33 @@ def _new_session(self):
         )
 
     def session(self, **kwargs):
-        """Check out a session from the pool.
+        """Return a context manager for a session from the pool.
 
-        :param kwargs: (optional) keyword arguments, passed through to
-                       the returned checkout.
+        .. warning::
+            This method bypasses the database session manager and will NOT use
+            multiplexed sessions even if they are enabled. For normal application
+            usage, prefer using database-level session APIs like
+            ``database.snapshot()`` or ``SessionCheckout(database)`` which
+            support multiplexed sessions.
 
-        :rtype: :class:`~google.cloud.spanner_v1.session.SessionCheckout`
-        :returns: a checkout instance, to be used as a context manager for
-                  accessing the session and returning it to the pool.
+        This method is primarily intended for pool testing and advanced use cases
+        where direct pool control is needed.
+
+        :type kwargs: dict
+        :param kwargs: additional keyword arguments to pass to the session
+
+        :rtype: :class:`SessionCheckout`
+        :returns: a context manager that yields a session from the pool
         """
+        import warnings
+
+        warnings.warn(
+            "pool.session() bypasses multiplexed session support. "
+            "Consider using database-level APIs like database.snapshot() "
+            "or SessionCheckout(database) instead.",
+            UserWarning,
+            stacklevel=2,
+        )
         return SessionCheckout(self, **kwargs)
 
 
@@ -537,7 +555,7 @@ def bind(self, database):
         request = BatchCreateSessionsRequest(
             database=database.name,
             session_count=self.size,
-            session_template=Session(creator_role=self.database_role),
+            session_template=SessionPB(creator_role=self.database_role),
         )
 
         span_event_attributes = {"kind": type(self).__name__}
@@ -789,5 +807,37 @@ def begin_pending_transactions(self):
             super(TransactionPingingPool, self).put(session)
 
 
+class SessionCheckout(object):
+    """Context manager for checking out a session from a pool.
+
+    .. warning::
+        This is a low-level API primarily intended for pool testing and
+        internal use. It bypasses the database session manager and will
+        NOT use multiplexed sessions even if enabled.
 
+        For normal application usage, prefer database-level APIs:
+        - ``database.snapshot()`` for read-only operations
+        - ``SessionCheckout(database)`` for general session management
+        - ``database.batch()`` for batch operations
+
+    :type pool: :class:`AbstractSessionPool`
+    :param pool: the pool to check out a session from
+    :type kwargs: dict
+    :param kwargs: additional keyword arguments for the session
+    """
 
+    def __init__(self, pool, **kwargs):
+        self._pool = pool
+        self._session = None
+        self._kwargs = kwargs
+
+    def __enter__(self):
+        """Check out a session from the pool."""
+        self._session = self._pool.get(**self._kwargs)
+        return self._session
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        """Return the session to the pool."""
+        if self._session is not None:
+            self._pool.put(self._session)
+            self._session = None
diff --git a/google/cloud/spanner_v1/session_options.py b/google/cloud/spanner_v1/session_options.py
@@ -66,9 +66,8 @@ def use_multiplexed(self, transaction_type: TransactionType) -> bool:
         """
 
         if transaction_type is TransactionType.READ_ONLY:
-            return (
-                self._is_multiplexed_enabled[transaction_type]
-                and self._getenv(self.ENV_VAR_ENABLE_MULTIPLEXED)
+            return self._is_multiplexed_enabled[transaction_type] and self._getenv(
+                self.ENV_VAR_ENABLE_MULTIPLEXED
             )
 
         elif transaction_type is TransactionType.PARTITIONED:
@@ -80,7 +79,7 @@ def use_multiplexed(self, transaction_type: TransactionType) -> bool:
 
         elif transaction_type is TransactionType.READ_WRITE:
             return False
- 
+
         raise ValueError(f"Transaction type {transaction_type} is not supported.")
 
     def disable_multiplexed(
@@ -104,11 +103,15 @@ def disable_multiplexed(
         if transaction_type is None:
             if logger:
                 logger.warning(
-                    disable_multiplexed_log_msg_fstring.format(transaction_type_value="all")
+                    disable_multiplexed_log_msg_fstring.format(
+                        transaction_type_value="all"
+                    )
                 )
             else:
                 print(
-                    disable_multiplexed_log_msg_fstring.format(transaction_type_value="all")
+                    disable_multiplexed_log_msg_fstring.format(
+                        transaction_type_value="all"
+                    )
                 )
             for transaction_type in TransactionType:
                 self._is_multiplexed_enabled[transaction_type] = False
@@ -123,7 +126,9 @@ def disable_multiplexed(
                 )
             else:
                 print(
-                    disable_multiplexed_log_msg_fstring.format(transaction_type_value=transaction_type.value)
+                    disable_multiplexed_log_msg_fstring.format(
+                        transaction_type_value=transaction_type.value
+                    )
                 )
             self._is_multiplexed_enabled[transaction_type] = False
             return
@@ -137,4 +142,4 @@ def _getenv(name: str) -> bool:
         considered false.
         """
         env_var = os.getenv(name, "").lower().strip()
-        return env_var in ["1", "true"]
+        return env_var in ["1", "true"]
diff --git a/google/cloud/spanner_v1/snapshot.py b/google/cloud/spanner_v1/snapshot.py
@@ -196,6 +196,7 @@ def _restart_on_unavailable(
 
         del item_buffer[:]
 
+
 def _handle_not_implemented_error(exception, database) -> None:
     """Handles NotImplementedError for the database. If the error is due to unsupported
     partitioned operations with multiplexed sessions, disables multiplexed sessions for
@@ -212,10 +213,13 @@ def _handle_not_implemented_error(exception, database) -> None:
         exception
     ):
         session_options = database.session_options
-        session_options.disable_multiplexed(database.logger, TransactionType.PARTITIONED)
+        session_options.disable_multiplexed(
+            database.logger, TransactionType.PARTITIONED
+        )
 
     raise exception
 
+
 class _SnapshotBase(_SessionWrapper):
     """Base class for Snapshot.
 
@@ -614,7 +618,6 @@ def wrapped_restart(*args, **kwargs):
                 return self._get_streamed_result_set(**get_streamed_result_set_args)
         return self._get_streamed_result_set(**get_streamed_result_set_args)
 
-
     def partition_read(
         self,
         table,
@@ -856,14 +859,14 @@ def attempt_tracking_method():
         return [partition.partition_token for partition in response.partitions]
 
     def _get_streamed_result_set(
-            self,
-            method,
-            request,
-            metadata,
-            trace_attributes,
-            column_info,
-            observability_options,
-            lazy_decode,
+        self,
+        method,
+        request,
+        metadata,
+        trace_attributes,
+        column_info,
+        observability_options,
+        lazy_decode,
     ):
         """Returns the streamed result set for a read or execute SQL request with the given arguments."""
 
@@ -901,6 +904,7 @@ def _get_streamed_result_set(
 
         return StreamedResultSet(**streamed_result_set_args)
 
+
 class Snapshot(_SnapshotBase):
     """Allow a set of reads / SQL statements with shared staleness.
 
diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py
diff --git a/tests/unit/test_pool.py b/tests/unit/test_pool.py
