Merge pull request quarkusio#13375 from jmartisk/master-issue-13326

Ability to exclude particular data sources from the health check
gsmet · Apr 1, 2021 · 23c529a · 23c529a
2 parents 8896cad + 1749f00
commit 23c529a
Show file tree

Hide file tree

Showing 10 changed files with 133 additions and 8 deletions.
diff --git a/core/deployment/src/main/java/io/quarkus/deployment/Capability.java b/core/deployment/src/main/java/io/quarkus/deployment/Capability.java
@@ -68,6 +68,7 @@ public enum Capability {
      */
     SMALLRYE_OPENTRACING,
     SPRING_WEB,
+    SMALLRYE_HEALTH,
     SMALLRYE_OPENAPI,
     OPENSHIFT_CLIENT;
 

diff --git a/extensions/agroal/deployment/pom.xml b/extensions/agroal/deployment/pom.xml
@@ -82,6 +82,11 @@
             <artifactId>quarkus-smallrye-metrics-deployment</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-smallrye-health-deployment</artifactId>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>

diff --git a/extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/AgroalProcessor.java b/extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/AgroalProcessor.java
@@ -33,6 +33,7 @@
 import io.quarkus.datasource.deployment.spi.DefaultDataSourceDbKindBuildItem;
 import io.quarkus.datasource.runtime.DataSourceBuildTimeConfig;
 import io.quarkus.datasource.runtime.DataSourcesBuildTimeConfig;
+import io.quarkus.datasource.runtime.DataSourcesExcludedFromHealthChecks;
 import io.quarkus.datasource.runtime.DataSourcesRuntimeConfig;
 import io.quarkus.deployment.Capabilities;
 import io.quarkus.deployment.Capability;
@@ -329,8 +330,19 @@ private String resolveDriver(String dataSourceName, String dbKind,
     }
 
     @BuildStep
-    HealthBuildItem addHealthCheck(DataSourcesBuildTimeConfig dataSourcesBuildTimeConfig) {
-        return new HealthBuildItem("io.quarkus.agroal.runtime.health.DataSourceHealthCheck",
-                dataSourcesBuildTimeConfig.healthEnabled);
+    @Record(ExecutionTime.STATIC_INIT)
+    HealthBuildItem addHealthCheck(Capabilities capabilities, DataSourcesBuildTimeConfig dataSourcesBuildTimeConfig,
+            AgroalRecorder recorder, BuildProducer<SyntheticBeanBuildItem> syntheticBeans) {
+        if (capabilities.isPresent(Capability.SMALLRYE_HEALTH)) {
+            syntheticBeans.produce(SyntheticBeanBuildItem.configure(DataSourcesExcludedFromHealthChecks.class)
+                    .scope(Singleton.class)
+                    .unremovable()
+                    .supplier(recorder.dataSourcesExcludedFromHealthChecks(dataSourcesBuildTimeConfig))
+                    .done());
+            return new HealthBuildItem("io.quarkus.agroal.runtime.health.DataSourceHealthCheck",
+                    dataSourcesBuildTimeConfig.healthEnabled);
+        } else {
+            return null;
+        }
     }
 }
diff --git a/...l/deployment/src/test/java/io/quarkus/agroal/test/DataSourceHealthCheckExclusionTest.java b/...l/deployment/src/test/java/io/quarkus/agroal/test/DataSourceHealthCheckExclusionTest.java
@@ -0,0 +1,26 @@
+package io.quarkus.agroal.test;
+
+import org.hamcrest.CoreMatchers;
+import org.jboss.shrinkwrap.api.ShrinkWrap;
+import org.jboss.shrinkwrap.api.spec.JavaArchive;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+import io.quarkus.test.QuarkusUnitTest;
+import io.restassured.RestAssured;
+
+public class DataSourceHealthCheckExclusionTest {
+
+    @RegisterExtension
+    static final QuarkusUnitTest config = new QuarkusUnitTest()
+            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class))
+            .withConfigurationResource("application-datasources-with-health.properties");
+
+    @Test
+    public void testDataSourceHealthCheckExclusion() {
+        RestAssured.when().get("/health/ready")
+                .then()
+                .body("status", CoreMatchers.equalTo("UP"));
+    }
+
+}
diff --git a/...sions/agroal/deployment/src/test/resources/application-datasources-with-health.properties b/...sions/agroal/deployment/src/test/resources/application-datasources-with-health.properties
@@ -0,0 +1,13 @@
+quarkus.datasource.db-kind=h2
+quarkus.datasource.username=username-default
+quarkus.datasource.jdbc.url=jdbc:h2:tcp://localhost/mem:default
+quarkus.datasource.jdbc.max-size=13
+
+# this data source is broken, but will be excluded from the health check, so the overall check should pass
+quarkus.datasource.brokenDS.db-kind=h2
+quarkus.datasource.brokenDS.username=username1
+quarkus.datasource.brokenDS.jdbc.driver=org.h2.Driver
+quarkus.datasource.brokenDS.jdbc.url=BROKEN
+quarkus.datasource.brokenDS.health-exclude=true
+
+quarkus.datasource.health.enabled=true
diff --git a/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AgroalRecorder.java b/extensions/agroal/runtime/src/main/java/io/quarkus/agroal/runtime/AgroalRecorder.java
@@ -1,8 +1,15 @@
 package io.quarkus.agroal.runtime;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
 import java.util.function.Supplier;
 
 import io.agroal.api.AgroalDataSource;
+import io.quarkus.datasource.common.runtime.DataSourceUtil;
+import io.quarkus.datasource.runtime.DataSourceBuildTimeConfig;
+import io.quarkus.datasource.runtime.DataSourcesBuildTimeConfig;
+import io.quarkus.datasource.runtime.DataSourcesExcludedFromHealthChecks;
 import io.quarkus.datasource.runtime.DataSourcesRuntimeConfig;
 import io.quarkus.runtime.annotations.Recorder;
 
@@ -29,4 +36,24 @@ public AgroalDataSource get() {
         };
     }
 
+    public Supplier<DataSourcesExcludedFromHealthChecks> dataSourcesExcludedFromHealthChecks(
+            DataSourcesBuildTimeConfig dataSourcesBuildTimeConfig) {
+        return new Supplier<DataSourcesExcludedFromHealthChecks>() {
+            @Override
+            public DataSourcesExcludedFromHealthChecks get() {
+                List<String> excludedNames = new ArrayList<>();
+                if (dataSourcesBuildTimeConfig.defaultDataSource.healthExclude) {
+                    excludedNames.add(DataSourceUtil.DEFAULT_DATASOURCE_NAME);
+                }
+                for (Map.Entry<String, DataSourceBuildTimeConfig> dataSource : dataSourcesBuildTimeConfig.namedDataSources
+                        .entrySet()) {
+                    if (dataSource.getValue().healthExclude) {
+                        excludedNames.add(dataSource.getKey());
+                    }
+                }
+                return new DataSourcesExcludedFromHealthChecks(excludedNames);
+            }
+        };
+    }
+
 }
diff --git a/.../agroal/runtime/src/main/java/io/quarkus/agroal/runtime/health/DataSourceHealthCheck.java b/.../agroal/runtime/src/main/java/io/quarkus/agroal/runtime/health/DataSourceHealthCheck.java
@@ -3,6 +3,7 @@
 import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 
@@ -17,6 +18,8 @@
 import org.eclipse.microprofile.health.Readiness;
 
 import io.quarkus.arc.Arc;
+import io.quarkus.datasource.common.runtime.DataSourceUtil;
+import io.quarkus.datasource.runtime.DataSourcesExcludedFromHealthChecks;
 
 @Readiness
 @ApplicationScoped
@@ -27,14 +30,21 @@ public class DataSourceHealthCheck implements HealthCheck {
     @PostConstruct
     protected void init() {
         Set<Bean<?>> beans = Arc.container().beanManager().getBeans(DataSource.class);
+        DataSourcesExcludedFromHealthChecks excluded = Arc.container().instance(DataSourcesExcludedFromHealthChecks.class)
+                .get();
+        List<String> excludedNames = excluded.getExcludedNames();
         for (Bean<?> bean : beans) {
             if (bean.getName() == null) {
-                // this is the default DataSource: retrieve it by type
-                DataSource defaultDs = Arc.container().instance(DataSource.class).get();
-                dataSources.put(DEFAULT_DS, defaultDs);
+                if (!excludedNames.contains(DataSourceUtil.DEFAULT_DATASOURCE_NAME)) {
+                    // this is the default DataSource: retrieve it by type
+                    DataSource defaultDs = Arc.container().instance(DataSource.class).get();
+                    dataSources.put(DEFAULT_DS, defaultDs);
+                }
             } else {
-                DataSource ds = (DataSource) Arc.container().instance(bean.getName()).get();
-                dataSources.put(bean.getName(), ds);
+                if (!excludedNames.contains(bean.getName())) {
+                    DataSource ds = (DataSource) Arc.container().instance(bean.getName()).get();
+                    dataSources.put(bean.getName(), ds);
+                }
             }
         }
     }

diff --git a/...source/runtime/src/main/java/io/quarkus/datasource/runtime/DataSourceBuildTimeConfig.java b/...source/runtime/src/main/java/io/quarkus/datasource/runtime/DataSourceBuildTimeConfig.java
@@ -22,4 +22,13 @@ public class DataSourceBuildTimeConfig {
     @ConfigItem
     public DevServicesBuildTimeConfig devservices;
 
+    /**
+     * Whether this particular data source should be excluded from the health check if
+     * the general health check for data sources is enabled.
+     * <p>
+     * By default, the health check includes all configured data sources (if it is enabled).
+     */
+    @ConfigItem(defaultValue = "false")
+    public boolean healthExclude;
+
 }
diff --git a/...time/src/main/java/io/quarkus/datasource/runtime/DataSourcesExcludedFromHealthChecks.java b/...time/src/main/java/io/quarkus/datasource/runtime/DataSourcesExcludedFromHealthChecks.java
@@ -0,0 +1,16 @@
+package io.quarkus.datasource.runtime;
+
+import java.util.List;
+
+public class DataSourcesExcludedFromHealthChecks {
+
+    private List<String> excludedNames;
+
+    public DataSourcesExcludedFromHealthChecks(List<String> names) {
+        this.excludedNames = names;
+    }
+
+    public List<String> getExcludedNames() {
+        return excludedNames;
+    }
+}
diff --git a/...ployment/src/main/java/io/quarkus/smallrye/health/deployment/SmallRyeHealthProcessor.java b/...ployment/src/main/java/io/quarkus/smallrye/health/deployment/SmallRyeHealthProcessor.java
@@ -46,6 +46,7 @@
 import io.quarkus.deployment.annotations.BuildStep;
 import io.quarkus.deployment.annotations.ExecutionTime;
 import io.quarkus.deployment.annotations.Record;
+import io.quarkus.deployment.builditem.CapabilityBuildItem;
 import io.quarkus.deployment.builditem.FeatureBuildItem;
 import io.quarkus.deployment.builditem.GeneratedResourceBuildItem;
 import io.quarkus.deployment.builditem.HotDeploymentWatchedFileBuildItem;
@@ -143,6 +144,11 @@ void healthCheck(BuildProducer<AdditionalBeanBuildItem> buildItemBuildProducer,
         }
     }
 
+    @BuildStep
+    public CapabilityBuildItem capability() {
+        return new CapabilityBuildItem(Capability.SMALLRYE_HEALTH);
+    }
+
     @BuildStep
     @Record(ExecutionTime.STATIC_INIT)
     @SuppressWarnings("unchecked")