Merge pull request #20657 from juazugas/add-agroal-config-flush-on-close

Add jdbc config if connections should be flushed when returning to the pool

extensions/agroal/deployment/src/test/java/io/quarkus/agroal/test/DefaultDataSourceConfigTest.java

@@ -4,6 +4,7 @@
 import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.entry;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.sql.SQLException;
@@ -60,6 +61,7 @@ private static void testDataSource(AgroalDataSource dataSource, String username,
                 agroalConnectionFactoryConfiguration.jdbcTransactionIsolation());
         assertTrue(agroalConnectionFactoryConfiguration.trackJdbcResources());
         assertTrue(dataSource.getConfiguration().metricsEnabled());
+        assertFalse(configuration.flushOnClose());
         assertEquals(newConnectionSql, agroalConnectionFactoryConfiguration.initialSql());
         assertThat(agroalConnectionFactoryConfiguration.jdbcProperties())
                 .contains(

extensions/agroal/deployment/src/test/java/io/quarkus/agroal/test/FlushOnCloseDataSourceConfigTest.java

@@ -0,0 +1,95 @@
+package io.quarkus.agroal.test;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+
+import javax.inject.Inject;
+import javax.inject.Singleton;
+
+import org.jboss.logging.Logger;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.asset.EmptyAsset;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.agroal.api.AgroalDataSource;
+import io.agroal.api.AgroalPoolInterceptor;
+import io.agroal.api.configuration.AgroalConnectionPoolConfiguration;
+import io.quarkus.test.QuarkusUnitTest;
+
+class FlushOnCloseDataSourceConfigTest {
+
+    private static final Logger LOGGER = Logger.getLogger(FlushOnCloseDataSourceConfigTest.class);
+
+    @Inject
+    AgroalDataSource defaultDataSource;
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest()
+            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class)
+                    .addAsManifestResource(EmptyAsset.INSTANCE, "beans.xml")
+                    .addClasses(FlushOnCloseDataSourceConfigTest.class)
+                    .addClasses(ClientUserTrackerInterceptor.class))
+            .withConfigurationResource("base.properties")
+            // Setting Compatibility MODE to MySQL to enable all the client-info properties on the connection
+            .overrideConfigKey("quarkus.datasource.jdbc.url", "jdbc:h2:tcp://localhost/mem:flushing;MODE=MySQL")
+            .overrideConfigKey("quarkus.datasource.jdbc.min-size", "1")
+            .overrideConfigKey("quarkus.datasource.jdbc.initial-size", "1")
+            .overrideConfigKey("quarkus.datasource.jdbc.max-size", "2")
+            .overrideConfigKey("quarkus.datasource.jdbc.flush-on-close", "true");
+
+    @Test
+    void testFlushOnCloseDataSourceConnection() throws SQLException {
+        AgroalConnectionPoolConfiguration configuration = defaultDataSource.getConfiguration().connectionPoolConfiguration();
+
+        assertTrue(configuration.flushOnClose());
+
+        try (Connection connection = defaultDataSource.getConnection()) {
+            connection.setClientInfo("ClientUser", "1");
+            connection.prepareStatement("SELECT 1").executeQuery();
+            try (Connection conn2 = defaultDataSource.getConnection()) {
+                conn2.setClientInfo("ClientUser", "2");
+                conn2.prepareStatement("SELECT 1").executeQuery();
+            }
+        }
+
+        try (Connection connection = defaultDataSource.getConnection()) {
+            assertEquals("1", connection.getClientInfo("ClientUser"));
+            try (Connection conn2 = defaultDataSource.getConnection()) {
+                // new connection should not contain any mark ClientUser
+                assertNull(conn2.getClientInfo("ClientUser"));
+            }
+        }
+    }
+
+    @Singleton
+    static class ClientUserTrackerInterceptor implements AgroalPoolInterceptor {
+
+        @Override
+        public void onConnectionAcquire(Connection connection) {
+            try {
+                LOGGER.info("connection acquired ClientUser:" + connection.getClientInfo("ClientUser"));
+            } catch (SQLException e) {
+                Assertions.fail(e);
+            }
+        }
+
+        @Override
+        public void onConnectionReturn(Connection connection) {
+            try {
+                // Connection marked "ClientUser:2" must not return to the pool.
+                assertNotEquals("2", connection.getClientInfo("ClientUser"));
+                LOGGER.info("connection returned ClientUser:" + connection.getClientInfo("ClientUser"));
+            } catch (SQLException e) {
+                Assertions.fail(e);
+            }
+        }
+    }
+}

extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSourceJdbcRuntimeConfig.java

@@ -89,6 +89,12 @@ public class DataSourceJdbcRuntimeConfig {
     @ConfigItem
     public boolean extendedLeakReport;
 
+    /**
+     * Allows connections to be flushed upon return to the pool. It's not enabled by default.
+     */
+    @ConfigItem
+    public boolean flushOnClose;
+
     /**
      * When enabled Agroal will be able to produce a warning when a connection is returned
      * to the pool without the application having closed all open statements.

extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSources.java

@@ -308,6 +308,7 @@ public boolean isValid(Connection connection) {
             poolConfiguration.transactionRequirement(dataSourceJdbcRuntimeConfig.transactionRequirement.get());
         }
         poolConfiguration.enhancedLeakReport(dataSourceJdbcRuntimeConfig.extendedLeakReport);
+        poolConfiguration.flushOnClose(dataSourceJdbcRuntimeConfig.flushOnClose);
     }
 
     public DataSourceBuildTimeConfig getDataSourceBuildTimeConfig(String dataSourceName) {