feat: Session leaks

google/cloud/spanner_dbapi/connection.py

@@ -583,7 +583,15 @@ def connect(
 
     instance = client.instance(instance_id)
     conn = Connection(
-        instance, instance.database(database_id, pool=pool) if database_id else None
+        instance,
+        instance.database(
+            database_id,
+            pool=pool,
+            logging_enabled=False,
+            close_inactive_transactions=False,
+        )
+        if database_id
+        else None,
     )
     if pool is not None:
         conn._own_pool = False

google/cloud/spanner_v1/_helpers.py

@@ -38,6 +38,10 @@
     + "numeric has a whole component with precision {}"
 )
 
+# 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

@@ -144,6 +144,10 @@ def _check_state(self):
         """
         if self.committed is not None:
             raise ValueError("Batch already committed")
+        if self._session is None:
+            raise ValueError(
+                "Transaction has been closed as it was running for more than 60 minutes"
+            )
 
     def commit(self, return_commit_stats=False, request_options=None):
         """Commit mutations to the database.

google/cloud/spanner_v1/database.py

@@ -108,11 +108,20 @@ class Database(object):
                  passed, the database will construct an instance of
                  :class:`~google.cloud.spanner_v1.pool.BurstyPool`.
 
+    :type logging_enabled: boolean
+    :param logging_enabled: (Optional) Represents whether the database
+                    has logging enabled or not. Default is True
+
     :type logger: :class:`logging.Logger`
     :param logger: (Optional) a custom logger that is used if `log_commit_stats`
                    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) 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`
@@ -141,7 +150,9 @@ def __init__(
         instance,
         ddl_statements=(),
         pool=None,
+        logging_enabled=True,
         logger=None,
+        close_inactive_transactions=False,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
@@ -159,7 +170,9 @@ def __init__(
         self._encryption_info = None
         self._default_leader = None
         self.log_commit_stats = False
+        self._logging_enabled = logging_enabled
         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
@@ -339,6 +352,14 @@ def database_role(self):
         """
         return self._database_role
 
+    @property
+    def logging_enabled(self):
+        """Whether the database has logging enabled. Default: True.
+        :rtype: bool
+        :returns: True if logging is enabled, else False.
+        """
+        return self._logging_enabled
+
     @property
     def reconciling(self):
         """Whether the database is currently reconciling.
@@ -366,7 +387,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 +402,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 +676,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 +1049,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 +1068,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 +1094,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

@@ -428,7 +428,9 @@ def database(
         database_id,
         ddl_statements=(),
         pool=None,
+        logging_enabled=True,
         logger=None,
+        close_inactive_transactions=False,
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
@@ -447,12 +449,20 @@ def database(
                     :class:`~google.cloud.spanner_v1.pool.AbstractSessionPool`.
         :param pool: (Optional) session pool to be used by database.
 
+        :type logging_enabled: boolean
+        :param logging_enabled: (Optional) Represents whether the database
+                        has logging enabled or not. Default is True
+
         :type logger: :class:`logging.Logger`
         :param logger: (Optional) a custom logger that is used if `log_commit_stats`
                        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) 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`
@@ -480,7 +490,9 @@ def database(
             self,
             ddl_statements=ddl_statements,
             pool=pool,
+            logging_enabled=logging_enabled,
             logger=logger,
+            close_inactive_transactions=close_inactive_transactions,
             encryption_config=encryption_config,
             database_dialect=database_dialect,
             database_role=database_role,