feat: support transactions management

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"])
 
@@ -37,11 +34,46 @@ class Connection:
     """
 
     def __init__(self, instance, database):
+        self._pool = BurstyPool()
+        self._pool.bind(database)
+
         self.instance = instance
         self.database = database
-        self.is_closed = False
 
         self._ddl_statements = []
+        self.transactions = []
+
+        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.
+
+        :rtype: :class:`google.cloud.spanner_v1.session.Session`
+        :returns: Cloud Spanner session object ready to use.
+        """
+        return self._pool.get()
 
     def cursor(self):
         self._raise_if_closed()
@@ -142,18 +174,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. Rollback
+        will be performed on all the pending transactions.
         """
+        for transaction in self.transactions:
+            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)
+        else:
+            for transaction in self.transactions:
+                transaction.commit()
+
+            self.transactions = []
 
     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)
+        else:
+            for transaction in self.transactions:
+                transaction.rollback()
+
+            self.transactions = []
 
     def __enter__(self):
         return self

google/cloud/spanner_dbapi/cursor.py

@@ -68,6 +68,7 @@ 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
 
@@ -88,6 +89,20 @@ 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)
+            self._itr = PeekIterator(self._res)
+            return
+
         # Classify whether this is a read-only SQL statement.
         try:
             classification = classify_stmt(sql)

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(