From ecc658ddd5dcb788801137ec5a0e99aace0556d4 Mon Sep 17 00:00:00 2001
From: Anamarija Talijanac <anamarija.medic@gmail.com>
Date: Mon, 4 Nov 2019 23:15:23 +0100
Subject: [PATCH] validationQuery option added for connection pool

---
 .../runtime/AbstractDataSourceProducer.java    | 18 ++++++++++++++++++
 .../runtime/DataSourceRuntimeConfig.java       |  6 ++++++
 2 files changed, 24 insertions(+)

diff --git a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AbstractDataSourceProducer.java b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AbstractDataSourceProducer.java
index 9199e9b51..f084f7937 100644
--- a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AbstractDataSourceProducer.java
+++ b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AbstractDataSourceProducer.java
@@ -1,6 +1,8 @@
 package io.quarkus.agroal.runtime;
 
+import java.sql.Connection;
 import java.sql.Driver;
+import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
@@ -167,6 +169,22 @@ public abstract class AbstractDataSourceProducer {
         if (dataSourceRuntimeConfig.backgroundValidationInterval.isPresent()) {
             poolConfiguration.validationTimeout(dataSourceRuntimeConfig.backgroundValidationInterval.get());
         }
+        if (dataSourceRuntimeConfig.validationQuerySql.isPresent()) {
+            String validationQuery = dataSourceRuntimeConfig.validationQuerySql.get();
+            poolConfiguration.connectionValidator(new ConnectionValidator() {
+
+                @Override
+                public boolean isValid(Connection connection) {
+                    try (Statement stmt = connection.createStatement()) {
+                        stmt.execute(validationQuery);
+                        return true;
+                    } catch (Exception e) {
+                        log.warn("Connection validation failed", e);
+                    }
+                    return false;
+                }
+            });
+        }
         if (dataSourceRuntimeConfig.idleRemovalInterval.isPresent()) {
             poolConfiguration.reapTimeout(dataSourceRuntimeConfig.idleRemovalInterval.get());
         }
diff --git a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSourceRuntimeConfig.java b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSourceRuntimeConfig.java
index ecff627c7..ca46829b3 100644
--- a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSourceRuntimeConfig.java
+++ b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/DataSourceRuntimeConfig.java
@@ -124,4 +124,10 @@ public class DataSourceRuntimeConfig {
      */
     @ConfigItem
     public Optional<String> newConnectionSql;
+
+    /**
+     * Query executed to validate a connection.
+     */
+    @ConfigItem
+    public Optional<String> validationQuerySql;
 }