feat: Add support for multiplexed sessions - part 1 (supporting classes and refactoring) by currantw · Pull Request #1329 · googleapis/python-spanner

currantw · 2025-03-26T20:31:07Z

Description

Adds SessionOptions and DatabaseSessionManager class support multiplexed session in future work.
Update Client, Instance, Database, Session and associated Checkout classes accordingly.
Add unit tests for new classes and update those for existing classes.

currantw · 2025-03-26T22:46:58Z

-        if isinstance(exc_val, NotFound):
-            # 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.create()


This check has been removed from MutationGroupsCheckout and SnapshotCheckout because:

exists only checks whether the session exists on the server (not whether the local session object exists)

This checkout wasn't for all context managers (not in BatchCheckout or SessionCheckout)

This seems to only be done as a slight optimization to potentially reduce overhead of session retrieval by adding it to the tail of a request instead of the head of another.

In reality, the optimization only plays a part if the session expires while a call was happening, and is unlikely to make a big difference.

exists checks are still performed by the session pool (for non-multiplexed session) and the database session manager (for multiplexed session).

We recommend removing this code, and instead depend on DatabaseSessionManager to take care of checking session existence and re-creating them if necessary, rather than (sometimes) handling this in the context manager.

currantw · 2025-03-26T23:25:05Z

+class SessionCheckout(object):
+    """Context manager for using a session from a database.
+
+    :type database: :class:`~google.cloud.spanner_v1.database.Database`
+    :param database: database to use the session from
+    """
+
+    _session = None  # Not checked out until '__enter__'.
+
+    def __init__(self, database):
+        if not isinstance(database, Database):
+            raise TypeError(
+                "{class_name} must receive an instance of {expected_class_name}. Received: {actual_class_name}".format(
+                    class_name=self.__class__.__name__,
+                    expected_class_name=Database.__name__,
+                    actual_class_name=database.__class__.__name__,
+                )
+            )
+
+        self._database = database
+
+    def __enter__(self):
+        self._session = self._database._session_manager.get_session_for_read_write()
+        return self._session
+
+    def __exit__(self, *ignored):
+        self._database._session_manager.put_session(self._session)


SessionCheckout has been moved from pool.py to database.py, and updated so that the session is checked out from the database rather than from the pool.

Furthermore, SessionCheckout's constructor previous accepted arbitrary keyword arguments, which were passed to the pool's get method. While AbstractSessionPool.get does not describe any keyword arguments, both FixedSizePool.get and PingingPool.get accept an optional timeout keyword argument, which can be used to override the default timeout.

Because DatabaseSessionManager now manages which session is returned, and it not guaranteed to return a session from the pool, we are suggesting removing the ability to provide arbitrary keyword arguments to the pool via SessionCheckout.__init__.

No keyword arguments are supported by the default pool implementation (BurstyPool), and users are already able to specify their own default timeout for FixedSizePool and PingingPool. Moreover, it is not clear to us where this functionality is even used: remaining uses of SessionCheckout do not support these keyword arguments.

Please let us know:

Is it acceptable to remove this functionality?

Would you also like to remove support for the timeout arguments to FixedSizePool and PingingPool, for consistency?

currantw · 2025-03-26T23:26:48Z

+        if not isinstance(database, Database):
+            raise TypeError(
+                "{class_name} must receive an instance of {expected_class_name}. Received: {actual_class_name}".format(
+                    class_name=self.__class__.__name__,
+                    expected_class_name=Database.__name__,
+                    actual_class_name=database.__class__.__name__,
+                )
+            )


For better error handling, added a type check to all *Checkout classes, along with corresponding unit tests.

currantw · 2025-03-26T23:27:53Z

+        database = Database(
+            database_id=self.DATABASE_ID, instance=_Instance(self.INSTANCE_NAME)
+        )


Because of the added type check in the constructors for *Checkout classes (see comment here), these tests have been updated to use a Database instance instead of _Database.

currantw · 2025-03-26T23:38:27Z

-    def test_context_mgr_session_not_found_error(self):
-        from google.cloud.exceptions import NotFound
-
-        database = _Database(self.DATABASE_NAME)
-        session = _Session(database, name="session-1")
-        session.exists = mock.MagicMock(return_value=False)
-        pool = database._pool = _Pool()
-        new_session = _Session(database, name="session-2")
-        new_session.create = mock.MagicMock(return_value=[])
-        pool._new_session = mock.MagicMock(return_value=new_session)
+        session_manager.put_session.assert_called_once_with(session)

-        pool.put(session)
-        checkout = self._make_one(database)
-
-        self.assertEqual(pool._session, session)
-        with self.assertRaises(NotFound):
-            with checkout as _:
-                raise NotFound("Session not found")
-        # Assert that session-1 was removed from pool and new session was added.
-        self.assertEqual(pool._session, new_session)
-
-    def test_context_mgr_table_not_found_error(self):
-        from google.cloud.exceptions import NotFound
-
-        database = _Database(self.DATABASE_NAME)
-        session = _Session(database, name="session-1")
-        session.exists = mock.MagicMock(return_value=True)
-        pool = database._pool = _Pool()
-        pool._new_session = mock.MagicMock(return_value=[])
-
-        pool.put(session)
-        checkout = self._make_one(database)
-
-        self.assertEqual(pool._session, session)
-        with self.assertRaises(NotFound):
-            with checkout as _:
-                raise NotFound("Table not found")
-        # Assert that session-1 was not removed from pool.
-        self.assertEqual(pool._session, session)
-        pool._new_session.assert_not_called()
-
-    def test_context_mgr_unknown_error(self):
-        database = _Database(self.DATABASE_NAME)
-        session = _Session(database)
-        pool = database._pool = _Pool()
-        pool._new_session = mock.MagicMock(return_value=[])
-        pool.put(session)
-        checkout = self._make_one(database)
-
-        class Testing(Exception):
-            pass
-
-        self.assertEqual(pool._session, session)
-        with self.assertRaises(Testing):
-            with checkout as _:
-                raise Testing("Unknown error.")
-        # Assert that session-1 was not removed from pool.
-        self.assertEqual(pool._session, session)
-        pool._new_session.assert_not_called()


These tests are no longer required because the context managers are no longer responsible for re-creating the session if it fails with a NotFound error. See this comment for more details: https://github.com/googleapis/python-spanner/pull/1329/files#r2015092908

currantw · 2025-03-26T23:39:02Z

-    def test_context_mgr_session_not_found_error(self):
-        from google.cloud.exceptions import NotFound

-        database = _Database(self.DATABASE_NAME)
-        session = _Session(database, name="session-1")
-        session.exists = mock.MagicMock(return_value=False)
-        pool = database._pool = _Pool()
-        new_session = _Session(database, name="session-2")
-        new_session.create = mock.MagicMock(return_value=[])
-        pool._new_session = mock.MagicMock(return_value=new_session)
+class TestSessionCheckout(_BaseTest):
+    def _get_target_class(self):
+        from google.cloud.spanner_v1.database import SessionCheckout
+
+        return SessionCheckout
+
+    def test_ctor(self):
+        from google.cloud.spanner_v1.database import Database
+
+        database = Database(
+            database_id=self.DATABASE_ID, instance=_Instance(self.INSTANCE_NAME)
+        )

-        pool.put(session)
        checkout = self._make_one(database)
+        self.assertIs(checkout._database, database)
+        self.assertIsNone(checkout._session)

-        self.assertEqual(pool._session, session)
-        with self.assertRaises(NotFound):
-            with checkout as _:
-                raise NotFound("Session not found")
-        # Assert that session-1 was removed from pool and new session was added.
-        self.assertEqual(pool._session, new_session)
+    def test_context_manager_success(self):
+        from google.cloud.spanner_v1.database import Database

-    def test_context_mgr_table_not_found_error(self):
-        from google.cloud.exceptions import NotFound
+        database = Database(
+            database_id=self.DATABASE_ID, instance=_Instance(self.INSTANCE_NAME)
+        )

-        database = _Database(self.DATABASE_NAME)
-        session = _Session(database, name="session-1")
-        session.exists = mock.MagicMock(return_value=True)
-        pool = database._pool = _Pool()
-        pool._new_session = mock.MagicMock(return_value=[])
+        session = _Session(database)
+        session_manager = database._session_manager
+        session_manager.get_session_for_read_write = mock.Mock(return_value=session)
+        session_manager.put_session = mock.Mock(return_value=None)

-        pool.put(session)
        checkout = self._make_one(database)

-        self.assertEqual(pool._session, session)
-        with self.assertRaises(NotFound):
-            with checkout as _:
-                raise NotFound("Table not found")
-        # Assert that session-1 was not removed from pool.
-        self.assertEqual(pool._session, session)
-        pool._new_session.assert_not_called()
-
-    def test_context_mgr_unknown_error(self):
-        database = _Database(self.DATABASE_NAME)
+        with checkout as borrowed:
+            session_manager.get_session_for_read_write.assert_called_once()
+            self.assertIs(borrowed, session)
+
+        session_manager.put_session.assert_called_once_with(session)
+
+    def test_context_manager_failure(self):
+        from google.cloud.spanner_v1.database import Database
+
+        database = Database(
+            database_id=self.DATABASE_ID, instance=_Instance(self.INSTANCE_NAME)
+        )
+
        session = _Session(database)
-        pool = database._pool = _Pool()
-        pool._new_session = mock.MagicMock(return_value=[])
-        pool.put(session)


These tests are no longer required because the context managers are no longer responsible for re-creating the session if it fails with a NotFound error. See this comment for more details: https://github.com/googleapis/python-spanner/pull/1329/files#r2015092908

currantw · 2025-03-26T23:56:34Z

-    def test_ctor_w_kwargs(self):
-        pool = _Pool()
-        checkout = self._make_one(pool, foo="bar", database_role="dummy-role")
-        self.assertIs(checkout._pool, pool)
-        self.assertIsNone(checkout._session)
-        self.assertEqual(
-            checkout._kwargs, {"foo": "bar", "database_role": "dummy-role"}
-        )
-
-    def test_context_manager_wo_kwargs(self):
-        session = object()
-        pool = _Pool(session)
-        checkout = self._make_one(pool)
-
-        self.assertEqual(len(pool._items), 1)
-        self.assertIs(pool._items[0], session)
-
-        with checkout as borrowed:
-            self.assertIs(borrowed, session)
-            self.assertEqual(len(pool._items), 0)
-
-        self.assertEqual(len(pool._items), 1)
-        self.assertIs(pool._items[0], session)
-        self.assertEqual(pool._got, {})
-
-    def test_context_manager_w_kwargs(self):
-        session = object()
-        pool = _Pool(session)
-        checkout = self._make_one(pool, foo="bar")
-
-        self.assertEqual(len(pool._items), 1)
-        self.assertIs(pool._items[0], session)
-
-        with checkout as borrowed:
-            self.assertIs(borrowed, session)
-            self.assertEqual(len(pool._items), 0)
-
-        self.assertEqual(len(pool._items), 1)
-        self.assertIs(pool._items[0], session)
-        self.assertEqual(pool._got, {"foo": "bar"})


These tests have mostly been moved to test_database.py (since SessionCheckout has moved to database.py). The "w_kwargs" tests have been removed because they are no longer supported: see this comment here: https://github.com/googleapis/python-spanner/pull/1329/files#r2015125513

currantw · 2025-03-26T23:57:03Z

-def _make_transaction(*args, **kw):
-    from google.cloud.spanner_v1.transaction import Transaction
-
-    txn = mock.create_autospec(Transaction)(*args, **kw)
-    txn.committed = None
-    txn.rolled_back = False
-    return txn


This method was unused so I just removed it.

currantw · 2025-03-26T23:57:18Z

-    def transaction(self):
-        txn = self._transaction = _make_transaction(self)
-        return txn
-
-    @property
-    def session_id(self):
-        return self._session_id


These were unused, so I just removed them.

…es and refactoring) - Adds SessionOptions and DatabaseSessionManager class support multiplexed session in future work. - Update Client, Instance, Database, Session and associated Checkout classes accordingly. - Add unit tests for new classes and update those for existing classes. Signed-off-by: currantw <taylor.curran@improving.com>

currantw requested review from a team March 26, 2025 20:31

blunderbuss-gcf bot assigned harshachinta Mar 26, 2025

product-auto-label bot added size: xl Pull request size is extra large. api: spanner Issues related to the googleapis/python-spanner API. labels Mar 26, 2025

currantw force-pushed the multiplexed_session_1 branch 2 times, most recently from 219c04e to 99b9a75 Compare March 26, 2025 22:35

currantw commented Mar 26, 2025

View reviewed changes

currantw force-pushed the multiplexed_session_1 branch from 99b9a75 to 89b4d01 Compare March 26, 2025 23:17

currantw commented Mar 26, 2025

View reviewed changes

currantw force-pushed the multiplexed_session_1 branch 2 times, most recently from 0730ac1 to 97d8c41 Compare March 28, 2025 16:57

rahul2393 changed the title ~~Add support for multiplexed sessions - part 1 (classes)~~ chore: add support for multiplexed sessions - part 1 (classes) Mar 28, 2025

currantw force-pushed the multiplexed_session_1 branch 2 times, most recently from f4238fc to 58d588a Compare March 31, 2025 16:21

currantw changed the title ~~chore: add support for multiplexed sessions - part 1 (classes)~~ feat: Add support for multiplexed sessions - part 1 (supporting classes and refactoring) Mar 31, 2025

rahul2393 added the kokoro:force-run Add this label to force Kokoro to re-run the tests. label Mar 31, 2025

yoshi-kokoro removed the kokoro:force-run Add this label to force Kokoro to re-run the tests. label Mar 31, 2025

currantw force-pushed the multiplexed_session_1 branch from 58d588a to 3bb0647 Compare April 1, 2025 19:48

currantw force-pushed the multiplexed_session_1 branch from 3bb0647 to a8093f4 Compare April 1, 2025 22:58

rahul2393 requested a review from kokoro-team April 3, 2025 04:10

currantw changed the base branch from main to multiplexed-sessions April 3, 2025 04:18

rahul2393 merged commit 8619929 into googleapis:multiplexed-sessions Apr 4, 2025
5 of 6 checks passed

rahul2393 mentioned this pull request Apr 4, 2025

feat: Add support for multiplexed sessions #1332

Closed

4 tasks

currantw deleted the multiplexed_session_1 branch April 8, 2025 23:54

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: Add support for multiplexed sessions - part 1 (supporting classes and refactoring)#1329

feat: Add support for multiplexed sessions - part 1 (supporting classes and refactoring)#1329
rahul2393 merged 1 commit intogoogleapis:multiplexed-sessionsfrom
currantw:multiplexed_session_1

currantw commented Mar 26, 2025 •

edited

Loading

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

currantw Mar 26, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

currantw commented Mar 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

currantw commented Mar 26, 2025 •

edited

Loading