feat: Add support for multiplexed sessions - part 1 (supporting classes and refactoring) #1329
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
219c04e to
99b9a75
Compare
currantw
commented
Mar 26, 2025
Comment on lines
-1271
to
-1276
| 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() |
There was a problem hiding this comment.
This check has been removed from MutationGroupsCheckout and SnapshotCheckout because:
existsonly 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
BatchCheckoutorSessionCheckout) - 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.
existschecks 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.
99b9a75 to
89b4d01
Compare
currantw
commented
Mar 26, 2025
Comment on lines
+1162
to
+1188
| 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) |
There was a problem hiding this comment.
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
timeoutarguments toFixedSizePoolandPingingPool, for consistency?
currantw
commented
Mar 26, 2025
Comment on lines
+1225
to
+1232
| 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__, | ||
| ) | ||
| ) |
There was a problem hiding this comment.
For better error handling, added a type check to all *Checkout classes, along with corresponding unit tests.
currantw
commented
Mar 26, 2025
Comment on lines
+1804
to
+1802
| database = Database( | ||
| database_id=self.DATABASE_ID, instance=_Instance(self.INSTANCE_NAME) | ||
| ) |
There was a problem hiding this comment.
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
commented
Mar 26, 2025
tests/unit/test_database.py
Outdated
Comment on lines
2046
to
2104
| 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() |
There was a problem hiding this comment.
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
commented
Mar 26, 2025
tests/unit/test_database.py
Outdated
Comment on lines
3044
to
3090
| 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) |
There was a problem hiding this comment.
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
commented
Mar 26, 2025
Comment on lines
-1092
to
-1131
| 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"}) |
There was a problem hiding this comment.
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
commented
Mar 26, 2025
Comment on lines
-1134
to
-1140
| 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 |
There was a problem hiding this comment.
This method was unused so I just removed it.
currantw
commented
Mar 26, 2025
Comment on lines
-1186
to
-1192
| def transaction(self): | ||
| txn = self._transaction = _make_transaction(self) | ||
| return txn | ||
|
|
||
| @property | ||
| def session_id(self): | ||
| return self._session_id |
There was a problem hiding this comment.
These were unused, so I just removed them.
0730ac1 to
97d8c41
Compare
f4238fc to
58d588a
Compare
58d588a to
3bb0647
Compare
…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>
3bb0647 to
a8093f4
Compare
8619929
into
googleapis:multiplexed-sessions
5 of 6 checks passed
4 tasks
lszinv
pushed a commit
to currantw/google-cloud-spanner-python
that referenced
this pull request
Apr 14, 2025
…es and refactoring) (googleapis#1329) - 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>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Description
SessionOptionsandDatabaseSessionManagerclass support multiplexed session in future work.Client,Instance,Database,Sessionand associatedCheckoutclasses accordingly.