add proxy test for trace exporter

Author: keshavdandeva

java-bigquery-jdbc/src/test/java/com/google/cloud/bigquery/jdbc/it/ITOpenTelemetryTest.java

@@ -23,7 +23,9 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import com.google.api.gax.core.FixedCredentialsProvider;
+import com.google.api.gax.grpc.InstantiatingGrpcChannelProvider;
 import com.google.api.gax.paging.Page;
+import com.google.api.gax.rpc.TransportChannelProvider;
 import com.google.auth.oauth2.GoogleCredentials;
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.bigquery.jdbc.BigQueryConnection;
@@ -36,7 +38,18 @@
 import com.google.devtools.cloudtrace.v1.Trace;
 import com.google.devtools.cloudtrace.v1.TraceSpan;
 import com.google.gson.JsonObject;
+import io.grpc.HttpConnectProxiedSocketAddress;
+import io.grpc.ProxiedSocketAddress;
+import io.grpc.ProxyDetector;
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.SpanContext;
+import io.opentelemetry.api.trace.TraceFlags;
+import io.opentelemetry.api.trace.TraceState;
+import io.opentelemetry.context.Scope;
+import io.opentelemetry.sdk.trace.IdGenerator;
 import java.io.File;
+import java.net.InetSocketAddress;
+import java.net.SocketAddress;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.sql.Connection;
@@ -45,6 +58,7 @@
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.List;
+import org.junit.jupiter.api.Tag;
 import org.junit.jupiter.api.Test;
 
 public class ITOpenTelemetryTest extends ITBase {
@@ -53,6 +67,7 @@ public class ITOpenTelemetryTest extends ITBase {
   private static final String CONNECTION_URL = connectionUrl;
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withOpenTelemetryGcpExporter() throws Exception {
 
     // Step 1: Connect with GCP Exporters enabled via DataSource
@@ -124,8 +139,8 @@ public void testExecute_withOpenTelemetryGcpExporter() throws Exception {
   }
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withErrorCorrelation() throws Exception {
-
     // Step 1: Connect with GCP Exporters enabled via DataSource
     DataSource ds = DataSource.fromUrl(CONNECTION_URL);
     ds.setEnableGcpTraceExporter(true);
@@ -168,6 +183,7 @@ public void testExecute_withErrorCorrelation() throws Exception {
   }
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withCustomCredentialsJson() throws Exception {
     JsonObject authJson = getAuthJson();
     DataSource ds = DataSource.fromUrl(CONNECTION_URL);
@@ -179,6 +195,7 @@ public void testExecute_withCustomCredentialsJson() throws Exception {
   }
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withCustomCredentialsFilePath() throws Exception {
     JsonObject authJson = getAuthJson();
     File tempFile = File.createTempFile("auth", ".json");
@@ -194,6 +211,7 @@ public void testExecute_withCustomCredentialsFilePath() throws Exception {
   }
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withHttpProtocol() throws Exception {
     JsonObject authJson = getAuthJson();
     System.setProperty("otel.exporter.otlp.protocol", "http/protobuf");
@@ -211,6 +229,7 @@ public void testExecute_withHttpProtocol() throws Exception {
   }
 
   @Test
+  @Tag("known_issue")
   public void testExecute_withGrpcProtocol() throws Exception {
     JsonObject authJson = getAuthJson();
     System.setProperty("otel.exporter.otlp.protocol", "grpc");
@@ -227,6 +246,39 @@ public void testExecute_withGrpcProtocol() throws Exception {
     }
   }
 
+  @Test
+  public void testExecute_withHttpProtocol_andDirectTraceVerification() throws Exception {
+    JsonObject authJson = getAuthJson();
+    System.setProperty("otel.exporter.otlp.protocol", "http/protobuf");
+
+    try {
+      DataSource ds = DataSource.fromUrl(CONNECTION_URL);
+      ds.setEnableGcpTraceExporter(true);
+      ds.setGcpTelemetryProjectId(PROJECT_ID);
+      ds.setGcpTelemetryCredentials(authJson.toString());
+
+      String traceId = IdGenerator.random().generateTraceId();
+      String spanId = IdGenerator.random().generateSpanId();
+      SpanContext parentContext =
+          SpanContext.create(traceId, spanId, TraceFlags.getSampled(), TraceState.getDefault());
+
+      try (Scope scope = Span.wrap(parentContext).makeCurrent()) {
+        try (Connection connection = ds.getConnection();
+            Statement statement = connection.createStatement()) {
+          String query = "SELECT 1;";
+          try (ResultSet rs = statement.executeQuery(query)) {
+            assertTrue(rs.next());
+          }
+        }
+      }
+
+      Trace trace = verifyAndFetchTrace(traceId);
+      assertNotNull(trace, "Trace must be found in Cloud Trace API: " + traceId);
+    } finally {
+      System.clearProperty("otel.exporter.otlp.protocol");
+    }
+  }
+
   private void verifyTraceDelivery(DataSource ds) throws Exception {
     ds.setEnableGcpLogExporter(true);
     ds.setLogLevel("5");
@@ -293,18 +345,42 @@ private Trace verifyAndFetchTrace(String traceId) throws Exception {
 
     GoogleCredentials credentials = getCredentials();
 
-    TraceServiceSettings settings =
+    TraceServiceSettings.Builder settingsBuilder =
         TraceServiceSettings.newBuilder()
-            .setCredentialsProvider(FixedCredentialsProvider.create(credentials))
-            .build();
+            .setCredentialsProvider(FixedCredentialsProvider.create(credentials));
+
+    DataSource ds = DataSource.fromUrl(CONNECTION_URL);
+    String proxyHost = ds.getProxyHost();
+    String proxyPortStr = ds.getProxyPort();
+    if (proxyHost != null && proxyPortStr != null) {
+      settingsBuilder.setTransportChannelProvider(
+          createProxyTransportChannelProvider(proxyHost, Integer.parseInt(proxyPortStr)));
+    }
 
-    try (TraceServiceClient traceClient = TraceServiceClient.create(settings)) {
+    try (TraceServiceClient traceClient = TraceServiceClient.create(settingsBuilder.build())) {
       Trace trace = fetchTraceWithRetry(traceClient, PROJECT_ID, hexTraceId);
       assertNotNull(trace, "Trace must be found in Cloud Trace API: " + hexTraceId);
       return trace;
     }
   }
 
+  private TransportChannelProvider createProxyTransportChannelProvider(String host, int port) {
+    return InstantiatingGrpcChannelProvider.newBuilder()
+        .setChannelConfigurator(
+            builder ->
+                builder.proxyDetector(
+                    new ProxyDetector() {
+                      @Override
+                      public ProxiedSocketAddress proxyFor(SocketAddress socketAddress) {
+                        return HttpConnectProxiedSocketAddress.newBuilder()
+                            .setProxyAddress(new InetSocketAddress(host, port))
+                            .setTargetAddress((InetSocketAddress) socketAddress)
+                            .build();
+                      }
+                    }))
+        .build();
+  }
+
   private <T> T pollWithRetry(java.util.concurrent.Callable<T> task) throws InterruptedException {
     int attempts = 0;
     int maxAttempts = 10;