From 14e4cac77fd9ba5cf421c56c636528ec77b82451 Mon Sep 17 00:00:00 2001
From: Ilya Gurov <ilya.faer@mail.ru>
Date: Thu, 5 Nov 2020 02:29:00 +0300
Subject: [PATCH] feat: clear session pool on connection close (#543)

Clear the pool on connection close to avoid keeping Cloud Spanner sessions
alive.
---
 google/cloud/spanner_dbapi/connection.py    | 13 ++++++++-
 tests/unit/spanner_dbapi/test_connection.py | 29 +++++++++++++++++----
 2 files changed, 36 insertions(+), 6 deletions(-)

diff --git a/google/cloud/spanner_dbapi/connection.py b/google/cloud/spanner_dbapi/connection.py
index b572c8573b..beb05a3173 100644
--- a/google/cloud/spanner_dbapi/connection.py
+++ b/google/cloud/spanner_dbapi/connection.py
@@ -43,6 +43,10 @@ def __init__(self, instance, database):
 
         self.is_closed = False
         self._autocommit = False
+        # indicator to know if the session pool used by
+        # this connection should be cleared on the
+        # connection close
+        self._own_pool = True
 
     @property
     def autocommit(self):
@@ -150,6 +154,9 @@ def close(self):
         ):
             self._transaction.rollback()
 
+        if self._own_pool:
+            self.database._pool.clear()
+
         self.is_closed = True
 
     def commit(self):
@@ -258,4 +265,8 @@ def connect(
     if not database.exists():
         raise ValueError("database '%s' does not exist." % database_id)
 
-    return Connection(instance, database)
+    conn = Connection(instance, database)
+    if pool is not None:
+        conn._own_pool = False
+
+    return conn
diff --git a/tests/unit/spanner_dbapi/test_connection.py b/tests/unit/spanner_dbapi/test_connection.py
index d545472c57..99aa0aa47b 100644
--- a/tests/unit/spanner_dbapi/test_connection.py
+++ b/tests/unit/spanner_dbapi/test_connection.py
@@ -34,14 +34,19 @@ def _get_client_info(self):
 
         return ClientInfo(user_agent=self.USER_AGENT)
 
-    def _make_connection(self):
+    def _make_connection(self, pool=None):
         from google.cloud.spanner_dbapi import Connection
         from google.cloud.spanner_v1.instance import Instance
 
         # We don't need a real Client object to test the constructor
         instance = Instance(self.INSTANCE, client=None)
         database = instance.database(self.DATABASE)
-        return Connection(instance, database)
+
+        conn = Connection(instance, database)
+        if pool is not None:
+            conn._own_pool = False
+
+        return conn
 
     def test_property_autocommit_setter(self):
         from google.cloud.spanner_dbapi import Connection
@@ -239,9 +244,7 @@ def test_run_prior_DDL_statements(self):
                 connection.run_prior_DDL_statements()
 
     def test_context(self):
-        from google.cloud.spanner_dbapi import Connection
-
-        connection = Connection(self.INSTANCE, self.DATABASE)
+        connection = self._make_connection()
         with connection as conn:
             self.assertEqual(conn, connection)
 
@@ -316,3 +319,19 @@ def test_sessions_pool(self):
             ):
                 connect("test-instance", database_id, pool=pool)
                 database_mock.assert_called_once_with(database_id, pool=pool)
+
+    def test_clearing_pool_on_close(self):
+        connection = self._make_connection()
+        with mock.patch.object(
+            connection.database._pool, "clear"
+        ) as pool_clear_mock:
+            connection.close()
+            pool_clear_mock.assert_called_once_with()
+
+    def test_global_pool(self):
+        connection = self._make_connection(pool=mock.Mock())
+        with mock.patch.object(
+            connection.database._pool, "clear"
+        ) as pool_clear_mock:
+            connection.close()
+            assert not pool_clear_mock.called