feat: Multiplexed sessions - Update BatchSnapshot to use database session manager.

Signed-off-by: Taylor Curran <taylor.curran@improving.com>

Author: currantw

google/cloud/spanner_v1/database.py

@@ -1481,11 +1481,15 @@ def from_dict(cls, database, mapping):
 
         :rtype: :class:`BatchSnapshot`
         """
+
         instance = cls(database)
-        session = instance._session = database.session()
-        session._session_id = mapping["session_id"]
+
+        session = instance._session = Session(database=database)
+        instance._session_id = session._session_id = mapping["session_id"]
+
         snapshot = instance._snapshot = session.snapshot()
-        snapshot._transaction_id = mapping["transaction_id"]
+        instance._transaction_id = snapshot._transaction_id = mapping["transaction_id"]
+
         return instance
 
     def to_dict(self):
@@ -1516,28 +1520,39 @@ def _get_session(self):
            all partitions have been processed.
         """
         if self._session is None:
-            # Use sessions manager for partition operations
-            transaction_type = TransactionType.PARTITIONED
-            self._session = self._database.sessions_manager.get_session(
-                transaction_type
-            )
+            database = self._database
+
+            # If the session ID is not specified, check out a new session for
+            # partitioned transactions from  the database session manager; otherwise,
+            # the session has already been checked out, so just create a session to
+            # represent it.
+            if self._session_id is None:
+                transaction_type = TransactionType.PARTITIONED
+                session = database.sessions_manager.get_session(transaction_type)
+                self._session_id = session.session_id
+
+            else:
+                session = Session(database=database)
+                session._session_id = self._session_id
 
-            if self._session_id is not None:
-                self._session._session_id = self._session_id
+            self._session = session
 
         return self._session
 
     def _get_snapshot(self):
         """Create snapshot if needed."""
+
         if self._snapshot is None:
             self._snapshot = self._get_session().snapshot(
                 read_timestamp=self._read_timestamp,
                 exact_staleness=self._exact_staleness,
                 multi_use=True,
                 transaction_id=self._transaction_id,
             )
+
             if self._transaction_id is None:
                 self._snapshot.begin()
+
         return self._snapshot
 
     def get_batch_transaction_id(self):

tests/unit/test_database.py

@@ -35,7 +35,9 @@
     _metadata_with_request_id,
 )
 from google.cloud.spanner_v1.request_id_header import REQ_RAND_PROCESS_ID
-from google.cloud.spanner_v1.session_options import SessionOptions
+from google.cloud.spanner_v1.session import Session
+from google.cloud.spanner_v1.session_options import SessionOptions, TransactionType
+from tests._builders import build_spanner_api
 
 DML_WO_PARAM = """
 DELETE FROM citizens
@@ -1509,8 +1511,6 @@ def test_execute_partitioned_dml_w_exclude_txn_from_change_streams(self):
         )
 
     def test_session_factory_defaults(self):
-        from google.cloud.spanner_v1.session import Session
-
         client = _Client()
         instance = _Instance(self.INSTANCE_NAME, client=client)
         pool = _Pool()
@@ -1524,8 +1524,6 @@ def test_session_factory_defaults(self):
         self.assertEqual(session.labels, {})
 
     def test_session_factory_w_labels(self):
-        from google.cloud.spanner_v1.session import Session
-
         client = _Client()
         instance = _Instance(self.INSTANCE_NAME, client=client)
         pool = _Pool()
@@ -2475,8 +2473,6 @@ def _make_database(**kwargs):
 
     @staticmethod
     def _make_session(**kwargs):
-        from google.cloud.spanner_v1.session import Session
-
         return mock.create_autospec(Session, instance=True, **kwargs)
 
     @staticmethod
@@ -2533,20 +2529,22 @@ def test_ctor_w_exact_staleness(self):
     def test_from_dict(self):
         klass = self._get_target_class()
         database = self._make_database()
-        session = database.session.return_value = self._make_session()
-        snapshot = session.snapshot.return_value = self._make_snapshot()
-        api_repr = {
-            "session_id": self.SESSION_ID,
-            "transaction_id": self.TRANSACTION_ID,
-        }
+        api = database.spanner_api = build_spanner_api()
+
+        batch_txn = klass.from_dict(
+            database,
+            {
+                "session_id": self.SESSION_ID,
+                "transaction_id": self.TRANSACTION_ID,
+            },
+        )
 
-        batch_txn = klass.from_dict(database, api_repr)
         self.assertIs(batch_txn._database, database)
-        self.assertIs(batch_txn._session, session)
-        self.assertEqual(session._session_id, self.SESSION_ID)
-        self.assertEqual(snapshot._transaction_id, self.TRANSACTION_ID)
-        snapshot.begin.assert_not_called()
-        self.assertIs(batch_txn._snapshot, snapshot)
+        self.assertEqual(batch_txn._session._session_id, self.SESSION_ID)
+        self.assertEqual(batch_txn._snapshot._transaction_id, self.TRANSACTION_ID)
+
+        api.create_session.assert_not_called()
+        api.begin_transaction.assert_not_called()
 
     def test_to_dict(self):
         database = self._make_database()
@@ -2574,8 +2572,6 @@ def test__get_session_new(self):
         batch_txn = self._make_one(database)
         self.assertIs(batch_txn._get_session(), session)
         # Verify that sessions_manager.get_session was called with PARTITIONED transaction type
-        from google.cloud.spanner_v1.session_options import TransactionType
-
         database.sessions_manager.get_session.assert_called_once_with(
             TransactionType.PARTITIONED
         )