feat: Add PingingPool check on session age

Mostafa Khalil · Mostafa Khalil · commit a1dcc9504bf3 · 2024-01-06T22:48:49.000+04:00
diff --git a/google/cloud/spanner_v1/pool.py b/google/cloud/spanner_v1/pool.py
@@ -385,6 +385,8 @@ class PingingPool(AbstractSessionPool):
     :param database_role: (Optional) user-assigned database_role for the session.
     """
 
+    SESSION_MAX_AGE = 28 * 24 * 60 * 60
+
     def __init__(
         self,
         size=10,
@@ -448,8 +450,9 @@ def get(self, timeout=None):
             timeout = self.default_timeout
 
         ping_after, session = self._sessions.get(block=True, timeout=timeout)
+        session_age = (_NOW() - session._created_at).total_seconds()
 
-        if _NOW() > ping_after:
+        if _NOW() > ping_after or session_age >= self.SESSION_MAX_AGE:
             # Using session.exists() guarantees the returned session exists.
             # session.ping() uses a cached result in the backend which could
             # result in a recently deleted session being returned.
@@ -481,6 +484,11 @@ def clear(self):
             else:
                 session.delete()
 
+    def _new_session(self):
+        session = super()._new_session()
+        session._created_at = _NOW()
+        return session
+
     def ping(self):
         """Refresh maybe-expired sessions in the pool.
 
diff --git a/tests/unit/test_pool.py b/tests/unit/test_pool.py
@@ -531,6 +531,48 @@ def test_get_hit_w_ping_expired(self):
         self.assertTrue(SESSIONS[0]._exists_checked)
         self.assertFalse(pool._sessions.full())
 
+    def test_get_hit_w_created(self):
+        import datetime
+
+        pool = self._make_one(size=4)
+        database = _Database("name")
+        SESSIONS = [_Session(database)] * 4
+        database._sessions.extend(SESSIONS)
+        pool.bind(database)
+
+        session_max_age = 28 * 24 * 60 * 60
+        SESSIONS[0]._created_at = datetime.datetime.utcnow() - datetime.timedelta(
+            seconds=session_max_age + 10
+        )
+
+        session = pool.get()
+
+        self.assertIs(session, SESSIONS[0])
+        self.assertTrue(session._exists_checked)
+        self.assertFalse(pool._sessions.full())
+
+    def test_get_hit_w_created_expired(self):
+        import datetime
+
+        pool = self._make_one(size=4)
+        database = _Database("name")
+        SESSIONS = [_Session(database)] * 5
+        database._sessions.extend(SESSIONS)
+        pool.bind(database)
+
+        session_max_age = 28 * 24 * 60 * 60
+        SESSIONS[0]._created_at = datetime.datetime.utcnow() - datetime.timedelta(
+            seconds=session_max_age
+        )
+        SESSIONS[0]._exists = False
+
+        session = pool.get()
+
+        self.assertIs(session, SESSIONS[4])
+        session.create.assert_called()
+        self.assertTrue(SESSIONS[0]._exists_checked)
+        self.assertFalse(pool._sessions.full())
+
     def test_get_empty_default_timeout(self):
         import queue
 
