feat: Session leaks

google/cloud/spanner_dbapi/connection.py

@@ -582,9 +582,18 @@ def connect(
             raise ValueError("project in url does not match client object project")
 
     instance = client.instance(instance_id)
-    conn = Connection(
-        instance, instance.database(database_id, pool=pool) if database_id else None
-    )
+
+    if database_id:
+        database = instance.database(
+            database_id,
+            pool=pool,
+            close_inactive_transactions=False,
+        )
+        database._pool.logging_enabled = False
+    else:
+        database = None
+
+    conn = Connection(instance, database)
     if pool is not None:
         conn._own_pool = False
 

google/cloud/spanner_v1/_helpers.py

@@ -38,6 +38,12 @@
     + "numeric has a whole component with precision {}"
 )
 
+LONG_RUNNING_TRANSACTION_ERR_MSG = "Transaction has been closed as it was running for more than 60 minutes. If transaction is expected to run long, run as batch or partitioned DML."
+
+# Constants
+DELETE_LONG_RUNNING_TRANSACTION_FREQUENCY_SEC = 120
+DELETE_LONG_RUNNING_TRANSACTION_THRESHOLD_SEC = 3600
+
 
 def _try_to_coerce_bytes(bytestring):
     """Try to coerce a byte string into the right thing based on Python

google/cloud/spanner_v1/batch.py

@@ -29,6 +29,7 @@
 from google.cloud.spanner_v1 import RequestOptions
 from google.cloud.spanner_v1._helpers import _retry
 from google.cloud.spanner_v1._helpers import _check_rst_stream_error
+from google.cloud.spanner_v1._helpers import LONG_RUNNING_TRANSACTION_ERR_MSG
 from google.api_core.exceptions import InternalServerError
 
 
@@ -144,6 +145,8 @@ def _check_state(self):
         """
         if self.committed is not None:
             raise ValueError("Batch already committed")
+        if self._session is None:
+            raise ValueError(LONG_RUNNING_TRANSACTION_ERR_MSG)
 
     def commit(self, return_commit_stats=False, request_options=None):
         """Commit mutations to the database.

google/cloud/spanner_v1/database.py

@@ -113,6 +113,11 @@ class Database(object):
                    is `True` to log commit statistics. If not passed, a logger
                    will be created when needed that will log the commit statistics
                    to stdout.
+
+    :type close_inactive_transactions: boolean
+    :param close_inactive_transactions: (Optional) If set to True, the database will automatically close inactive transactions that have been running for longer than 60 minutes which may cause session leaks.
+                    By default, this is set to False.
+
     :type encryption_config:
         :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
         or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
@@ -142,6 +147,7 @@ def __init__(
         ddl_statements=(),
         pool=None,
         logger=None,
+        close_inactive_transactions=False,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
@@ -160,6 +166,7 @@ def __init__(
         self._default_leader = None
         self.log_commit_stats = False
         self._logger = logger
+        self._close_inactive_transactions = close_inactive_transactions
         self._encryption_config = encryption_config
         self._database_dialect = database_dialect
         self._database_role = database_role
@@ -366,7 +373,7 @@ def enable_drop_protection(self, value):
     def logger(self):
         """Logger used by the database.
 
-        The default logger will log commit stats at the log level INFO using
+        The default logger will log at the log level INFO using
         `sys.stderr`.
 
         :rtype: :class:`logging.Logger` or `None`
@@ -381,6 +388,14 @@ def logger(self):
             self._logger.addHandler(ch)
         return self._logger
 
+    @property
+    def close_inactive_transactions(self):
+        """Whether the database has has closing inactive transactions enabled. Default: False.
+        :rtype: bool
+        :returns: True if closing inactive transactions is enabled, else False.
+        """
+        return self._close_inactive_transactions
+
     @property
     def spanner_api(self):
         """Helper for session-related API calls."""
@@ -647,7 +662,7 @@ def execute_partitioned_dml(
             )
 
         def execute_pdml():
-            with SessionCheckout(self._pool) as session:
+            with SessionCheckout(self._pool, isLongRunning=True) as session:
 
                 txn = api.begin_transaction(
                     session=session.name, options=txn_options, metadata=metadata
@@ -1020,6 +1035,7 @@ def __enter__(self):
         """Begin ``with`` block."""
         session = self._session = self._database._pool.get()
         batch = self._batch = Batch(session)
+        self._session._transaction = batch
         if self._request_options.transaction_tag:
             batch.transaction_tag = self._request_options.transaction_tag
         return batch
@@ -1038,7 +1054,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
                     "CommitStats: {}".format(self._batch.commit_stats),
                     extra={"commit_stats": self._batch.commit_stats},
                 )
-            self._database._pool.put(self._session)
+            if self._batch._session is not None:
+                self._database._pool.put(self._session)
+            self._session._transaction = None
 
 
 class SnapshotCheckout(object):
@@ -1062,22 +1080,27 @@ class SnapshotCheckout(object):
     def __init__(self, database, **kw):
         self._database = database
         self._session = None
+        self._snapshot = None
         self._kw = kw
 
     def __enter__(self):
         """Begin ``with`` block."""
         session = self._session = self._database._pool.get()
-        return Snapshot(session, **self._kw)
+        self._snapshot = Snapshot(session, **self._kw)
+        self._session._transaction = self._snapshot
+        return self._snapshot
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         """End ``with`` block."""
-        if isinstance(exc_val, NotFound):
-            # If NotFound exception occurs inside the with block
-            # then we validate if the session still exists.
-            if not self._session.exists():
-                self._session = self._database._pool._new_session()
-                self._session.create()
-        self._database._pool.put(self._session)
+        if self._snapshot._session is not None:
+            if isinstance(exc_val, NotFound):
+                # If NotFound exception occurs inside the with block
+                # then we validate if the session still exists.
+                if not self._session.exists():
+                    self._session = self._database._pool._new_session()
+                    self._session.create()
+            self._database._pool.put(self._session)
+        self._session._transaction = None
 
 
 class BatchSnapshot(object):

google/cloud/spanner_v1/instance.py

@@ -429,6 +429,7 @@ def database(
         ddl_statements=(),
         pool=None,
         logger=None,
+        close_inactive_transactions=False,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
@@ -453,6 +454,10 @@ def database(
                        will be created when needed that will log the commit statistics
                        to stdout.
 
+        :type close_inactive_transactions: boolean
+        :param close_inactive_transactions: (Optional) Represents whether the database
+                        has close inactive transactions enabled or not. Default is False
+
         :type encryption_config:
             :class:`~google.cloud.spanner_admin_database_v1.types.EncryptionConfig`
             or :class:`~google.cloud.spanner_admin_database_v1.types.RestoreDatabaseEncryptionConfig`
@@ -481,6 +486,7 @@ def database(
             ddl_statements=ddl_statements,
             pool=pool,
             logger=logger,
+            close_inactive_transactions=close_inactive_transactions,
             encryption_config=encryption_config,
             database_dialect=database_dialect,
             database_role=database_role,