Skip to content

Commit

Permalink
feat: clear session pool on connection close (#543)
Browse files Browse the repository at this point in the history
Clear the pool on connection close to avoid keeping Cloud Spanner sessions
alive.
  • Loading branch information
Ilya Gurov committed Nov 4, 2020
1 parent 927e178 commit 14e4cac
Show file tree
Hide file tree
Showing 2 changed files with 36 additions and 6 deletions.
13 changes: 12 additions & 1 deletion google/cloud/spanner_dbapi/connection.py
Expand Up @@ -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):
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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
29 changes: 24 additions & 5 deletions tests/unit/spanner_dbapi/test_connection.py
Expand Up @@ -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
Expand Down Expand Up @@ -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)

Expand Down Expand Up @@ -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

0 comments on commit 14e4cac

Please sign in to comment.