feat: Multiplexed sessions - Move Batch and Transaction attributes from class attributes to instance attributes.

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

Author: currantw

google/cloud/spanner_v1/batch.py

@@ -14,8 +14,10 @@
 
 """Context manager for Cloud Spanner batched writes."""
 import functools
+from datetime import datetime
+from typing import List
 
-from google.cloud.spanner_v1 import CommitRequest
+from google.cloud.spanner_v1 import CommitRequest, CommitResponse
 from google.cloud.spanner_v1 import Mutation
 from google.cloud.spanner_v1 import TransactionOptions
 from google.cloud.spanner_v1 import BatchWriteRequest
@@ -47,13 +49,17 @@ class _BatchBase(_SessionWrapper):
     :param session: the session used to perform the commit
     """
 
-    transaction_tag = None
-    _read_only = False
-
     def __init__(self, session):
         super(_BatchBase, self).__init__(session)
-        self._mutations = []
 
+        self._mutations: List[Mutation] = []
+        self.transaction_tag: str = None
+
+        self.committed: datetime = None
+        """Timestamp at which the batch was successfully committed."""
+        self.commit_stats: CommitResponse.CommitStats = None
+
+    # TODO multiplexed - cleanup
     def _check_state(self):
         """Helper for :meth:`commit` et al.
 
@@ -148,10 +154,7 @@ def delete(self, table, keyset):
 class Batch(_BatchBase):
     """Accumulate mutations for transmission during :meth:`commit`."""
 
-    committed = None
-    commit_stats = None
-    """Timestamp at which the batch was successfully committed."""
-
+    # TODO multiplexed - cleanup
     def _check_state(self):
         """Helper for :meth:`commit` et al.
 
@@ -163,6 +166,7 @@ def _check_state(self):
         if self.committed is not None:
             raise ValueError("Batch already committed")
 
+    # TODO multiplexed - cleanup kwargs
     def commit(
         self,
         return_commit_stats=False,
@@ -205,7 +209,10 @@ def commit(
         :rtype: datetime
         :returns: timestamp of the committed changes.
         """
+
+        # TODO multiplexed - cleanup
         self._check_state()
+
         database = self._session._database
         api = database.spanner_api
         metadata = _metadata_with_prefix(database.name)
@@ -282,6 +289,8 @@ def wrapped_method(*args, **kwargs):
 
     def __enter__(self):
         """Begin ``with`` block."""
+
+        # TODO multiplexed - cleanup
         self._check_state()
 
         return self
@@ -317,19 +326,18 @@ class MutationGroups(_SessionWrapper):
     :param session: the session used to perform the commit
     """
 
-    committed = None
-
     def __init__(self, session):
         super(MutationGroups, self).__init__(session)
-        self._mutation_groups = []
+        self._mutation_groups: List[MutationGroup] = []
+        self.committed: bool = False
 
     def _check_state(self):
         """Checks if the object's state is valid for making API requests.
 
         :raises: :exc:`ValueError` if the object's state is invalid for making
                  API requests.
         """
-        if self.committed is not None:
+        if self.committed:
             raise ValueError("MutationGroups already committed")
 
     def group(self):
@@ -358,10 +366,14 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
         :rtype: :class:`Iterable[google.cloud.spanner_v1.types.BatchWriteResponse]`
         :returns: a sequence of responses for each batch.
         """
+
+        # TODO multiplexed - cleanup
         self._check_state()
 
-        database = self._session._database
+        session = self._session
+        database = session._database
         api = database.spanner_api
+
         metadata = _metadata_with_prefix(database.name)
         if database._route_to_leader_enabled:
             metadata.append(
@@ -374,7 +386,7 @@ def batch_write(self, request_options=None, exclude_txn_from_change_streams=Fals
             request_options = RequestOptions(request_options)
 
         request = BatchWriteRequest(
-            session=self._session.name,
+            session=session.name,
             mutation_groups=self._mutation_groups,
             request_options=request_options,
             exclude_txn_from_change_streams=exclude_txn_from_change_streams,
@@ -409,6 +421,7 @@ def wrapped_method(*args, **kwargs):
                     InternalServerError: _check_rst_stream_error,
                 },
             )
+
         self.committed = True
         return response
 

google/cloud/spanner_v1/transaction.py

@@ -57,12 +57,10 @@ class Transaction(_SnapshotBase, _BatchBase):
     :raises ValueError: if session has an existing transaction
     """
 
-    committed = None
-    """Timestamp at which the transaction was successfully committed."""
-    rolled_back = False
-    commit_stats = None
-    exclude_txn_from_change_streams = False
-    isolation_level = TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+    exclude_txn_from_change_streams: bool = False
+    isolation_level: TransactionOptions.IsolationLevel = (
+        TransactionOptions.IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED
+    )
 
     # Override defaults from _SnapshotBase.
     _multi_use: bool = True
@@ -74,6 +72,7 @@ def __init__(self, session):
             raise ValueError("Session has existing transaction.")
 
         super(Transaction, self).__init__(session)
+        self.rolled_back: bool = False
 
     # TODO multiplexed - remove
     def _check_state(self):
@@ -97,6 +96,8 @@ def _make_txn_selector(self):
             :class:`~.transaction_pb2.TransactionSelector`
         :returns: a selector configured for read-write transaction semantics.
         """
+
+        # TODO multiplexed - remove
         self._check_state()
 
         if self._transaction_id is None: