🍒 9774 - Add a feature flag to disable extra calls to the DB in JDBC instrumentation

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DriverInstrumentation.java

@@ -64,6 +64,7 @@ public void methodAdvice(MethodTransformer transformer) {
   }
 
   public static class DriverAdvice {
+
     @Advice.OnMethodExit(suppress = Throwable.class)
     public static void addDBInfo(
         @Advice.Argument(0) final String url,
@@ -75,8 +76,10 @@ public static void addDBInfo(
       }
       String connectionUrl = url;
       Properties connectionProps = props;
-      if (null == props
-          || !Boolean.parseBoolean(props.getProperty("oracle.jdbc.useShardingDriverConnection"))) {
+      if (JDBCDecorator.FETCH_DB_METADATA_ON_CONNECT
+          && (null == props
+              || !Boolean.parseBoolean(
+                  props.getProperty("oracle.jdbc.useShardingDriverConnection")))) {
         try {
           DatabaseMetaData metaData = connection.getMetaData();
           connectionUrl = metaData.getURL();

dd-java-agent/instrumentation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/JDBCDecorator.java

@@ -28,6 +28,7 @@
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.HashSet;
+import java.util.Properties;
 import java.util.Set;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -61,6 +62,10 @@ public class JDBCDecorator extends DatabaseClientDecorator<DBInfo> {
       Config.get().isDbmTracePreparedStatements();
   public static final boolean DBM_ALWAYS_APPEND_SQL_COMMENT =
       Config.get().isDbmAlwaysAppendSqlComment();
+  public static final boolean FETCH_DB_METADATA_ON_CONNECT =
+      Config.get().isDbMetadataFetchingOnConnectEnabled();
+  private static final boolean FETCH_DB_METADATA_ON_QUERY =
+      Config.get().isDbMetadataFetchingOnQueryEnabled();
 
   private volatile boolean warnedAboutDBMPropagationMode = false; // to log a warning only once
 
@@ -183,7 +188,7 @@ public static DBInfo parseDBInfo(
         } catch (Throwable ignore) {
         }
         if (dbInfo == null) {
-          // couldn't find DBInfo anywhere, so fall back to default
+          // couldn't find DBInfo from a previous call anywhere, so we try to fetch it from the DB
           dbInfo = parseDBInfoFromConnection(connection);
         }
         // store the DBInfo on the outermost connection instance to avoid future searches
@@ -202,7 +207,7 @@ public String getDbService(final DBInfo dbInfo) {
   }
 
   public static DBInfo parseDBInfoFromConnection(final Connection connection) {
-    if (connection == null) {
+    if (connection == null || !FETCH_DB_METADATA_ON_QUERY) {
       // we can log here, but it risks to be too verbose
       return DBInfo.DEFAULT;
     }
@@ -211,16 +216,19 @@ public static DBInfo parseDBInfoFromConnection(final Connection connection) {
       final DatabaseMetaData metaData = connection.getMetaData();
       final String url;
       if (metaData != null && (url = metaData.getURL()) != null) {
+        Properties clientInfo = null;
         try {
-          dbInfo = JDBCConnectionUrlParser.extractDBInfo(url, connection.getClientInfo());
+          clientInfo = connection.getClientInfo();
         } catch (final Throwable ex) {
-          // getClientInfo is likely not allowed.
-          dbInfo = JDBCConnectionUrlParser.extractDBInfo(url, null);
+          // getClientInfo is likely not allowed, we can still extract info from the url alone
+          log.debug("Could not get client info from DB", ex);
         }
+        dbInfo = JDBCConnectionUrlParser.extractDBInfo(url, clientInfo);
       } else {
         dbInfo = DBInfo.DEFAULT;
       }
     } catch (final SQLException se) {
+      log.debug("Could not get metadata from DB", se);
       dbInfo = DBInfo.DEFAULT;
     }
     return dbInfo;

dd-java-agent/instrumentation/jdbc/src/test/groovy/DriverInstrumentationMetadataFetchingTest.groovy

@@ -0,0 +1,271 @@
+import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
+import static datadog.trace.api.config.TraceInstrumentationConfig.DB_METADATA_FETCHING_ON_CONNECT
+
+import datadog.trace.agent.test.InstrumentationSpecification
+import datadog.trace.api.DDSpanTypes
+import datadog.trace.bootstrap.instrumentation.api.Tags
+import java.sql.Connection
+import java.sql.DatabaseMetaData
+import java.sql.SQLException
+import java.sql.Statement
+import test.TestConnection
+import test.TestDatabaseMetaData
+import test.TestDriver
+
+abstract class DriverInstrumentationMetadataFetchingTestBase extends InstrumentationSpecification {
+
+  def "test dbInfo extraction from metadata or URL with username"() {
+    setup:
+    def originalUrl = "jdbc:postgresql://original-host:1234/originaldb"
+    def metadataUrl = "jdbc:postgresql://metadata-host:5678/metadatadb"
+
+    def metadata = new TestDatabaseMetaData() {
+        @Override
+        String getURL() throws SQLException {
+          return metadataUrl
+        }
+
+        @Override
+        String getUserName() throws SQLException {
+          return "metadata-user"
+        }
+      }
+
+    def testConnection = new TestConnection(false)
+    testConnection.setMetaData(metadata)
+
+    def driver = new TestDriver() {
+        @Override
+        Connection connect(String url, Properties info) {
+          return testConnection
+        }
+      }
+
+    def props = new Properties()
+    props.put("user", "original-user")
+
+    // Expected values depend on whether metadata fetching is enabled
+    def expectedDbInstance = shouldFetchMetadataOnConnect() ? "metadatadb" : "originaldb"
+    def expectedDbUser = shouldFetchMetadataOnConnect() ? "metadata-user" : "original-user"
+
+    when:
+    def connection = driver.connect(originalUrl, props)
+    Statement statement = connection.createStatement()
+    runUnderTrace("parent") {
+      statement.execute("SELECT 1")
+    }
+
+    then:
+    assertTraces(1) {
+      trace(2) {
+        span(0) {
+          operationName "parent"
+        }
+        span(1) {
+          operationName "postgresql.query"
+          spanType DDSpanTypes.SQL
+          childOf span(0)
+          tags(false) {
+            "$Tags.DB_INSTANCE" expectedDbInstance
+            "$Tags.DB_USER" expectedDbUser
+          }
+        }
+      }
+    }
+
+    cleanup:
+    statement?.close()
+    connection?.close()
+  }
+
+  def "test driver connect with null metadata URL"() {
+    setup:
+    def originalUrl = "jdbc:postgresql://original-host:5432/originaldb"
+
+    def metadata = new TestDatabaseMetaData() {
+        @Override
+        String getURL() throws SQLException {
+          return null
+        }
+      }
+
+    def testConnection = new TestConnection(false)
+    testConnection.setMetaData(metadata)
+
+    def driver = new TestDriver() {
+        @Override
+        Connection connect(String url, Properties info) {
+          return testConnection
+        }
+      }
+
+    def props = new Properties()
+
+    when:
+    def connection = driver.connect(originalUrl, props)
+    Statement statement = connection.createStatement()
+    runUnderTrace("parent") {
+      statement.execute("SELECT 1")
+    }
+
+    then:
+    assertTraces(1) {
+      trace(2) {
+        span(0) {
+          operationName "parent"
+        }
+        span(1) {
+          operationName "postgresql.query"
+          spanType DDSpanTypes.SQL
+          childOf span(0)
+          tags(false) {
+            // Should fallback to original URL regardless of flag when metadata URL is null
+            "$Tags.DB_INSTANCE" "originaldb"
+          }
+        }
+      }
+    }
+
+    cleanup:
+    statement?.close()
+    connection?.close()
+  }
+
+  def "test driver connect with metadata exception"() {
+    setup:
+    def originalUrl = "jdbc:postgresql://original-host:5432/originaldb"
+
+    def testConnection = new TestConnection(false) {
+        @Override
+        DatabaseMetaData getMetaData() throws SQLException {
+          throw new SQLException("Test exception")
+        }
+      }
+
+    def driver = new TestDriver() {
+        @Override
+        Connection connect(String url, Properties info) {
+          return testConnection
+        }
+      }
+
+    def props = new Properties()
+
+    when:
+    def connection = driver.connect(originalUrl, props)
+    Statement statement = connection.createStatement()
+    runUnderTrace("parent") {
+      statement.execute("SELECT 1")
+    }
+
+    then:
+    assertTraces(1) {
+      trace(2) {
+        span(0) {
+          operationName "parent"
+        }
+        span(1) {
+          operationName "postgresql.query"
+          spanType DDSpanTypes.SQL
+          childOf span(0)
+          tags(false) {
+            // Should fallback to original URL when exception occurs
+            "$Tags.DB_INSTANCE" "originaldb"
+          }
+        }
+      }
+    }
+
+    cleanup:
+    statement?.close()
+    connection?.close()
+  }
+
+  def "test driver connect with Oracle sharding driver"() {
+    setup:
+    def originalUrl = "jdbc:oracle:thin:@original-host:1521:orcl"
+    def metadataUrl = "jdbc:oracle:thin:@metadata-host:1521:orcl"
+
+    def metadata = new TestDatabaseMetaData()
+    metadata.setURL(metadataUrl)
+
+    def testConnection = new TestConnection(false)
+    testConnection.setMetaData(metadata)
+
+    def driver = new TestDriver() {
+        @Override
+        Connection connect(String url, Properties info) {
+          return testConnection
+        }
+      }
+
+    def props = new Properties()
+    props.setProperty("oracle.jdbc.useShardingDriverConnection", "true")
+
+    when:
+    def connection = driver.connect(originalUrl, props)
+    Statement statement = connection.createStatement()
+    runUnderTrace("parent") {
+      statement.execute("SELECT 1")
+    }
+
+    then:
+    assertTraces(1) {
+      trace(2) {
+        span(0) {
+          operationName "parent"
+        }
+        span(1) {
+          operationName "oracle.query"
+          spanType DDSpanTypes.SQL
+          childOf span(0)
+          tags(false) {
+            // Should use original URL for Oracle sharding driver regardless of flag
+            "$Tags.DB_INSTANCE" "orcl"
+          }
+        }
+      }
+    }
+
+    cleanup:
+    statement?.close()
+    connection?.close()
+  }
+
+  abstract boolean shouldFetchMetadataOnConnect()
+}
+
+/**
+ * Test with metadata fetching enabled on connect (default behavior)
+ */
+class DriverInstrumentationWithMetadataOnConnectForkedTest extends DriverInstrumentationMetadataFetchingTestBase {
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(DB_METADATA_FETCHING_ON_CONNECT, "true")
+  }
+
+  @Override
+  boolean shouldFetchMetadataOnConnect() {
+    return true
+  }
+}
+
+/**
+ * Test with metadata fetching disabled on connect
+ */
+class DriverInstrumentationWithoutMetadataOnConnectForkedTest extends DriverInstrumentationMetadataFetchingTestBase {
+
+  @Override
+  void configurePreAgent() {
+    super.configurePreAgent()
+    injectSysConfig(DB_METADATA_FETCHING_ON_CONNECT, "false")
+  }
+
+  @Override
+  boolean shouldFetchMetadataOnConnect() {
+    return false
+  }
+}
+