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

@@ -10,15 +10,12 @@
 import warnings
 
 from google.cloud import spanner_v1
+from google.cloud.spanner_v1.pool import BurstyPool
 
 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"])
 
@@ -34,14 +31,87 @@ 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):
+    def __init__(self, instance, database, pool=None):
+        self._pool = pool or BurstyPool()
+        self._pool.bind(database)
+
         self.instance = instance
         self.database = database
-        self.is_closed = False
 
         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
+
+    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.
+        """
+        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()
@@ -142,18 +212,34 @@ 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.
         """
+        if (
+            self._transaction
+            and not self._transaction.committed
+            and not self._transaction.rolled_back
+        ):
+            self._transaction.rollback()
+
         self.__dbhandle = None
         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,15 @@ 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)
+                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_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(