feat: support transactions management

google/cloud/spanner_dbapi/__init__.py

@@ -6,8 +6,11 @@
 
 """Connection-based DB API for Cloud Spanner."""
 
+import atexit
+
 from google.cloud import spanner_v1
 
+from google.cloud.spanner_dbapi import config
 from .connection import Connection
 from .exceptions import (
     DatabaseError,
@@ -93,7 +96,17 @@ def connect(
     if not database.exists():
         raise ValueError("database '%s' does not exist." % database_id)
 
-    return Connection(instance, database)
+    if config.default_pool is None:
+        config.default_pool = spanner_v1.BurstyPool()
+
+    return Connection(instance, database, config.default_pool)
+
+
+@atexit.register
+def _cleanup():
+    """Clear the sessions pool on a program termination."""
+    if config.default_pool is not None:
+        config.default_pool.clear()
 
 
 __all__ = [

google/cloud/spanner_dbapi/config.py

@@ -0,0 +1,12 @@
+# Copyright 2020 Google LLC
+#
+# Use of this source code is governed by a BSD-style
+# license that can be found in the LICENSE file or at
+# https://developers.google.com/open-source/licenses/bsd
+
+"""Configurations module for DB API."""
+
+# Cloud Spanner sessions pool, used by
+# default for the whole DB API package.
+# Is lazily initiated on a connection creation.
+default_pool = None

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
+from .exceptions import InterfaceError, ProgrammingError
 
-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,125 @@ 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.instance = instance
-        self.database = database
-        self.is_closed = False
+    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._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
+
+    @database.setter
+    def database(self, _):
+        """Related database setter.
+
+        Restricts replacing the related database.
+        """
+        raise ProgrammingError(
+            "Replacing the related database of the existing connection is restricted."
+        )
+
+    @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
+
+    @instance.setter
+    def instance(self, _):
+        """Related instance setter.
+
+        Restricts replacing the related instance.
+        """
+        raise ProgrammingError(
+            "Replacing the related instance of the existing connection is restricted."
+        )
+
+    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 instead.
+
+        :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 +250,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

@@ -11,7 +11,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_dbapi import config, connect, Connection
 
 
 def _make_credentials():
@@ -108,3 +108,20 @@ def test_connect_database_id(self):
                 database_mock.assert_called_once_with(DATABASE, pool=mock.ANY)
 
         self.assertIsInstance(connection, Connection)
+
+    def test_pool_reuse(self):
+        DATABASE = "test-database"
+        DATABASE2 = "test-database2"
+
+        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", DATABASE)
+
+                self.assertIsNotNone(config.default_pool)
+                self.assertEqual(connection._pool, config.default_pool)
+
+                connection2 = connect("test-instance", DATABASE2)
+                self.assertEqual(connection._pool, connection2._pool)

tests/spanner_dbapi/test_connection.py

@@ -11,7 +11,11 @@
 
 # import google.cloud.spanner_dbapi.exceptions as dbapi_exceptions
 
-from google.cloud.spanner_dbapi import Connection, InterfaceError
+from google.cloud.spanner_dbapi import (
+    Connection,
+    InterfaceError,
+    ProgrammingError,
+)
 from google.cloud.spanner_dbapi.connection import AUTOCOMMIT_MODE_WARNING
 from google.cloud.spanner_v1.database import Database
 from google.cloud.spanner_v1.instance import Instance
@@ -49,8 +53,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 +65,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(ProgrammingError):
+            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(ProgrammingError):
+            connection.instance = None