Fix for MySQL multi-host internally_generated connection URLs

DataDog · Mar 28, 2024 · 0613d39 · 0613d39
1 parent 6167d5e
commit 0613d39
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 1 deletion.
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;

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
@@ -91,8 +91,22 @@ public static void addDBInfo(
       } catch (Throwable ignored) {
         // use original values
       }
+      Connection connWithContext = connection;
+      if (connectionUrl.startsWith("**internally_generated**", connectionUrl.indexOf("://") + 3)) {
+        // Treat MySql multi-host connections specifically.
+        // Use the original connection URL for parsing.
+        connectionUrl = url;
+        try {
+          // Find an underlying connection object to which to attach the DBInfo. Otherwise, it may
+          // be attached to a wrapper or a proxy connection object that is different from the one
+          // available in other parts of the instrumentation, resulting in a loss of context.
+          connWithContext = connection.unwrap(Connection.class);
+        } catch (Throwable t) {
+          // 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"