feat: Multiplexed sessions - Update Connection to use multiplexed sessions, add unit tests.

Signed-off-by: Taylor Curran <[email protected]>

Author: currantw

google/cloud/spanner_dbapi/connection.py

@@ -28,6 +28,7 @@
 from google.cloud.spanner_dbapi.transaction_helper import TransactionRetryHelper
 from google.cloud.spanner_dbapi.cursor import Cursor
 from google.cloud.spanner_v1 import RequestOptions, TransactionOptions
+from google.cloud.spanner_v1.session_options import TransactionType
 from google.cloud.spanner_v1.snapshot import Snapshot
 
 from google.cloud.spanner_dbapi.exceptions import (
@@ -356,8 +357,16 @@ def _session_checkout(self):
         """
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
+
         if not self._session:
-            self._session = self.database._pool.get()
+            transaction_type = (
+                TransactionType.READ_ONLY
+                if self.read_only
+                else TransactionType.READ_WRITE
+            )
+            self._session = self.database._sessions_manager.get_session(
+                transaction_type
+            )
 
         return self._session
 
@@ -368,9 +377,11 @@ def _release_session(self):
         """
         if self._session is None:
             return
+
         if self.database is None:
             raise ValueError("Database needs to be passed for this operation")
-        self.database._pool.put(self._session)
+
+        self.database._sessions_manager.put_session(self._session)
         self._session = None
 
     def transaction_checkout(self):
@@ -432,7 +443,7 @@ def close(self):
             self._transaction.rollback()
 
         if self._own_pool and self.database:
-            self.database._pool.clear()
+            self.database._sessions_manager._pool.clear()
 
         self.is_closed = True
 

tests/_builders.py

@@ -11,8 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+from logging import Logger
 from mock import create_autospec
+from typing import Mapping
+
+from google.cloud.spanner_dbapi import Connection
+from google.cloud.spanner_v1 import SpannerClient
+from google.cloud.spanner_v1.client import Client
+from google.cloud.spanner_v1.database import Database
+from google.cloud.spanner_v1.instance import Instance
+from google.cloud.spanner_v1.session import Session
 
 # Default values used to populate required or expected attributes.
 # Tests should not depend on them: if a test requires a specific
@@ -22,67 +30,78 @@
 _DATABASE_ID = "default-database-id"
 
 
-def build_logger():
-    """Builds and returns a logger for testing."""
-    from logging import Logger
-
-    return create_autospec(Logger, instance=True)
-
-
-# Client objects
-# --------------
-
-
-def build_client(**kwargs):
+def build_client(**kwargs: Mapping) -> Client:
     """Builds and returns a client for testing using the given arguments.
     If a required argument is not provided, a default value will be used."""
-    from google.cloud.spanner_v1 import Client
 
     if "project" not in kwargs:
         kwargs["project"] = _PROJECT_ID
 
     return Client(**kwargs)
 
 
-def build_database(**kwargs):
+def build_connection(**kwargs: Mapping) -> Connection:
+    """Builds and returns a connection for testing using the given arguments.
+    If a required argument is not provided, a default value will be used."""
+
+    if "instance" not in kwargs:
+        kwargs["instance"] = build_instance()
+
+    if "database" not in kwargs:
+        kwargs["database"] = build_database(instance=kwargs["instance"])
+
+    return Connection(**kwargs)
+
+
+def build_database(**kwargs: Mapping) -> Database:
     """Builds and returns a database for testing using the given arguments.
-    If a required argument is not provided, a default value will be used.."""
-    from google.cloud.spanner_v1.database import Database
+    If a required argument is not provided, a default value will be used."""
 
     if "database_id" not in kwargs:
         kwargs["database_id"] = _DATABASE_ID
 
     if "logger" not in kwargs:
         kwargs["logger"] = build_logger()
 
-    if "instance" not in kwargs or isinstance(kwargs["instance"], dict):
-        instance_args = kwargs.pop("instance", {})
-        kwargs["instance"] = build_instance(**instance_args)
+    if "instance" not in kwargs:
+        kwargs["instance"] = build_instance()
 
     database = Database(**kwargs)
     database._spanner_api = build_spanner_api()
 
     return database
 
 
-def build_instance(**kwargs):
+def build_instance(**kwargs: Mapping) -> Instance:
     """Builds and returns an instance for testing using the given arguments.
     If a required argument is not provided, a default value will be used."""
-    from google.cloud.spanner_v1.instance import Instance
 
     if "instance_id" not in kwargs:
         kwargs["instance_id"] = _INSTANCE_ID
 
-    if "client" not in kwargs or isinstance(kwargs["client"], dict):
-        client_args = kwargs.pop("client", {})
-        kwargs["client"] = build_client(**client_args)
+    if "client" not in kwargs:
+        kwargs["client"] = build_client()
 
     return Instance(**kwargs)
 
 
-def build_spanner_api():
+def build_logger() -> Logger:
+    """Builds and returns a logger for testing."""
+    return create_autospec(Logger, instance=True)
+
+
+def build_session(**kwargs: Mapping) -> Session:
+    """Builds and returns a session for testing using the given arguments.
+    If a required argument is not provided, a default value will be used."""
+
+    if "database" not in kwargs:
+        kwargs["database"] = build_database()
+
+    return Session(**kwargs)
+
+
+def build_spanner_api() -> SpannerClient:
     """Builds and returns a mock Spanner Client API for testing using the given arguments.
     Commonly used methods are mocked to return default values."""
-    from google.cloud.spanner_v1 import SpannerClient
 
     return create_autospec(SpannerClient, instance=True)

tests/unit/spanner_dbapi/test_connection.py

@@ -37,6 +37,8 @@
     ClientSideStatementType,
     AutocommitDmlMode,
 )
+from google.cloud.spanner_v1.session_options import TransactionType
+from tests._builders import build_connection, build_session
 
 PROJECT = "test-project"
 INSTANCE = "test-instance"
@@ -151,42 +153,48 @@ def test_read_only_connection(self):
         connection.read_only = False
         self.assertFalse(connection.read_only)
 
-    @staticmethod
-    def _make_pool():
-        from google.cloud.spanner_v1.pool import AbstractSessionPool
+    def test__session_checkout_read_only(self):
+        connection = build_connection(read_only=True)
+        database = connection._database
+        sessions_manager = database._sessions_manager
 
-        return mock.create_autospec(AbstractSessionPool)
+        expected_session = build_session(database=database)
+        sessions_manager.get_session = mock.MagicMock(return_value=expected_session)
 
-    @mock.patch("google.cloud.spanner_v1.database.Database")
-    def test__session_checkout(self, mock_database):
-        pool = self._make_pool()
-        mock_database._pool = pool
-        connection = Connection(INSTANCE, mock_database)
+        actual_session = connection._session_checkout()
+
+        self.assertEqual(actual_session, expected_session)
+        sessions_manager.get_session.assert_called_once_with(TransactionType.READ_ONLY)
+
+    def test__session_checkout_read_write(self):
+        connection = build_connection(read_only=False)
+        database = connection._database
+        sessions_manager = database._sessions_manager
+
+        expected_session = build_session(database=database)
+        sessions_manager.get_session = mock.MagicMock(return_value=expected_session)
 
-        connection._session_checkout()
-        pool.get.assert_called_once_with()
-        self.assertEqual(connection._session, pool.get.return_value)
+        actual_session = connection._session_checkout()
 
-        connection._session = "db_session"
-        connection._session_checkout()
-        self.assertEqual(connection._session, "db_session")
+        self.assertEqual(actual_session, expected_session)
+        sessions_manager.get_session.assert_called_once_with(TransactionType.READ_WRITE)
 
     def test_session_checkout_database_error(self):
         connection = Connection(INSTANCE)
 
         with pytest.raises(ValueError):
             connection._session_checkout()
 
-    @mock.patch("google.cloud.spanner_v1.database.Database")
-    def test__release_session(self, mock_database):
-        pool = self._make_pool()
-        mock_database._pool = pool
-        connection = Connection(INSTANCE, mock_database)
-        connection._session = "session"
+    def test__release_session(self):
+        connection = build_connection()
+        sessions_manager = connection._database._sessions_manager
+
+        session = connection._session = build_session(database=connection._database)
+        put_session = sessions_manager.put_session = mock.MagicMock()
 
         connection._release_session()
-        pool.put.assert_called_once_with("session")
-        self.assertIsNone(connection._session)
+
+        put_session.assert_called_once_with(session)
 
     def test_release_session_database_error(self):
         connection = Connection(INSTANCE)