feat: support transactions management (#535)

Add transaction management, including utils for handling spanner sessions and connections.

Co-authored-by: MF2199 <38331387+mf2199@users.noreply.github.com>
Co-authored-by: Chris Kleinknecht <libc@google.com>

google/cloud/spanner_dbapi/__init__.py

@@ -49,7 +49,12 @@
 
 
 def connect(
-    instance_id, database_id, project=None, credentials=None, user_agent=None
+    instance_id,
+    database_id,
+    project=None,
+    credentials=None,
+    pool=None,
+    user_agent=None,
 ):
     """
     Create a connection to Cloud Spanner database.
@@ -71,6 +76,13 @@ def connect(
                         If none are specified, the client will attempt to ascertain
                         the credentials from the environment.
 
+    :type pool: Concrete subclass of
+                :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`.
+    :param pool: (Optional). Session pool to be used by database.
+
+    :type user_agent: :class:`str`
+    :param user_agent: (Optional) User agent to be used with this connection requests.
+
     :rtype: :class:`google.cloud.spanner_dbapi.connection.Connection`
     :returns: Connection object associated with the given Cloud Spanner resource.
 
@@ -87,9 +99,7 @@ def connect(
     if not instance.exists():
         raise ValueError("instance '%s' does not exist." % instance_id)
 
-    database = instance.database(
-        database_id, pool=spanner_v1.pool.BurstyPool()
-    )
+    database = instance.database(database_id, pool=pool)
     if not database.exists():
         raise ValueError("database '%s' does not exist." % database_id)
 

google/cloud/spanner_dbapi/connection.py

@@ -14,11 +14,7 @@
 from .cursor import Cursor
 from .exceptions import InterfaceError
 
-AUTOCOMMIT_MODE_WARNING = (
-    "This method is non-operational, as Cloud Spanner"
-    "DB API always works in `autocommit` mode."
-    "See https://github.com/googleapis/python-spanner-django#transaction-management-isnt-supported"
-)
+AUTOCOMMIT_MODE_WARNING = "This method is non-operational in autocommit mode"
 
 ColumnDetails = namedtuple("column_details", ["null_ok", "spanner_type"])
 
@@ -37,11 +33,98 @@ class Connection:
     """
 
     def __init__(self, instance, database):
-        self.instance = instance
-        self.database = database
-        self.is_closed = False
+        self._instance = instance
+        self._database = database
 
         self._ddl_statements = []
+        self._transaction = None
+        self._session = None
+
+        self.is_closed = False
+        self._autocommit = False
+
+    @property
+    def autocommit(self):
+        """Autocommit mode flag for this connection.
+
+        :rtype: bool
+        :returns: Autocommit mode flag value.
+        """
+        return self._autocommit
+
+    @autocommit.setter
+    def autocommit(self, value):
+        """Change this connection autocommit mode.
+
+        :type value: bool
+        :param value: New autocommit mode state.
+        """
+        if value and not self._autocommit:
+            self.commit()
+
+        self._autocommit = value
+
+    @property
+    def database(self):
+        """Database to which this connection relates.
+
+        :rtype: :class:`~google.cloud.spanner_v1.database.Database`
+        :returns: The related database object.
+        """
+        return self._database
+
+    @property
+    def instance(self):
+        """Instance to which this connection relates.
+
+        :rtype: :class:`~google.cloud.spanner_v1.instance.Instance`
+        :returns: The related instance object.
+        """
+        return self._instance
+
+    def _session_checkout(self):
+        """Get a Cloud Spanner session from the pool.
+
+        If there is already a session associated with
+        this connection, it'll be used instead.
+
+        :rtype: :class:`google.cloud.spanner_v1.session.Session`
+        :returns: Cloud Spanner session object ready to use.
+        """
+        if not self._session:
+            self._session = self.database._pool.get()
+
+        return self._session
+
+    def _release_session(self):
+        """Release the currently used Spanner session.
+
+        The session will be returned into the sessions pool.
+        """
+        self.database._pool.put(self._session)
+        self._session = None
+
+    def transaction_checkout(self):
+        """Get a Cloud Spanner transaction.
+
+        Begin a new transaction, if there is no transaction in
+        this connection yet. Return the begun one otherwise.
+
+        The method is non operational in autocommit mode.
+
+        :rtype: :class:`google.cloud.spanner_v1.transaction.Transaction`
+        :returns: A Cloud Spanner transaction object, ready to use.
+        """
+        if not self.autocommit:
+            if (
+                not self._transaction
+                or self._transaction.committed
+                or self._transaction.rolled_back
+            ):
+                self._transaction = self._session_checkout().transaction()
+                self._transaction.begin()
+
+            return self._transaction
 
     def cursor(self):
         self._raise_if_closed()
@@ -142,18 +225,33 @@ def get_table_column_schema(self, table_name):
     def close(self):
         """Close this connection.
 
-        The connection will be unusable from this point forward.
+        The connection will be unusable from this point forward. If the
+        connection has an active transaction, it will be rolled back.
         """
-        self.__dbhandle = None
+        if (
+            self._transaction
+            and not self._transaction.committed
+            and not self._transaction.rolled_back
+        ):
+            self._transaction.rollback()
+
         self.is_closed = True
 
     def commit(self):
         """Commit all the pending transactions."""
-        warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
+        if self.autocommit:
+            warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
+        elif self._transaction:
+            self._transaction.commit()
+            self._release_session()
 
     def rollback(self):
         """Rollback all the pending transactions."""
-        warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
+        if self.autocommit:
+            warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
+        elif self._transaction:
+            self._transaction.rollback()
+            self._release_session()
 
     def __enter__(self):
         return self

google/cloud/spanner_dbapi/cursor.py

@@ -91,6 +91,7 @@ def execute(self, sql, args=None):
         # Classify whether this is a read-only SQL statement.
         try:
             classification = classify_stmt(sql)
+
             if classification == STMT_DDL:
                 self._connection.append_ddl_statement(sql)
                 return
@@ -99,6 +100,17 @@ def execute(self, sql, args=None):
             # any prior DDL statements were run.
             self._run_prior_DDL_statements()
 
+            if not self._connection.autocommit:
+                transaction = self._connection.transaction_checkout()
+
+                sql, params = sql_pyformat_args_to_spanner(sql, args)
+
+                self._res = transaction.execute_sql(
+                    sql, params, param_types=get_param_types(params)
+                )
+                self._itr = PeekIterator(self._res)
+                return
+
             if classification == STMT_NON_UPDATING:
                 self.__handle_DQL(sql, args or None)
             elif classification == STMT_INSERT:

tests/spanner_dbapi/test_connect.py

@@ -12,6 +12,7 @@
 import google.auth.credentials
 from google.api_core.gapic_v1.client_info import ClientInfo
 from google.cloud.spanner_dbapi import connect, Connection
+from google.cloud.spanner_v1.pool import FixedSizePool
 
 
 def _make_credentials():
@@ -43,7 +44,7 @@ def test_connect(self):
                     "test-database",
                     PROJECT,
                     CREDENTIALS,
-                    USER_AGENT,
+                    user_agent=USER_AGENT,
                 )
 
                 self.assertIsInstance(connection, Connection)
@@ -108,3 +109,27 @@ def test_connect_database_id(self):
                 database_mock.assert_called_once_with(DATABASE, pool=mock.ANY)
 
         self.assertIsInstance(connection, Connection)
+
+    def test_default_sessions_pool(self):
+        with mock.patch("google.cloud.spanner_v1.instance.Instance.database"):
+            with mock.patch(
+                "google.cloud.spanner_v1.instance.Instance.exists",
+                return_value=True,
+            ):
+                connection = connect("test-instance", "test-database")
+
+                self.assertIsNotNone(connection.database._pool)
+
+    def test_sessions_pool(self):
+        database_id = "test-database"
+        pool = FixedSizePool()
+
+        with mock.patch(
+            "google.cloud.spanner_v1.instance.Instance.database"
+        ) as database_mock:
+            with mock.patch(
+                "google.cloud.spanner_v1.instance.Instance.exists",
+                return_value=True,
+            ):
+                connect("test-instance", database_id, pool=pool)
+                database_mock.assert_called_once_with(database_id, pool=pool)

tests/spanner_dbapi/test_connection.py

@@ -49,8 +49,9 @@ def test_close(self):
             connection.cursor()
 
     @mock.patch("warnings.warn")
-    def test_transaction_management_warnings(self, warn_mock):
+    def test_transaction_autocommit_warnings(self, warn_mock):
         connection = self._make_connection()
+        connection.autocommit = True
 
         connection.commit()
         warn_mock.assert_called_with(
@@ -60,3 +61,19 @@ def test_transaction_management_warnings(self, warn_mock):
         warn_mock.assert_called_with(
             AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2
         )
+
+    def test_database_property(self):
+        connection = self._make_connection()
+        self.assertIsInstance(connection.database, Database)
+        self.assertEqual(connection.database, connection._database)
+
+        with self.assertRaises(AttributeError):
+            connection.database = None
+
+    def test_instance_property(self):
+        connection = self._make_connection()
+        self.assertIsInstance(connection.instance, Instance)
+        self.assertEqual(connection.instance, connection._instance)
+
+        with self.assertRaises(AttributeError):
+            connection.instance = None