dbapi: make global session pool and tie sessions to it

odeke-em · odeke-em · commit 1142dafae52a · 2020-02-26T17:34:19.000-08:00
Apparently invoking spanner_v1.Database.session() by-passes using the provided session pool as advised at: googleapis/python-spanner#10 (comment) so this change using the logic that's provided with the context manager, but our use case to use Transactions with ORMs hadn't been thought of, so we have to use private methods and the methods used by spanner_v1.pool.SessionCheckout. Fixes #291
diff --git a/packages/django-google-spanner/spanner/dbapi/__init__.py b/packages/django-google-spanner/spanner/dbapi/__init__.py
@@ -30,6 +30,8 @@
 # are working properly, we'll update the threadsafety level.
 threadsafety = 0
 
+global_session_pool = spanner.pool.BurstyPool()
+
 
 def connect(project=None, instance=None, database=None, credentials_uri=None, user_agent=None):
     """
@@ -68,11 +70,31 @@ def connect(project=None, instance=None, database=None, credentials_uri=None, us
     if not client_instance.exists():
         raise ProgrammingError("instance '%s' does not exist." % instance)
 
-    db = client_instance.database(database)
+    db = client_instance.database(database, pool=global_session_pool)
     if not db.exists():
         raise ProgrammingError("database '%s' does not exist." % database)
 
-    return Connection(db)
+    # Correctly retrieve a session from the global session pool.
+    # See:
+    #   * https://github.com/orijtech/django-spanner/issues/291
+    #   * https://github.com/googleapis/python-spanner/issues/10#issuecomment-585056760
+    #
+    # Adapted from:
+    #   https://bit.ly/3c8MK6p: python-spanner, Git hash 997a03477b07ec39c7184
+    #   google/cloud/spanner_v1/pool.py#L514-L535
+    # TODO: File a bug to googleapis/python-spanner asking for a convenience
+    # method since invoke database.session() gives the wrong result
+    # yet requires a context manager wrapped with SessionCheckout
+    # and needs accessing private methods, which leaks the details of the
+    # implementation in order to try to use it correctly.
+    pool = db._pool
+    session_checkout = spanner.pool.SessionCheckout(pool)
+    session = session_checkout.__enter__()
+    if not session.exists():
+        session.create()
+    return_session = lambda: session_checkout.__exit__() # noqa
+
+    return Connection(db, session, return_session)
 
 
 __all__ = [
diff --git a/packages/django-google-spanner/spanner/dbapi/connection.py b/packages/django-google-spanner/spanner/dbapi/connection.py
@@ -10,11 +10,9 @@
 
 
 class Connection(object):
-    def __init__(self, db_handle):
-        sess = db_handle.session()
-        if not sess.exists():
-            sess.create()
-        self.__sess = sess
+    def __init__(self, db_handle, session, discard_session):
+        self.__sess = session
+        self.__discard_session = discard_session
         self.__txn = None
         self.__dbhandle = db_handle
         self.__closed = False
@@ -37,7 +35,8 @@ def __enter__(self):
 
     def __clear(self):
         self.__dbhandle = None
-        self.__sess.delete()
+        self.__discard_session()
+        self.__sess = None
 
     def __exit__(self, etype, value, traceback):
         self.__raise_if_already_closed()
