Use a property to enable OTEL for JDBC & replace driver with datasource

Author: michalvavrik

docs/src/main/asciidoc/opentelemetry.adoc

@@ -243,15 +243,16 @@ The https://github.com/open-telemetry/opentelemetry-java-instrumentation/tree/ma
 implementation("io.opentelemetry.instrumentation:opentelemetry-jdbc")
 ----
 
-As it uses a dedicated JDBC driver, you must configure your datasource and Hibernate ORM to use it.
+As it uses a dedicated JDBC datasource wrapper, you must enable JDBC tracing for your datasource:
 
 [source, properties]
 ----
+# enable tracing
+quarkus.datasource.jdbc.telemetry=true
+
+# configure datasource
 quarkus.datasource.db-kind=postgresql
-# add ':otel' to your database URL
-quarkus.datasource.jdbc.url=jdbc:otel:postgresql://localhost:5432/mydatabase
-# use the 'OpenTelemetryDriver' instead of the one for your database
-quarkus.datasource.jdbc.driver=io.opentelemetry.instrumentation.jdbc.OpenTelemetryDriver
+quarkus.datasource.jdbc.url=jdbc:postgresql://localhost:5432/mydatabase
 ----
 
 == Additional configuration

extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/AgroalProcessor.java

@@ -1,5 +1,7 @@
 package io.quarkus.agroal.deployment;
 
+import static io.quarkus.deployment.Capability.OPENTELEMETRY_TRACER;
+
 import java.sql.Driver;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -20,6 +22,7 @@
 import io.agroal.api.AgroalDataSource;
 import io.agroal.api.AgroalPoolInterceptor;
 import io.quarkus.agroal.DataSource;
+import io.quarkus.agroal.runtime.AgroalOpenTelemetryRecorder;
 import io.quarkus.agroal.runtime.AgroalRecorder;
 import io.quarkus.agroal.runtime.DataSourceJdbcBuildTimeConfig;
 import io.quarkus.agroal.runtime.DataSourceSupport;
@@ -77,8 +80,10 @@ void build(
             List<JdbcDriverBuildItem> jdbcDriverBuildItems,
             BuildProducer<ReflectiveClassBuildItem> reflectiveClass,
             BuildProducer<NativeImageResourceBuildItem> resource,
+            Capabilities capabilities,
             BuildProducer<ExtensionSslNativeSupportBuildItem> sslNativeSupport,
             BuildProducer<AggregatedDataSourceBuildTimeConfigBuildItem> aggregatedConfig,
+            BuildProducer<OpenTelemetryJDBCInstrumentationBuildItem> otelInstrumentationActiveProducer,
             CurateOutcomeBuildItem curateOutcomeBuildItem) throws Exception {
         if (dataSourcesBuildTimeConfig.driver.isPresent() || dataSourcesBuildTimeConfig.url.isPresent()) {
             throw new ConfigurationException(
@@ -96,6 +101,7 @@ void build(
             return;
         }
 
+        boolean otelJdbcInstrumentationActive = false;
         for (AggregatedDataSourceBuildTimeConfigBuildItem aggregatedDataSourceBuildTimeConfig : aggregatedDataSourceBuildTimeConfigs) {
             validateBuildTimeConfig(aggregatedDataSourceBuildTimeConfig);
 
@@ -105,13 +111,23 @@ void build(
                                 DataSources.TRACING_DRIVER_CLASSNAME));
             }
 
+            if (aggregatedDataSourceBuildTimeConfig.getJdbcConfig().telemetry && !otelJdbcInstrumentationActive) {
+                otelJdbcInstrumentationActive = true;
+            }
+
             reflectiveClass
                     .produce(new ReflectiveClassBuildItem(true, false,
                             aggregatedDataSourceBuildTimeConfig.getResolvedDriverClass()));
 
             aggregatedConfig.produce(aggregatedDataSourceBuildTimeConfig);
         }
 
+        if (otelJdbcInstrumentationActive && capabilities.isPresent(OPENTELEMETRY_TRACER)) {
+            // at least one datasource is using OpenTelemetry JDBC instrumentation,
+            // therefore we need to prepare OpenTelemetry data source wrapper
+            otelInstrumentationActiveProducer.produce(new OpenTelemetryJDBCInstrumentationBuildItem());
+        }
+
         // For now, we can't push the security providers to Agroal so we need to include
         // the service file inside the image. Hopefully, we will get an entry point to
         // resolve them at build time and push them to Agroal soon.
@@ -200,6 +216,8 @@ void generateDataSourceSupportBean(AgroalRecorder recorder,
             List<AggregatedDataSourceBuildTimeConfigBuildItem> aggregatedBuildTimeConfigBuildItems,
             SslNativeConfigBuildItem sslNativeConfig,
             Capabilities capabilities,
+            Optional<OpenTelemetryJDBCInstrumentationBuildItem> otelInstrumentationActive,
+            AgroalOpenTelemetryRecorder agroalOpenTelemetryRecorder,
             BuildProducer<AdditionalBeanBuildItem> additionalBeans,
             BuildProducer<SyntheticBeanBuildItem> syntheticBeanBuildItemBuildProducer,
             BuildProducer<UnremovableBeanBuildItem> unremovableBeans) {
@@ -210,6 +228,13 @@ void generateDataSourceSupportBean(AgroalRecorder recorder,
             return;
         }
 
+        if (otelInstrumentationActive.isPresent()) {
+            // prepare OpenTelemetry datasource wrapper
+            // this code must only run when the OpenTelemetry is active
+            // to avoid native failures
+            agroalOpenTelemetryRecorder.prepareOpenTelemetryAgroalDatasource();
+        }
+
         // make a DataSourceProducer bean
         additionalBeans.produce(AdditionalBeanBuildItem.builder().addBeanClasses(DataSources.class).setUnremovable()
                 .setDefaultScope(DotNames.SINGLETON).build());

extensions/agroal/deployment/src/main/java/io/quarkus/agroal/deployment/OpenTelemetryJDBCInstrumentationBuildItem.java

@@ -0,0 +1,9 @@
+package io.quarkus.agroal.deployment;
+
+import io.quarkus.builder.item.SimpleBuildItem;
+
+/**
+ * Purely marker build item that tells us to prepare OpenTelemetry JDBC instrumentation.
+ */
+final class OpenTelemetryJDBCInstrumentationBuildItem extends SimpleBuildItem {
+}

extensions/agroal/runtime/pom.xml

@@ -71,6 +71,13 @@
             <scope>test</scope>
         </dependency>
 
+        <!-- Required for OpenTelemetry JDBC instrumentation -->
+        <dependency>
+            <groupId>io.opentelemetry.instrumentation</groupId>
+            <artifactId>opentelemetry-jdbc</artifactId>
+            <optional>true</optional>
+        </dependency>
+
     </dependencies>
 
     <build>

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

@@ -0,0 +1,24 @@
+package io.quarkus.agroal.runtime;
+
+import java.util.function.Function;
+
+import io.agroal.api.AgroalDataSource;
+import io.quarkus.runtime.annotations.Recorder;
+
+@Recorder
+public class AgroalOpenTelemetryRecorder {
+
+    /**
+     * Tell {@link DataSources} to use the OpenTelemetry datasource wrapper for
+     * data sources with activated telemetry. We need to set this way only when
+     * optional OpenTelemetry JDBC instrumentation dependency is present to avoid native failures.
+     */
+    public void prepareOpenTelemetryAgroalDatasource() {
+        DataSources.setOpenTelemetryDatasourceTransformer(new Function<AgroalDataSource, AgroalDataSource>() {
+            @Override
+            public AgroalDataSource apply(AgroalDataSource agroalDataSource) {
+                return new OpenTelemetryAgroalDataSource(agroalDataSource);
+            }
+        });
+    }
+}

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

@@ -43,4 +43,10 @@ public class DataSourceJdbcBuildTimeConfig {
      */
     @ConfigItem(defaultValue = "false")
     public boolean tracing = false;
+
+    /**
+     * Enable OpenTelemetry JDBC instrumentation.
+     */
+    @ConfigItem(defaultValue = "false")
+    public boolean telemetry = false;
 }

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

@@ -145,4 +145,10 @@ public class DataSourceJdbcRuntimeConfig {
     @ConfigItem
     public DataSourceJdbcTracingRuntimeConfig tracing = new DataSourceJdbcTracingRuntimeConfig();
 
+    /**
+     * Enable OpenTelemetry JDBC instrumentation.
+     */
+    @ConfigItem(name = "telemetry.enabled", defaultValueDocumentation = "false if quarkus.datasource.jdbc.telemetry=false and true if quarkus.datasource.jdbc.telemetry=true")
+    public Optional<Boolean> telemetry;
+
 }

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

@@ -7,6 +7,7 @@
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Objects;
 import java.util.ServiceLoader;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
@@ -68,6 +69,7 @@ public class DataSources {
     public static final String TRACING_DRIVER_CLASSNAME = "io.opentracing.contrib.jdbc.TracingDriver";
     private static final String JDBC_URL_PREFIX = "jdbc:";
     private static final String JDBC_TRACING_URL_PREFIX = "jdbc:tracing:";
+    private static volatile Function<AgroalDataSource, AgroalDataSource> OTEL_DATASOURCE_TRANSFORMER = null;
 
     private final DataSourcesBuildTimeConfig dataSourcesBuildTimeConfig;
     private final DataSourcesRuntimeConfig dataSourcesRuntimeConfig;
@@ -252,6 +254,12 @@ public AgroalDataSource doCreateDataSource(String dataSourceName) {
             dataSource.setPoolInterceptors(interceptorList);
         }
 
+        if (dataSourceJdbcBuildTimeConfig.telemetry && dataSourceJdbcRuntimeConfig.telemetry.orElse(true)) {
+            // active OpenTelemetry JDBC instrumentation by wrapping AgroalDatasource
+            // use the transformer as we can't reference optional OpenTelemetry classes here
+            dataSource = OTEL_DATASOURCE_TRANSFORMER.apply(dataSource);
+        }
+
         return dataSource;
     }
 
@@ -445,4 +453,12 @@ public void stop() {
         }
     }
 
+    /**
+     * Set a function that will wrap {@link AgroalDataSource} with the OpenTelemetry datasource.
+     */
+    static void setOpenTelemetryDatasourceTransformer(Function<AgroalDataSource, AgroalDataSource> otelDatasourceTransformer) {
+        Objects.requireNonNull(otelDatasourceTransformer);
+        OTEL_DATASOURCE_TRANSFORMER = otelDatasourceTransformer;
+    }
+
 }

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

@@ -0,0 +1,66 @@
+package io.quarkus.agroal.runtime;
+
+import java.sql.SQLException;
+import java.sql.ShardingKeyBuilder;
+import java.util.Collection;
+import java.util.List;
+
+import io.agroal.api.AgroalDataSource;
+import io.agroal.api.AgroalDataSourceMetrics;
+import io.agroal.api.AgroalPoolInterceptor;
+import io.agroal.api.configuration.AgroalDataSourceConfiguration;
+import io.opentelemetry.api.GlobalOpenTelemetry;
+import io.opentelemetry.instrumentation.jdbc.datasource.OpenTelemetryDataSource;
+
+/**
+ * The {@link AgroalDataSource} wrapper that actives OpenTelemetry JDBC instrumentation.
+ */
+public class OpenTelemetryAgroalDataSource extends OpenTelemetryDataSource implements AgroalDataSource {
+
+    private final AgroalDataSource delegate;
+
+    public OpenTelemetryAgroalDataSource(AgroalDataSource delegate) {
+        super(delegate, GlobalOpenTelemetry.get());
+        this.delegate = delegate;
+    }
+
+    @Override
+    public boolean isHealthy(boolean newConnection) throws SQLException {
+        return delegate.isHealthy(newConnection);
+    }
+
+    @Override
+    public AgroalDataSourceConfiguration getConfiguration() {
+        return delegate.getConfiguration();
+    }
+
+    @Override
+    public AgroalDataSourceMetrics getMetrics() {
+        return delegate.getMetrics();
+    }
+
+    @Override
+    public void flush(FlushMode mode) {
+        delegate.flush(mode);
+    }
+
+    @Override
+    public void setPoolInterceptors(Collection<? extends AgroalPoolInterceptor> interceptors) {
+        delegate.setPoolInterceptors(interceptors);
+    }
+
+    @Override
+    public List<AgroalPoolInterceptor> getPoolInterceptors() {
+        return delegate.getPoolInterceptors();
+    }
+
+    @Override
+    public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
+        return delegate.createShardingKeyBuilder();
+    }
+
+    @Override
+    public void close() {
+        delegate.close();
+    }
+}

extensions/opentelemetry/deployment/pom.xml

@@ -111,11 +111,6 @@
             <artifactId>quarkus-agroal-deployment</artifactId>
             <scope>test</scope>
         </dependency>
-        <dependency>
-            <groupId>io.quarkus</groupId>
-            <artifactId>quarkus-jdbc-h2-deployment</artifactId>
-            <scope>test</scope>
-        </dependency>
         <dependency>
             <groupId>io.opentelemetry.instrumentation</groupId>
             <artifactId>opentelemetry-jdbc</artifactId>
@@ -137,6 +132,17 @@
             <artifactId>opentelemetry-aws-xray</artifactId>
             <scope>test</scope>
         </dependency>
+        <dependency>
+            <groupId>io.quarkus</groupId>
+            <artifactId>quarkus-jdbc-h2-deployment</artifactId>
+            <scope>test</scope>
+        </dependency>
+
+        <!-- Needed for InMemorySpanExporter to verify captured traces -->
+        <dependency>
+            <groupId>io.opentelemetry</groupId>
+            <artifactId>opentelemetry-sdk-testing</artifactId>
+        </dependency>
     </dependencies>
 
     <build>