initial changes to integrate trace propagation

Author: sfc-gh-mkarjikar

pom.xml

@@ -28,6 +28,16 @@
   </properties>
 
   <dependencies>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-api</artifactId>
+      <version>1.42.1</version>
+    </dependency>
+    <dependency>
+      <groupId>io.opentelemetry</groupId>
+      <artifactId>opentelemetry-context</artifactId>
+      <version>1.42.1</version>
+    </dependency>
     <dependency>
       <groupId>org.bouncycastle</groupId>
       <artifactId>bcpkix-jdk18on</artifactId>

public_pom.xml

@@ -4,7 +4,7 @@
 
     <groupId>net.snowflake</groupId>
     <artifactId>snowflake-jdbc</artifactId>
-    <version>1.0-SNAPSHOT</version>
+    <version>1.0-SNAPSHOT-CUSTOM</version>
     <packaging>jar</packaging>
     <name>Snowflake JDBC Driver</name>
     <description>Snowflake JDBC Driver</description>

src/main/java/net/snowflake/client/core/HttpUtil.java

@@ -12,6 +12,11 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Strings;
 import com.microsoft.azure.storage.OperationContext;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
@@ -794,13 +799,24 @@ private static String executeRequestInternal(
     StringWriter writer = null;
     CloseableHttpResponse response = null;
     Stopwatch stopwatch = null;
+    W3CTraceContextPropagator tracePropagator = W3CTraceContextPropagator.getInstance();
 
     if (logger.isDebugEnabled()) {
       stopwatch = new Stopwatch();
       stopwatch.start();
     }
 
     try {
+      // Propagate trace context
+      tracePropagator.inject(
+        Context.current(),
+        httpRequest,
+        (request, name, value) -> request.setHeader(name, value));
+
+      // logger.info("Attempting to inject trace context into request: {}", httpRequest.toString());
+      logger.info("executeRequestInternal HttpUtil context: TraceID: {}, SpanID: {}", 
+        Span.current().getSpanContext().getTraceId(), Span.current().getSpanContext().getSpanId());
+      logger.info("executeRequestInternal HttpUtil thread id: {}", String.valueOf(Thread.currentThread().getId()));
       response =
           RestRequest.execute(
               httpClient,

src/main/java/net/snowflake/client/core/SFStatement.java

@@ -11,6 +11,11 @@
 import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
 
 import com.fasterxml.jackson.databind.JsonNode;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.api.trace.propagation.W3CTraceContextPropagator;
+import io.opentelemetry.context.Context;
+
 import java.sql.SQLException;
 import java.sql.Statement;
 import java.util.ArrayList;
@@ -22,6 +27,8 @@
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
+
+import net.minidev.json.JSONObject;
 import net.snowflake.client.core.BasicEvent.QueryState;
 import net.snowflake.client.core.bind.BindException;
 import net.snowflake.client.core.bind.BindUploader;
@@ -127,6 +134,14 @@ private SFBaseResultSet executeQuery(
       ExecTimeTelemetryData execTimeData)
       throws SQLException, SFException {
     sanityCheckQuery(sql);
+    
+    // Check if the TelemetryService is holding the current context. If NOT, then update the
+    // TelemteryService context by calling updateContext.
+    // Use this code: TelemetryService.getInstance().updateContext(session.getSnowflakeConnectionString());
+    logger.info("executeQuery(with params) SFStatement thread id: {}", String.valueOf(Thread.currentThread().getId()));
+    logger.info("executeQuery(with params) SFStatement context: TraceID = {}, SpanID = {}", 
+        Span.current().getSpanContext().getTraceId(),
+        Span.current().getSpanContext().getSpanId());
 
     String trimmedSql = sql.trim();
 

src/main/java/net/snowflake/client/jdbc/SnowflakeConnectionV1.java

@@ -8,6 +8,10 @@
 import static net.snowflake.client.jdbc.ErrorCode.INVALID_CONNECT_STRING;
 
 import com.google.common.base.Strings;
+
+import io.opentelemetry.api.trace.Span;
+import io.opentelemetry.context.Context;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.sql.Array;
@@ -196,6 +200,9 @@ private void executeImmediate(String stmtText) throws SQLException {
   @Override
   public Statement createStatement() throws SQLException {
     raiseSQLExceptionIfConnectionIsClosed();
+    logger.info("createStatement SnowflakeConnectionV1 thread id: {}", String.valueOf(Thread.currentThread().getId()));
+    logger.info("createStatement SnowflakeStatementV1 context: TraceID: {}, SpanID: {}", 
+        Span.current().getSpanContext().getTraceId(), Span.current().getSpanContext().getSpanId());
     Statement stmt = createStatement(ResultSet.TYPE_FORWARD_ONLY, ResultSet.CONCUR_READ_ONLY);
     openStatements.add(stmt);
     return stmt;

src/main/java/net/snowflake/client/jdbc/SnowflakeStatementV1.java

@@ -20,6 +20,8 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
+
+import io.opentelemetry.api.trace.Span;
 import net.snowflake.client.core.CancellationReason;
 import net.snowflake.client.core.ExecTimeTelemetryData;
 import net.snowflake.client.core.ParameterBindingDTO;
@@ -146,9 +148,11 @@ protected void raiseSQLExceptionIfStatementIsClosed() throws SQLException {
    */
   @Override
   public ResultSet executeQuery(String sql) throws SQLException {
+    logger.info("executeQuery SnowflakeStatementV1 thread id: {}", String.valueOf(Thread.currentThread().getId()));
+      logger.info("executeQuery SnowflakeStatementV1 context: TraceID: {}, SpanID: {}", 
+          Span.current().getSpanContext().getTraceId(), Span.current().getSpanContext().getSpanId());
     ExecTimeTelemetryData execTimeData =
         new ExecTimeTelemetryData("ResultSet Statement.executeQuery(String)", this.batchID);
-
     raiseSQLExceptionIfStatementIsClosed();
     ResultSet rs = executeQueryInternal(sql, false, null, execTimeData);
     execTimeData.setQueryEnd();