OpenTelemetryDriver don't require DriverManager for underlying drivers

closes: #7028
open-telemetry · Nov 7, 2022 · 2b39337 · 2b39337
1 parent cd64119
commit 2b39337
Show file tree

Hide file tree

Showing 2 changed files with 76 additions and 1 deletion.
diff --git a/...jdbc/library/src/main/java/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriver.java b/...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> DRIVERS = new ArrayList<>();
 
   static {
     try {
@@ -99,8 +102,49 @@ public static boolean isRegistered() {
     return REGISTERED.get();
   }
 
+
+  /**
+   * Register underlying driver that wasn't registered against {@link DriverManager}. This is mainly
+   * useful when drivers are initialized at build time, and you don't want to postpone class
+   * initialization (f.e. runtime initialization could make driver incompatible with other systems).
+   * Drivers registered with this method has higher priority over those registered against
+   * {@link DriverManager}.
+   *
+   * @param driver {@link Driver} that should be registered
+   */
+  public static void registerDriver(Driver driver) {
+    if (driver != null) {
+      DRIVERS.add(driver);
+    }
+  }
+
+  /**
+   * Unregister driver added via {@link #registerDriver(Driver)}.
+   *
+   * @param driver {@link Driver} that should be unregistered
+   * @return true if the driver was unregistered
+   */
+  public static boolean deregisterDriver(Driver driver) {
+    return DRIVERS.remove(driver);
+  }
+
+  /**
+   * Find driver that accepts {@code realUrl}. Drivers registered against {@link #DRIVERS} are
+   * preferred over {@link DriverManager} drivers.
+   */
   private static Driver findDriver(String realUrl) {
-    for (Driver candidate : Collections.list(DriverManager.getDrivers())) {
+    Driver driver = null;
+    if (!DRIVERS.isEmpty()) {
+      driver = findDriver(realUrl, DRIVERS);
+    }
+    if (driver == null) {
+      driver = findDriver(realUrl, Collections.list(DriverManager.getDrivers()));
+    }
+    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;

diff --git a/...rary/src/test/groovy/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriverTest.groovy b/...rary/src/test/groovy/io/opentelemetry/instrumentation/jdbc/OpenTelemetryDriverTest.groovy
@@ -121,6 +121,31 @@ class OpenTelemetryDriverTest extends Specification {
     connection == null
   }
 
+  def "verify registration of local driver"() {
+    when:
+    deregisterTestDriver()
+    TestDriver driver = new TestDriver()
+    OpenTelemetryDriver.INSTANCE.registerDriver(driver)
+    def connection = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null)
+    OpenTelemetryDriver.INSTANCE.deregisterDriver(driver)
+
+    then:
+    connection != null
+    connection instanceof OpenTelemetryConnection
+  }
+
+  def "verify deregistration of local driver"() {
+    when:
+    deregisterTestDriver()
+    TestDriver driver = new TestDriver()
+    OpenTelemetryDriver.INSTANCE.registerDriver(driver)
+    OpenTelemetryDriver.INSTANCE.deregisterDriver(driver)
+    def connection = OpenTelemetryDriver.INSTANCE.connect("jdbc:otel:test:", null)
+
+    then:
+    connection == null
+  }
+
   def "verify connection with accepted url"() {
     when:
     registerTestDriver()
@@ -177,4 +202,10 @@ class OpenTelemetryDriverTest extends Specification {
     }
   }
 
+  private static void deregisterTestDriver() {
+    if ((DriverManager.drivers.any { it instanceof TestDriver })) {
+      DriverManager.deregisterDriver(new TestDriver())
+    }
+  }
+
 }