[reference] eval: blind regeneration of PostgreSQL JDBC driver (toolkit output)

dd-java-agent/instrumentation/postgresql/postgresql-42.0/build.gradle

@@ -0,0 +1,26 @@
+muzzle {
+  pass {
+    group = "org.postgresql"
+    module = "postgresql"
+    versions = "[42.0.0,)"
+  }
+}
+
+apply from: "$rootDir/gradle/java.gradle"
+
+addTestSuiteForDir('latestDepTest', 'test')
+
+dependencies {
+  compileOnly group: 'org.postgresql', name: 'postgresql', version: '42.0.0'
+
+  testImplementation group: 'org.postgresql', name: 'postgresql', version: '42.0.0'
+  testImplementation group: 'org.testcontainers', name: 'postgresql', version: libs.versions.testcontainers.get()
+
+  latestDepTestImplementation group: 'org.postgresql', name: 'postgresql', version: '42.+'
+}
+
+tasks.withType(Test).configureEach {
+  usesService(testcontainersLimit)
+  environment 'TESTCONTAINERS_RYUK_DISABLED', 'true'
+  jvmArgs '-Dtestcontainers.reuse.enable=false'
+}

dd-java-agent/instrumentation/postgresql/postgresql-42.0/src/main/java/datadog/trace/instrumentation/postgresql/PgPreparedStatementAdvice.java

@@ -0,0 +1,97 @@
+package datadog.trace.instrumentation.postgresql;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.DECORATE;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.JAVA_POSTGRESQL;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.POSTGRESQL_QUERY;
+
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
+import java.sql.Statement;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.asm.Advice;
+
+public class PgPreparedStatementAdvice {
+
+  public static final Map<Statement, String> PREPARED_SQL =
+      Collections.synchronizedMap(new WeakHashMap<Statement, String>());
+
+  public static final class ConstructorAdvice {
+
+    @Advice.OnMethodExit(suppress = Throwable.class)
+    public static void onExit(
+        @Advice.This final Statement statement, @Advice.Argument(1) final Object query) {
+      // In PostgreSQL JDBC 42.0+, argument[1] is CachedQuery, use toString() to get SQL
+      if (query != null) {
+        PREPARED_SQL.put(statement, query.toString());
+      }
+    }
+  }
+
+  public static final class ExecuteAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(@Advice.This final Statement statement) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class);
+      if (callDepth > 0) {
+        return null;
+      }
+
+      final AgentSpan span = startSpan(JAVA_POSTGRESQL.toString(), POSTGRESQL_QUERY);
+      DECORATE.afterStart(span);
+
+      DBInfo dbInfo = InstrumentationContext.get(Statement.class, DBInfo.class).get(statement);
+      if (dbInfo == null) {
+        dbInfo = PgStatementAdvice.ExecuteQueryAdvice.extractDbInfo(statement);
+        if (dbInfo != null) {
+          InstrumentationContext.get(Statement.class, DBInfo.class).put(statement, dbInfo);
+        }
+      }
+      if (dbInfo != null) {
+        DECORATE.onConnection(span, dbInfo);
+        if (dbInfo.getPort() != null) {
+          DECORATE.setPeerPort(span, dbInfo.getPort());
+        }
+      }
+
+      final String sql = PREPARED_SQL.get(statement);
+      if (sql != null) {
+        final DBQueryInfo dbQueryInfo = DBQueryInfo.ofPreparedStatement(sql);
+        DECORATE.onStatement(span, dbQueryInfo.getSql());
+        span.setTag(Tags.DB_OPERATION, dbQueryInfo.getOperation());
+
+        // DBM: inject SQL comment with trace context for prepared statements
+        // For prepared statements, we inject via modifying the stored SQL that will
+        // be sent to the server. The comment is only added to spans' context,
+        // the actual SQL modification for prepared statements happens at connection level.
+        PostgreSQLSQLCommenter.inject(sql, span, dbInfo);
+      }
+
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      final AgentSpan span = scope.span();
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+      }
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+      CallDepthThreadLocalMap.reset(Statement.class);
+    }
+  }
+}

dd-java-agent/instrumentation/postgresql/postgresql-42.0/src/main/java/datadog/trace/instrumentation/postgresql/PgPreparedStatementInstrumentation.java

@@ -0,0 +1,37 @@
+package datadog.trace.instrumentation.postgresql;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isConstructor;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+
+public final class PgPreparedStatementInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "org.postgresql.jdbc.PgPreparedStatement";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isConstructor(),
+        "datadog.trace.instrumentation.postgresql.PgPreparedStatementAdvice$ConstructorAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("executeQuery")).and(takesArguments(0)),
+        "datadog.trace.instrumentation.postgresql.PgPreparedStatementAdvice$ExecuteAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("executeUpdate")).and(takesArguments(0)),
+        "datadog.trace.instrumentation.postgresql.PgPreparedStatementAdvice$ExecuteAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("execute")).and(takesArguments(0)),
+        "datadog.trace.instrumentation.postgresql.PgPreparedStatementAdvice$ExecuteAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("executeBatch")).and(takesArguments(0)),
+        "datadog.trace.instrumentation.postgresql.PgPreparedStatementAdvice$ExecuteAdvice");
+  }
+}

dd-java-agent/instrumentation/postgresql/postgresql-42.0/src/main/java/datadog/trace/instrumentation/postgresql/PgStatementAdvice.java

@@ -0,0 +1,165 @@
+package datadog.trace.instrumentation.postgresql;
+
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.activateSpan;
+import static datadog.trace.bootstrap.instrumentation.api.AgentTracer.startSpan;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.DECORATE;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.JAVA_POSTGRESQL;
+import static datadog.trace.instrumentation.postgresql.PostgreSQLDecorator.POSTGRESQL_QUERY;
+
+import datadog.trace.bootstrap.CallDepthThreadLocalMap;
+import datadog.trace.bootstrap.InstrumentationContext;
+import datadog.trace.bootstrap.instrumentation.api.AgentScope;
+import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import datadog.trace.bootstrap.instrumentation.api.Tags;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBInfo;
+import datadog.trace.bootstrap.instrumentation.jdbc.DBQueryInfo;
+import datadog.trace.bootstrap.instrumentation.jdbc.JDBCConnectionUrlParser;
+import java.sql.Connection;
+import java.sql.DatabaseMetaData;
+import java.sql.Statement;
+import java.util.Collections;
+import java.util.Map;
+import java.util.WeakHashMap;
+import net.bytebuddy.asm.Advice;
+
+public class PgStatementAdvice {
+
+  public static final Map<Statement, String> BATCH_SQL =
+      Collections.synchronizedMap(new WeakHashMap<Statement, String>());
+
+  public static final class ExecuteQueryAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(
+        @Advice.This final Statement statement,
+        @Advice.Argument(value = 0, readOnly = false) String sql) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class);
+      if (callDepth > 0) {
+        return null;
+      }
+
+      final AgentSpan span = startSpan(JAVA_POSTGRESQL.toString(), POSTGRESQL_QUERY);
+      DECORATE.afterStart(span);
+
+      DBInfo dbInfo = InstrumentationContext.get(Statement.class, DBInfo.class).get(statement);
+      if (dbInfo == null) {
+        dbInfo = extractDbInfo(statement);
+        if (dbInfo != null) {
+          InstrumentationContext.get(Statement.class, DBInfo.class).put(statement, dbInfo);
+        }
+      }
+      if (dbInfo != null) {
+        DECORATE.onConnection(span, dbInfo);
+        if (dbInfo.getPort() != null) {
+          DECORATE.setPeerPort(span, dbInfo.getPort());
+        }
+      }
+
+      final DBQueryInfo dbQueryInfo = DBQueryInfo.ofStatement(sql);
+      DECORATE.onStatement(span, dbQueryInfo.getSql());
+      span.setTag(Tags.DB_OPERATION, dbQueryInfo.getOperation());
+
+      // DBM: inject SQL comment with trace context
+      sql = PostgreSQLSQLCommenter.inject(sql, span, dbInfo);
+
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      final AgentSpan span = scope.span();
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+      }
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+      CallDepthThreadLocalMap.reset(Statement.class);
+    }
+
+    public static DBInfo extractDbInfo(final Statement statement) {
+      try {
+        final Connection connection = statement.getConnection();
+        if (connection != null) {
+          final DatabaseMetaData metaData = connection.getMetaData();
+          final String url = metaData.getURL();
+          final String user = metaData.getUserName();
+          DBInfo dbInfo = JDBCConnectionUrlParser.parse(url, null);
+          if (user != null && !user.isEmpty()) {
+            dbInfo = dbInfo.toBuilder().user(user).build();
+          }
+          return dbInfo;
+        }
+      } catch (final Exception ignored) {
+        // Unable to extract connection info
+      }
+      return null;
+    }
+  }
+
+  public static final class AddBatchAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static void onEnter(
+        @Advice.This final Statement statement, @Advice.Argument(0) final String sql) {
+      BATCH_SQL.put(statement, sql);
+    }
+  }
+
+  public static final class ExecuteBatchAdvice {
+
+    @Advice.OnMethodEnter(suppress = Throwable.class)
+    public static AgentScope onEnter(@Advice.This final Statement statement) {
+      final int callDepth = CallDepthThreadLocalMap.incrementCallDepth(Statement.class);
+      if (callDepth > 0) {
+        return null;
+      }
+
+      final AgentSpan span = startSpan(JAVA_POSTGRESQL.toString(), POSTGRESQL_QUERY);
+      DECORATE.afterStart(span);
+
+      DBInfo dbInfo = InstrumentationContext.get(Statement.class, DBInfo.class).get(statement);
+      if (dbInfo == null) {
+        dbInfo = ExecuteQueryAdvice.extractDbInfo(statement);
+        if (dbInfo != null) {
+          InstrumentationContext.get(Statement.class, DBInfo.class).put(statement, dbInfo);
+        }
+      }
+      if (dbInfo != null) {
+        DECORATE.onConnection(span, dbInfo);
+        if (dbInfo.getPort() != null) {
+          DECORATE.setPeerPort(span, dbInfo.getPort());
+        }
+      }
+
+      final String batchSql = BATCH_SQL.get(statement);
+      if (batchSql != null) {
+        final DBQueryInfo dbQueryInfo = DBQueryInfo.ofStatement(batchSql);
+        DECORATE.onStatement(span, dbQueryInfo.getSql());
+        span.setTag(Tags.DB_OPERATION, dbQueryInfo.getOperation());
+      }
+
+      return activateSpan(span);
+    }
+
+    @Advice.OnMethodExit(onThrowable = Throwable.class, suppress = Throwable.class)
+    public static void onExit(
+        @Advice.Enter final AgentScope scope, @Advice.Thrown final Throwable throwable) {
+      if (scope == null) {
+        return;
+      }
+      final AgentSpan span = scope.span();
+      if (throwable != null) {
+        DECORATE.onError(span, throwable);
+      }
+      DECORATE.beforeFinish(span);
+      scope.close();
+      span.finish();
+      CallDepthThreadLocalMap.reset(Statement.class);
+    }
+  }
+}

dd-java-agent/instrumentation/postgresql/postgresql-42.0/src/main/java/datadog/trace/instrumentation/postgresql/PgStatementInstrumentation.java

@@ -0,0 +1,53 @@
+package datadog.trace.instrumentation.postgresql;
+
+import static datadog.trace.agent.tooling.bytebuddy.matcher.NameMatchers.named;
+import static net.bytebuddy.matcher.ElementMatchers.isMethod;
+import static net.bytebuddy.matcher.ElementMatchers.isPublic;
+import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
+import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
+
+import datadog.trace.agent.tooling.Instrumenter;
+
+public final class PgStatementInstrumentation
+    implements Instrumenter.ForSingleType, Instrumenter.HasMethodAdvice {
+
+  @Override
+  public String instrumentedType() {
+    return "org.postgresql.jdbc.PgStatement";
+  }
+
+  @Override
+  public void methodAdvice(MethodTransformer transformer) {
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("executeQuery"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class)),
+        "datadog.trace.instrumentation.postgresql.PgStatementAdvice$ExecuteQueryAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("executeUpdate"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class)),
+        "datadog.trace.instrumentation.postgresql.PgStatementAdvice$ExecuteQueryAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("execute"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class)),
+        "datadog.trace.instrumentation.postgresql.PgStatementAdvice$ExecuteQueryAdvice");
+    transformer.applyAdvice(
+        isMethod()
+            .and(isPublic())
+            .and(named("addBatch"))
+            .and(takesArguments(1))
+            .and(takesArgument(0, String.class)),
+        "datadog.trace.instrumentation.postgresql.PgStatementAdvice$AddBatchAdvice");
+    transformer.applyAdvice(
+        isMethod().and(isPublic()).and(named("executeBatch")).and(takesArguments(0)),
+        "datadog.trace.instrumentation.postgresql.PgStatementAdvice$ExecuteBatchAdvice");
+  }
+}