feat: support transactions management

google/cloud/spanner_dbapi/__init__.py

@@ -47,6 +47,10 @@
 # threadsafety level.
 threadsafety = 1
 
+# Cloud Spanner sessions pool, used by
+# default for the whole DB API package
+default_pool = spanner_v1.BurstyPool()
+
 
 def connect(
     instance_id, database_id, project=None, credentials=None, user_agent=None
@@ -93,7 +97,7 @@ def connect(
     if not database.exists():
         raise ValueError("database '%s' does not exist." % database_id)
 
-    return Connection(instance, database)
+    return Connection(instance, database, default_pool)
 
 
 __all__ = [

google/cloud/spanner_dbapi/connection.py

@@ -31,17 +31,21 @@ class Connection:
 
     :type database: :class:`~google.cloud.spanner_v1.database.Database`
     :param database: Cloud Spanner database to connect to.
+
+    :type pool: :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`
+    :param pool: (Optional) Cloud Spanner sessions pool.
     """
 
-    def __init__(self, instance, database):
-        self._pool = BurstyPool()
+    def __init__(self, instance, database, pool=None):
+        self._pool = pool or BurstyPool()
         self._pool.bind(database)
 
         self.instance = instance
         self.database = database
 
         self._ddl_statements = []
-        self.transactions = []
+        self._transaction = None
+        self._session = None
 
         self.is_closed = False
         self._autocommit = False
@@ -67,13 +71,47 @@ def autocommit(self, value):
 
         self._autocommit = value
 
-    def session_checkout(self):
-        """Get a Cloud Spanner session.
+    def _session_checkout(self):
+        """Get a Cloud Spanner session from a pool.
+
+        If there is already a session associated with
+        this connection, it'll be used otherwise.
 
         :rtype: :class:`google.cloud.spanner_v1.session.Session`
         :returns: Cloud Spanner session object ready to use.
         """
-        return self._pool.get()
+        if not self._session:
+            self._session = self._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._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.
+
+        :rtype: :class:`google.cloud.spanner_v1.transaction.Transaction`
+        :returns: 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()
@@ -174,11 +212,11 @@ def get_table_column_schema(self, table_name):
     def close(self):
         """Close this connection.
 
-        The connection will be unusable from this point forward. Rollback
-        will be performed on all the pending transactions.
+        The connection will be unusable from this point forward. If the
+        connection has an active transaction, it will be rolled back.
         """
-        for transaction in self.transactions:
-            transaction.rollback()
+        if self._transaction and not self._transaction.committed:
+            self._transaction.rollback()
 
         self.__dbhandle = None
         self.is_closed = True
@@ -187,21 +225,17 @@ def commit(self):
         """Commit all the pending transactions."""
         if self.autocommit:
             warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
-        else:
-            for transaction in self.transactions:
-                transaction.commit()
-
-            self.transactions = []
+        elif self._transaction:
+            self._transaction.commit()
+            self._release_session()
 
     def rollback(self):
         """Rollback all the pending transactions."""
         if self.autocommit:
             warnings.warn(AUTOCOMMIT_MODE_WARNING, UserWarning, stacklevel=2)
-        else:
-            for transaction in self.transactions:
-                transaction.rollback()
-
-            self.transactions = []
+        elif self._transaction:
+            self._transaction.rollback()
+            self._release_session()
 
     def __enter__(self):
         return self

google/cloud/spanner_dbapi/cursor.py

@@ -68,7 +68,6 @@ def __init__(self, connection):
         self._connection = connection
         self._is_closed = False
 
-        self.transaction = None
         # the number of rows to fetch at a time with fetchmany()
         self.arraysize = 1
 
@@ -90,16 +89,9 @@ def execute(self, sql, args=None):
         self._res = None
 
         if not self._connection.autocommit:
-            if (
-                not self.transaction
-                or self.transaction.committed
-                or self.transaction.rolled_back
-            ):
-                self.transaction = self._connection.session_checkout().transaction()
-                self.transaction.begin()
-                self._connection.transactions.append(self.transaction)
-
-            self._res = self.transaction.execute_sql(sql)
+            transaction = self._connection.transaction_checkout()
+
+            self._res = transaction.execute_sql(sql)
             self._itr = PeekIterator(self._res)
             return
 

tests/system/test_system.py

@@ -169,6 +169,9 @@ def test_commit(self):
 
         self.assertEqual(got_rows, [want_row])
 
+        cursor.close()
+        conn.close()
+
     def test_rollback(self):
         """Test rollbacking a transaction with several statements."""
         want_row = (2, "first-name", "last-name", "test.email@domen.ru")
@@ -208,6 +211,9 @@ def test_rollback(self):
 
         self.assertEqual(got_rows, [want_row])
 
+        cursor.close()
+        conn.close()
+
     def test_autocommit_mode_change(self):
         """Test auto committing a transaction on `autocommit` mode change."""
         want_row = (
@@ -238,10 +244,12 @@ def test_autocommit_mode_change(self):
         # read the resulting data from the database
         cursor.execute("SELECT * FROM contacts")
         got_rows = cursor.fetchall()
-        conn.commit()
 
         self.assertEqual(got_rows, [want_row])
 
+        cursor.close()
+        conn.close()
+
     def test_rollback_on_connection_closing(self):
         """
         When closing a connection all the pending transactions
@@ -280,6 +288,9 @@ def test_rollback_on_connection_closing(self):
 
         self.assertEqual(got_rows, [want_row])
 
+        cursor.close()
+        conn.close()
+
 
 def clear_table(transaction):
     """Clear the test table."""