Adding tests that verify the recover callbacks and improving the two previous tests with asserts on the session id.

Author: steveellis

domain-models-runtime/src/main/java/org/folio/rest/persist/PostgresClient.java

@@ -508,6 +508,11 @@ public static void closeAllClients() {
     PG_POOLS_READER.clear();
   }
 
+  public static void clearConnectionCache() {
+    CACHED_CONNECTION_MANAGER.closeAllConnectionsInCache();
+    CACHED_CONNECTION_MANAGER.clearCache();
+  }
+
 
   private void init() throws Exception {
 

domain-models-runtime/src/main/java/org/folio/rest/persist/cache/CachedConnectionManager.java

@@ -39,6 +39,10 @@ public void clearCache() {
     LOG.debug("Cleared connection cache");
   }
 
+  public void closeAllConnectionsInCache() {
+    connectionCache.closeAll();
+  }
+
   public void removeFromCache(CachedPgConnection connection)  {
     connectionCache.remove(connection);
   }

domain-models-runtime/src/main/java/org/folio/rest/persist/cache/ConnectionCache.java

@@ -102,6 +102,14 @@ public void clear() {
     }
   }
 
+  public void closeAll() {
+    synchronized (cache) {
+      for (CachedPgConnection conn : cache) {
+        conn.getWrappedConnection().close();
+      }
+    }
+  }
+
   public int size() {
     synchronized (cache) {
       return cache.size();

domain-models-runtime/src/test/java/org/folio/rest/persist/cache/CachedConnectionManagerIT.java

@@ -2,7 +2,6 @@
 
 import io.vertx.ext.unit.TestContext;
 import io.vertx.ext.unit.junit.VertxUnitRunner;
-import io.vertx.pgclient.PgConnection;
 import io.vertx.pgclient.impl.PgConnectionImpl;
 import org.folio.postgres.testing.PostgresTesterContainer;
 import org.folio.rest.impl.TenantAPI;
@@ -11,6 +10,7 @@
 import org.folio.rest.persist.PostgresClientHelper;
 import org.folio.rest.tools.utils.VertxUtils;
 import org.junit.AfterClass;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 import org.junit.runner.RunWith;
@@ -32,41 +32,48 @@ public static void afterClass() {
     PostgresClientHelper.setSharedPgPool(false);
   }
 
+  @Before
+  public void beforeEach() {
+    PostgresClientHelper.setSharedPgPool(true);
+    PostgresClient.clearConnectionCache();
+  }
+
   @Test
   public void shouldNotCreateNestedWrappers(TestContext context) {
-    PostgresClientHelper.setSharedPgPool(true);
     tenantPost(new TenantAPI(), context, null, "tenant1");
     tenantPost(new TenantAPI(), context, null, "tenant2");
 
     var async = context.async();
     var pgClient1 = PostgresClient.getInstance(vertx, "tenant1");
     pgClient1.getConnection().onComplete(context.asyncAssertSuccess(conn1 -> {
+      var originalSessionId = ((CachedPgConnection) conn1).getSessionId();
       // Closing the connection returns it to the cache so it can be recycled.
       conn1.close();
 
       var pgClient2 = PostgresClient.getInstance(vertx, "tenant2");
       pgClient2.getConnection().onComplete(context.asyncAssertSuccess(conn2 -> {
         var recycledCached = (CachedPgConnection) conn2;
 
+        assertEquals(originalSessionId, recycledCached.getSessionId());
         // Should still wrap the REAL connection, not the cached one
         assertEquals(PgConnectionImpl.class, recycledCached.getWrappedConnection().getClass());
         assertNotEquals(CachedPgConnection.class, recycledCached.getWrappedConnection().getClass());
 
-        async.complete();
+        conn2.close().onComplete(context.asyncAssertSuccess(v -> async.complete()));
       }));
     }));
   }
 
   @Test
   public void shouldUpdateConnectionMetadataAfterConnectionSucceeds(TestContext context) {
-    PostgresClientHelper.setSharedPgPool(true);
     tenantPost(new TenantAPI(), context, null, "tenant1");
     tenantPost(new TenantAPI(), context, null, "tenant2");
 
     var async = context.async();
     var pgClient1 = PostgresClient.getInstance(vertx, "tenant1");
     pgClient1.getConnection().onComplete(context.asyncAssertSuccess(conn1 -> {
       var recycledCached1 = (CachedPgConnection) conn1;
+      var originalSessionId = recycledCached1.getSessionId();
 
       assertEquals("tenant1", recycledCached1.getTenantId());
       assertEquals("tenant1_raml_module_builder", recycledCached1.getSchemaName());
@@ -82,10 +89,62 @@ public void shouldUpdateConnectionMetadataAfterConnectionSucceeds(TestContext co
       pgClient2.getConnection().onComplete(context.asyncAssertSuccess(conn2 -> {
         var recycledCached2 = (CachedPgConnection) conn2;
 
+        assertEquals(originalSessionId, recycledCached2.getSessionId());
         assertEquals("tenant2", recycledCached2.getTenantId());
         assertEquals("tenant2_raml_module_builder", recycledCached2.getSchemaName());
 
-        async.complete();
+        conn2.close().onComplete(context.asyncAssertSuccess(v -> async.complete()));
+      }));
+    }));
+  }
+
+  @Test
+  public void shouldRecoverOnSetRoleFailure(TestContext context) {
+    tenantPost(new TenantAPI(), context, null, "tenant1");
+
+    var async = context.async();
+    var pgClient1 = PostgresClient.getInstance(vertx, "tenant1");
+    pgClient1.getConnection().onComplete(context.asyncAssertSuccess(conn1 -> {
+      var cachedConn = (CachedPgConnection) conn1;
+      var originalSessionId = cachedConn.getSessionId();
+
+      // Closing the connection returns it to the cache so it can be recycled.
+      conn1.close();
+
+      // Try to get a connection for a tenant that does not exist. This will cause the SET ROLE to fail.
+      var pgClient2 = PostgresClient.getInstance(vertx, "nonexistent");
+      pgClient2.getConnection().onComplete(context.asyncAssertFailure(failure -> {
+        // After the failure, the connection should have been returned to the cache with its original tenant.
+        // Let's try to get it again for the original tenant.
+        pgClient1.getConnection().onComplete(context.asyncAssertSuccess(conn3 -> {
+          var cachedConn3 = (CachedPgConnection) conn3;
+
+          // The connection should be the same one we got before.
+          assertEquals(originalSessionId, cachedConn3.getSessionId());
+          assertEquals("tenant1", cachedConn3.getTenantId());
+          assertEquals("tenant1_raml_module_builder", cachedConn3.getSchemaName());
+
+          conn3.close().onComplete(context.asyncAssertSuccess(v -> async.complete()));
+        }));
+      }));
+    }));
+  }
+
+  @Test
+  public void shouldRecoverOnSetRoleFailureForNewConnection(TestContext context) {
+    tenantPost(new TenantAPI(), context, null, "diku");
+
+    var async = context.async();
+    // Try to get a connection for a tenant that does not exist. This will cause the SET ROLE to fail.
+    var pgClient1 = PostgresClient.getInstance(vertx, "nonexistent");
+    pgClient1.getConnection().onComplete(context.asyncAssertFailure(failure -> {
+      // After the failure, the underlying connection should have been closed and returned to the pool.
+      // Let's try to get a connection for a valid tenant to prove the pool is healthy.
+      var pgClient2 = PostgresClient.getInstance(vertx, "diku");
+      pgClient2.getConnection().onComplete(context.asyncAssertSuccess(conn -> {
+        // The connection should be a valid one.
+        assertEquals("diku", ((CachedPgConnection) conn).getTenantId());
+        conn.close().onComplete(context.asyncAssertSuccess(v -> async.complete()));
       }));
     }));
   }