OpenTelemetryDriver don't require DriverManager for underlying drivers

instrumentation/jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriver.java

@@ -32,6 +32,8 @@
 import java.sql.DriverPropertyInfo;
 import java.sql.SQLException;
 import java.sql.SQLFeatureNotSupportedException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.Properties;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -49,6 +51,7 @@ public final class OpenTelemetryDriver implements Driver {
 
   private static final String URL_PREFIX = "jdbc:otel:";
   private static final AtomicBoolean REGISTERED = new AtomicBoolean();
+  private static final Collection<Driver> DRIVER_CANDIDATES = new ArrayList<>();
 
   static {
     try {
@@ -99,8 +102,52 @@ public static boolean isRegistered() {
     return REGISTERED.get();
   }
 
-  private static Driver findDriver(String realUrl) {
-    for (Driver candidate : Collections.list(DriverManager.getDrivers())) {
+  /**
+   * Add driver candidate that wasn't registered against {@link DriverManager}. This is mainly
+   * useful when drivers are initialized at native image build time, and you don't want to postpone
+   * class initialization, for runtime initialization could make driver incompatible with other
+   * systems. Drivers registered with this method have higher priority over those registered against
+   * {@link DriverManager}.
+   *
+   * @param driver {@link Driver} that should be registered
+   */
+  public static void addDriverCandidate(Driver driver) {
+    if (driver != null) {
+      DRIVER_CANDIDATES.add(driver);
+    }
+  }
+
+  /**
+   * Remove driver added via {@link #addDriverCandidate(Driver)}.
+   *
+   * @param driver {@link Driver} that should be unregistered
+   * @return true if the driver was unregistered
+   */
+  public static boolean removeDriverCandidate(Driver driver) {
+    return DRIVER_CANDIDATES.remove(driver);
+  }
+
+  /**
+   * Find driver that accepts {@code realUrl}. Drivers registered against {@link #DRIVER_CANDIDATES}
+   * are preferred over {@link DriverManager} drivers.
+   */
+  static Driver findDriver(String realUrl) {
+    Driver driver = null;
+    if (!DRIVER_CANDIDATES.isEmpty()) {
+      driver = findDriver(realUrl, DRIVER_CANDIDATES);
+    }
+    if (driver == null) {
+      driver = findDriver(realUrl, Collections.list(DriverManager.getDrivers()));
+    }
+    if (driver == null) {
+      throw new IllegalStateException("Unable to find a driver that accepts url: " + realUrl);
+    }
+
+    return driver;
+  }
+
+  private static Driver findDriver(String realUrl, Collection<Driver> drivers) {
+    for (Driver candidate : drivers) {
       try {
         if (!(candidate instanceof OpenTelemetryDriver) && candidate.acceptsURL(realUrl)) {
           return candidate;
@@ -109,8 +156,7 @@ private static Driver findDriver(String realUrl) {
         // intentionally ignore exception
       }
     }
-
-    throw new IllegalStateException("Unable to find a driver that accepts url: " + realUrl);
+    return null;
   }
 
   /**

instrumentation/jdbc/library/src/test/groovy/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriverTest.groovy

@@ -9,6 +9,7 @@ import io.opentelemetry.instrumentation.api.instrumenter.Instrumenter
 import io.opentelemetry.instrumentation.jdbc.internal.OpenTelemetryConnection
 import spock.lang.Specification
 
+import java.sql.Driver
 import java.sql.DriverManager
 import java.sql.SQLFeatureNotSupportedException
 
@@ -121,6 +122,83 @@ class OpenTelemetryDriverTest extends Specification {
     connection == null
   }
 
+  def "verify add driver candidate"() {
+    when:
+    deregisterTestDriver()
+    TestDriver driver = new TestDriver()
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(driver)
+    def connection = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null)
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(driver)
+
+    then:
+    connection != null
+    connection instanceof OpenTelemetryConnection
+  }
+
+  def "verify remove driver candidate"() {
+    when:
+    deregisterTestDriver()
+    TestDriver driver = new TestDriver()
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(driver)
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(driver)
+    def connection = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null)
+
+    then:
+    connection == null
+    def e = thrown(IllegalStateException)
+    e.message == "Unable to find a driver that accepts url: jdbc:test:"
+  }
+
+  def "verify driver candidate has higher priority"() {
+    when:
+    deregisterTestDriver()
+    TestDriver localDriver = new TestDriver()
+    TestDriver globalDriver = new TestDriver()
+
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(localDriver)
+    DriverManager.registerDriver(globalDriver)
+    Driver winner = OpenTelemetryDriver.INSTANCE.findDriver("jdbc:test:")
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(localDriver)
+
+    then:
+    winner == localDriver
+    winner != globalDriver
+  }
+
+  def "verify two clashing driver candidates"() {
+    when:
+    TestDriver localDriver1 = new TestDriver()
+    TestDriver localDriver2 = new TestDriver()
+
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(localDriver1)
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(localDriver2)
+    Driver winner = OpenTelemetryDriver.INSTANCE.findDriver("jdbc:test:")
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(localDriver1)
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(localDriver2)
+
+    then:
+    winner == localDriver1
+    winner != localDriver2
+  }
+
+  def "verify drivers in DriverManager are used as fallback"() {
+    when:
+    registerTestDriver()
+    TestDriver localDriver = new TestDriver() {
+      boolean acceptsURL(String url) {
+        return false
+      }
+    }
+    OpenTelemetryDriver.INSTANCE.addDriverCandidate(localDriver)
+
+    Driver winner = OpenTelemetryDriver.INSTANCE.findDriver("jdbc:test:")
+    OpenTelemetryDriver.INSTANCE.removeDriverCandidate(localDriver)
+
+    then:
+    winner != null
+    winner != localDriver
+  }
+
   def "verify connection with accepted url"() {
     when:
     registerTestDriver()
@@ -177,4 +255,10 @@ class OpenTelemetryDriverTest extends Specification {
     }
   }
 
+  private static void deregisterTestDriver() {
+    if ((DriverManager.drivers.any { it instanceof TestDriver })) {
+      DriverManager.deregisterDriver(new TestDriver())
+    }
+  }
+
 }