feat: Multiplexed sessions - Update session to remove class attribute…

…s, add TODOs, and make `Session._transaction` default to None. Plus add some `Optional` typing hints.

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

google/cloud/spanner_v1/batch.py

@@ -15,7 +15,7 @@
 """Context manager for Cloud Spanner batched writes."""
 import functools
 from datetime import datetime
-from typing import List
+from typing import List, Optional
 
 from google.cloud.spanner_v1 import CommitRequest, CommitResponse
 from google.cloud.spanner_v1 import Mutation
@@ -53,11 +53,11 @@ def __init__(self, session):
         super(_BatchBase, self).__init__(session)
 
         self._mutations: List[Mutation] = []
-        self.transaction_tag: str = None
+        self.transaction_tag: Optional[str] = None
 
-        self.committed: datetime = None
+        self.committed: Optional[datetime] = None
         """Timestamp at which the batch was successfully committed."""
-        self.commit_stats: CommitResponse.CommitStats = None
+        self.commit_stats: Optional[CommitResponse.CommitStats] = None
 
     # TODO multiplexed - cleanup
     def _check_state(self):

google/cloud/spanner_v1/session.py

@@ -17,6 +17,7 @@
 from functools import total_ordering
 import time
 from datetime import datetime
+from typing import MutableMapping, Optional
 
 from google.api_core.exceptions import Aborted
 from google.api_core.exceptions import GoogleAPICallError
@@ -69,17 +70,20 @@ class Session(object):
     :param is_multiplexed: (Optional) whether this session is a multiplexed session.
     """
 
-    _session_id = None
-    _transaction = None
-
     def __init__(self, database, labels=None, database_role=None, is_multiplexed=False):
         self._database = database
+        self._session_id: Optional[str] = None
+
+        # TODO multiplexed - remove
+        self._transaction: Optional[Transaction] = None
+
         if labels is None:
             labels = {}
-        self._labels = labels
-        self._database_role = database_role
-        self._is_multiplexed = is_multiplexed
-        self._last_use_time = datetime.utcnow()
+
+        self._labels: MutableMapping[str, str] = labels
+        self._database_role: Optional[str] = database_role
+        self._is_multiplexed: bool = is_multiplexed
+        self._last_use_time: datetime = datetime.utcnow()
 
     def __lt__(self, other):
         return self._session_id < other._session_id
@@ -100,7 +104,7 @@ def is_multiplexed(self):
 
     @property
     def last_use_time(self):
-        """ "Approximate last use time of this session
+        """Approximate last use time of this session
 
         :rtype: datetime
         :returns: the approximate last use time of this session"""
@@ -157,27 +161,28 @@ def create(self):
 
         if self._session_id is not None:
             raise ValueError("Session ID already set by back-end")
-        api = self._database.spanner_api
-        metadata = _metadata_with_prefix(self._database.name)
-        if self._database._route_to_leader_enabled:
+
+        database = self._database
+        api = database.spanner_api
+
+        metadata = _metadata_with_prefix(database.name)
+        if database._route_to_leader_enabled:
             metadata.append(
-                _metadata_with_leader_aware_routing(
-                    self._database._route_to_leader_enabled
-                )
+                _metadata_with_leader_aware_routing(database._route_to_leader_enabled)
             )
 
-        request = CreateSessionRequest(database=self._database.name)
-        if self._database.database_role is not None:
-            request.session.creator_role = self._database.database_role
+        create_session_request = CreateSessionRequest(database=database.name)
+        if database.database_role is not None:
+            create_session_request.session.creator_role = database.database_role
 
         if self._labels:
-            request.session.labels = self._labels
+            create_session_request.session.labels = self._labels
 
         # Set the multiplexed field for multiplexed sessions
         if self._is_multiplexed:
-            request.session.multiplexed = True
+            create_session_request.session.multiplexed = True
 
-        observability_options = getattr(self._database, "observability_options", None)
+        observability_options = getattr(database, "observability_options", None)
         span_name = (
             "CloudSpanner.CreateMultiplexedSession"
             if self._is_multiplexed
@@ -191,9 +196,9 @@ def create(self):
             metadata=metadata,
         ) as span, MetricsCapture():
             session_pb = api.create_session(
-                request=request,
-                metadata=self._database.metadata_with_request_id(
-                    self._database._next_nth_request,
+                request=create_session_request,
+                metadata=database.metadata_with_request_id(
+                    database._next_nth_request,
                     1,
                     metadata,
                     span,
@@ -462,6 +467,7 @@ def batch(self):
 
         return Batch(self)
 
+    # TODO multiplexed - remove
     def transaction(self):
         """Create a transaction to perform a set of reads with shared staleness.
 
@@ -474,7 +480,7 @@ def transaction(self):
 
         if self._transaction is not None:
             self._transaction.rolled_back = True
-            del self._transaction
+            self._transaction = None
 
         txn = self._transaction = Transaction(self)
         return txn
@@ -531,6 +537,7 @@ def run_in_transaction(self, func, *args, **kw):
             observability_options=observability_options,
         ) as span, MetricsCapture():
             while True:
+                # TODO multiplexed - remove
                 if self._transaction is None:
                     txn = self.transaction()
                     txn.transaction_tag = transaction_tag
@@ -552,8 +559,11 @@ def run_in_transaction(self, func, *args, **kw):
 
                     return_value = func(txn, *args, **kw)
 
+                # TODO multiplexed: store previous transaction ID.
                 except Aborted as exc:
-                    del self._transaction
+                    # TODO multiplexed - remove
+                    self._transaction = None
+
                     if span:
                         delay_seconds = _get_retry_delay(
                             exc.errors[0],
@@ -573,7 +583,9 @@ def run_in_transaction(self, func, *args, **kw):
                     )
                     continue
                 except GoogleAPICallError:
-                    del self._transaction
+                    # TODO multiplexed - remove
+                    self._transaction = None
+
                     add_span_event(
                         span,
                         "User operation failed due to GoogleAPICallError, not retrying",
@@ -596,7 +608,9 @@ def run_in_transaction(self, func, *args, **kw):
                         max_commit_delay=max_commit_delay,
                     )
                 except Aborted as exc:
-                    del self._transaction
+                    # TODO multiplexed - remove
+                    self._transaction = None
+
                     if span:
                         delay_seconds = _get_retry_delay(
                             exc.errors[0],
@@ -615,7 +629,9 @@ def run_in_transaction(self, func, *args, **kw):
                         exc, deadline, attempts, default_retry_delay=default_retry_delay
                     )
                 except GoogleAPICallError:
-                    del self._transaction
+                    # TODO multiplexed - remove
+                    self._transaction = None
+
                     add_span_event(
                         span,
                         "Transaction.commit failed due to GoogleAPICallError, not retrying",

google/cloud/spanner_v1/snapshot.py

@@ -16,7 +16,7 @@
 
 import functools
 import threading
-from typing import List, Union
+from typing import List, Union, Optional
 
 from google.protobuf.struct_pb2 import Struct
 from google.cloud.spanner_v1 import (
@@ -225,12 +225,12 @@ def __init__(self, session):
         self._read_request_count: int = 0
 
         # Identifier for the transaction.
-        self._transaction_id: bytes = None
+        self._transaction_id: Optional[bytes] = None
 
         # Precommit tokens are returned for transactions with
         # multiplexed sessions. The precommit token with the
         # highest sequence number is included in the  commit request.
-        self._precommit_token: MultiplexedSessionPrecommitToken = None
+        self._precommit_token: Optional[MultiplexedSessionPrecommitToken] = None
 
         # Operations within a transaction can be performed using multiple
         # threads, so we need to use a lock when updating the transaction.

google/cloud/spanner_v1/transaction.py

@@ -259,7 +259,7 @@ def wrapped_method(*args, **kwargs):
                     attempt.increment()
                     rollback_method = functools.partial(
                         api.rollback,
-                        session=self._session.name,
+                        session=session.name,
                         transaction_id=self._transaction_id,
                         metadata=database.metadata_with_request_id(
                             nth_request,
@@ -278,7 +278,7 @@ def wrapped_method(*args, **kwargs):
         self.rolled_back = True
 
         # TODO multiplexed - remove
-        del self._session._transaction
+        self._session._transaction = None
 
     def commit(
         self, return_commit_stats=False, request_options=None, max_commit_delay=None
@@ -396,7 +396,7 @@ def before_next_retry(nth_retry, delay_in_seconds):
             self.commit_stats = response_pb.commit_stats
 
         # TODO multiplexed - remove
-        del self._session._transaction
+        self._session._transaction = None
 
         return self.committed