feat(cassandra4): more attributes

fixes #1298
open-telemetry · Dec 14, 2020 · 021c905 · 021c905
1 parent c800c0d
commit 021c905
Showing 1 changed file with 87 additions and 1 deletion.
diff --git a/...opentelemetry/javaagent/instrumentation/cassandra/v4_0/CassandraDatabaseClientTracer.java b/...opentelemetry/javaagent/instrumentation/cassandra/v4_0/CassandraDatabaseClientTracer.java
@@ -7,7 +7,10 @@
 
 import com.datastax.oss.driver.api.core.CqlIdentifier;
 import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
+import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
 import com.datastax.oss.driver.api.core.cql.ExecutionInfo;
+import com.datastax.oss.driver.api.core.cql.Statement;
 import com.datastax.oss.driver.api.core.metadata.Node;
 import io.opentelemetry.api.trace.Span;
 import io.opentelemetry.context.Context;
@@ -16,8 +19,25 @@
 import io.opentelemetry.javaagent.instrumentation.api.db.DbSystem;
 import java.net.InetSocketAddress;
 import java.net.SocketAddress;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+import javax.annotation.Nullable;
 
 public class CassandraDatabaseClientTracer extends DatabaseClientTracer<CqlSession, String> {
+
+  private static final Pattern tableNameRegex =
+      Pattern.compile(".*(?:FROM|INTO|UPDATE)\\s+([A-Z1-9_]+\\.([A-Z1-9_]+)|([A-Z1-9_]+))",
+          Pattern.CASE_INSENSITIVE);
+
+  // TODO remove when semantic attributes are added
+  private static final String QUERY_PAGE_SIZE = "queryPageSize";
+  private static final String QUERY_CL = "queryCL";
+  private static final String DB_HOST_DC = "dbHostDC";
+  private static final String DB_HOST_ID = "dbHostID";
+  private static final String QUERY_TABLE = "table";
+  private static final String QUERY_IDEMPOTENCE = "idempotence";
+  private static final String QUERY_SPECULATIVE_EXECUTION_COUNT = "speculativeExecutionCount";
+
   private static final CassandraDatabaseClientTracer TRACER = new CassandraDatabaseClientTracer();
 
   public static CassandraDatabaseClientTracer tracer() {
@@ -49,14 +69,80 @@ protected InetSocketAddress peerAddress(CqlSession cqlSession) {
     return null;
   }
 
+  @Override
+  protected Span onConnection(Span span, CqlSession cqlSession) {
+    span = super.onConnection(span, cqlSession);
+    DriverExecutionProfile config = cqlSession.getContext().getConfig()
+        .getDefaultProfile();
+    // may be overwritten by statement, but take the default for now
+    int pageSize = config.getInt(DefaultDriverOption.REQUEST_PAGE_SIZE);
+    if (pageSize > 0) {
+      span.setAttribute(QUERY_PAGE_SIZE, pageSize);
+    }
+    // may be overwritten by statement, but take the default for now
+    span.setAttribute(QUERY_CL,
+        config.getString(DefaultDriverOption.REQUEST_CONSISTENCY));
+    return span;
+  }
+
   public void onResponse(Context context, ExecutionInfo executionInfo) {
+    Span span = Span.fromContext(context);
     Node coordinator = executionInfo.getCoordinator();
     if (coordinator != null) {
       SocketAddress socketAddress = coordinator.getEndPoint().resolve();
       if (socketAddress instanceof InetSocketAddress) {
-        Span span = Span.fromContext(context);
         NetPeerUtils.INSTANCE.setNetPeer(span, ((InetSocketAddress) socketAddress));
       }
+      if (coordinator.getDatacenter() != null) {
+        span.setAttribute(DB_HOST_DC, coordinator.getDatacenter());
+      }
+      if (coordinator.getHostId() != null) {
+        span.setAttribute(DB_HOST_ID, coordinator.getHostId().toString());
+      }
+    }
+    span.setAttribute(QUERY_SPECULATIVE_EXECUTION_COUNT,
+        executionInfo.getSpeculativeExecutionCount());
+
+    Statement<?> statement = executionInfo.getStatement();
+    // override connection default if present
+    if (statement.getConsistencyLevel() != null) {
+      span.setAttribute(QUERY_CL, statement.getConsistencyLevel().name());
+    }
+    // override connection default if present
+    if (statement.getPageSize() > 0) {
+      span.setAttribute(QUERY_PAGE_SIZE, statement.getPageSize());
+    }
+    span.setAttribute(QUERY_IDEMPOTENCE, isIdempotent(statement));
+  }
+
+  @Override
+  protected void onStatement(Span span, String statement) {
+    super.onStatement(span, statement);
+    String table = extractTableNameFromQuery(statement);
+    if (table != null) {
+      span.setAttribute(QUERY_TABLE, table);
+    }
+  }
+
+  @Nullable
+  private static String extractTableNameFromQuery(String query) {
+    String tableName = null;
+    Matcher matcher = tableNameRegex.matcher(query);
+    if (matcher.find()) {
+      if (matcher.group(2) != null) {
+        tableName = matcher.group(2);
+      } else {
+        tableName = matcher.group(1);
+      }
+    }
+    return tableName;
+  }
+
+  private static boolean isIdempotent(Statement<?> statement) {
+    Boolean idempotent = statement.isIdempotent();
+    if (idempotent != null) {
+      return idempotent;
     }
+    return false;
   }
 }