Fix for MySQL multi-host internally_generated connection urls

DataDog · Mar 27, 2024 · d804bcf · d804bcf
1 parent 6167d5e
commit d804bcf
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 2 deletions.
diff --git a/...p/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/JDBCConnectionUrlParser.java b/...p/src/main/java/datadog/trace/bootstrap/instrumentation/jdbc/JDBCConnectionUrlParser.java
@@ -216,6 +216,12 @@ DBInfo.Builder doParse(final String jdbcUrl, final DBInfo.Builder builder) {
   MARIA_SUBPROTO() {
     @Override
     DBInfo.Builder doParse(final String jdbcUrl, final DBInfo.Builder builder) {
+      if (jdbcUrl.startsWith("**internally_generated**")) {
+        // there is nothing to parse
+        builder.host(null);
+        builder.port(null);
+        return builder;
+      }
       final int hostEndLoc;
       final int clusterSepLoc = jdbcUrl.indexOf(',');
       final int ipv6End = jdbcUrl.startsWith("[") ? jdbcUrl.indexOf(']') : -1;
@@ -824,7 +830,7 @@ public static DBInfo parse(String connectionUrl, final Properties props) {
       }
       return GENERIC_URL_LIKE.doParse(connectionUrl, parsedProps).build();
     } catch (final Exception e) {
-      ExceptionLogger.LOGGER.debug("Error parsing URL", e);
+      ExceptionLogger.LOGGER.debug("Error parsing URL {}", jdbcUrl, e);
       return parsedProps.build();
     }
   }

diff --git a/...entation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DriverInstrumentation.java b/...entation/jdbc/src/main/java/datadog/trace/instrumentation/jdbc/DriverInstrumentation.java
@@ -15,6 +15,7 @@
 import datadog.trace.bootstrap.instrumentation.jdbc.JDBCConnectionUrlParser;
 import java.sql.Connection;
 import java.sql.DatabaseMetaData;
+import java.sql.SQLException;
 import java.util.Map;
 import java.util.Properties;
 import net.bytebuddy.asm.Advice;
@@ -91,8 +92,21 @@ public static void addDBInfo(
       } catch (Throwable ignored) {
         // use original values
       }
+      if (connectionUrl.contains("**internally_generated**")) {
+        // Use the original connection URL for parsing
+        connectionUrl = url;
+      }
+      Connection connWithContext = connection;
+      try {
+        // Find an underlying connection object to which to attach the DBInfo. Otherwise, it may be
+        // attached to a wrapper connection object that is different from the one used in other
+        // parts of the instrumentation, resulting in a loss of context.
+        connWithContext = connection.unwrap(Connection.class);
+      } catch (SQLException e) {
+        // ignore
+      }
       DBInfo dbInfo = JDBCConnectionUrlParser.extractDBInfo(connectionUrl, connectionProps);
-      InstrumentationContext.get(Connection.class, DBInfo.class).put(connection, dbInfo);
+      InstrumentationContext.get(Connection.class, DBInfo.class).put(connWithContext, dbInfo);
     }
   }
 }
diff --git a/dd-java-agent/instrumentation/jdbc/src/test/groovy/JDBCConnectionUrlParserTest.groovy b/dd-java-agent/instrumentation/jdbc/src/test/groovy/JDBCConnectionUrlParserTest.groovy
@@ -70,6 +70,9 @@ class JDBCConnectionUrlParserTest extends AgentTestRunner {
     "jdbc:mysql://my.host:22/mydb?user=myuser&password=PW"                                                                                                                                                                                      | null     | "mysql"      | null          | "myuser"      | "my.host"                                                           | 22    | null                               | "mydb"
     "jdbc:mysql://127.0.0.1:22/mydb?user=myuser&password=PW"                                                                                                                                                                                    | stdProps | "mysql"      | null          | "myuser"      | "127.0.0.1"                                                         | 22    | null                               | "mydb"
 
+    // accept **internally_generated**
+    "jdbc:mysql:loadbalance://**internally_generated**1711497942886**"                                                                                                                                                                          | null     | "mysql"      | "loadbalance" | null          | null                                                                | null  | null                               | null
+
     //https://github.com/awslabs/aws-mysql-jdbc#connection-url-descriptions
     "jdbc:mysql:aws://my.host"                                                                                                                                                                                                                  | null     | "mysql"      | null          | null          | "my.host"                                                           | 3306  | null                               | null
     "jdbc:mysql:aws://my.host:22/mydb"                                                                                                                                                                                                          | null     | "mysql"      | null          | null          | "my.host"                                                           | 22    | null                               | "mydb"