Skip to content

Commit e6257ec

Browse files
amarzialiamarziali
committed
Cassandra: allow extracting keyspace from statement result
1 parent f1d359b commit e6257ec

File tree

9 files changed

+160
-55
lines changed

9 files changed

+160
-55
lines changed

‎dd-java-agent/agent-bootstrap/src/main/java/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecorator.java

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -70,14 +70,7 @@ public String getDbType() {
7070
public AgentSpan onConnection(final AgentSpan span, final CONNECTION connection) {
7171
if (connection != null) {
7272
span.setTag(Tags.DB_USER, dbUser(connection));
73-
final String instanceName = dbInstance(connection);
74-
span.setTag(Tags.DB_INSTANCE, instanceName);
75-
76-
String serviceName = dbClientService(instanceName);
77-
if (null != serviceName) {
78-
span.setServiceName(serviceName);
79-
}
80-
73+
onInstance(span, dbInstance(connection));
8174
CharSequence hostName = dbHostname(connection);
8275
if (hostName != null) {
8376
span.setTag(Tags.PEER_HOSTNAME, hostName);
@@ -90,6 +83,17 @@ public AgentSpan onConnection(final AgentSpan span, final CONNECTION connection)
9083
return span;
9184
}
9285

86+
protected AgentSpan onInstance(final AgentSpan span, final String dbInstance) {
87+
if (dbInstance != null) {
88+
span.setTag(Tags.DB_INSTANCE, dbInstance);
89+
String serviceName = dbClientService(dbInstance);
90+
if (null != serviceName) {
91+
span.setServiceName(serviceName);
92+
}
93+
}
94+
return span;
95+
}
96+
9397
public String dbService(final String dbType, final String instanceName) {
9498
if (instanceName != null && Config.get().isDbClientSplitByInstance()) {
9599
return dbClientService(instanceName);

‎dd-java-agent/instrumentation/datastax-cassandra-3.8/src/test/groovy/CassandraClientTest.groovy

Lines changed: 34 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
22
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
3+
import static datadog.trace.api.config.TraceInstrumentationConfig.CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED
34
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
45

56
import com.datastax.driver.core.Cluster
@@ -57,8 +58,11 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
5758
def "test sync"() {
5859
setup:
5960

60-
Session session = cluster.connect(keyspace)
61+
Session session = keyspace ? cluster.connect(keyspace) : cluster.connect()
6162
injectSysConfig(DB_CLIENT_HOST_SPLIT_BY_INSTANCE, "$renameService")
63+
if (extractFromStatement) {
64+
injectSysConfig(CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED, "true")
65+
}
6266

6367
when:
6468
session.execute(statement)
@@ -71,27 +75,37 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
7175
}
7276
}
7377
trace(1) {
74-
cassandraSpan(it, statement, keyspace, renameService)
78+
cassandraSpan(it, statement, expectedKeySpace, renameService)
7579
}
7680
}
7781

7882
cleanup:
7983
session.close()
8084

8185
where:
82-
statement | keyspace | renameService
83-
"DROP KEYSPACE IF EXISTS sync_test" | null | false
84-
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | true
85-
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )" | "sync_test" | false
86-
"INSERT INTO sync_test.users (id, name) values (uuid(), 'alice')" | "sync_test" | false
87-
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "sync_test" | true
86+
statement | keyspace | expectedKeySpace | renameService | extractFromStatement
87+
"DROP KEYSPACE IF EXISTS sync_test" | null | null | false | true
88+
"DROP KEYSPACE IF EXISTS a_ks" | null | null | false | true
89+
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | true
90+
"CREATE KEYSPACE a_ks WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | false | true
91+
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )" | "sync_test" | "sync_test" | false | true
92+
"INSERT INTO sync_test.users (id, name) values (uuid(), 'alice')" | "sync_test" | "sync_test" | false | true
93+
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "sync_test" | "sync_test" | true | true
94+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "sync_test" | false | true
95+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "sync_test" | true | true
96+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | "sync_test" | false | true
97+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | "sync_test" | true | true
98+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | null | false | false
8899
}
89100

90101
def "test async"() {
91102
setup:
92103

93104
def callbackExecuted = new CountDownLatch(1)
94105
injectSysConfig(DB_CLIENT_HOST_SPLIT_BY_INSTANCE, "$renameService")
106+
if (extractFromStatement) {
107+
injectSysConfig(CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED, "true")
108+
}
95109

96110
when:
97111
Session session = cluster.connect(keyspace)
@@ -117,7 +131,7 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
117131
trace(3) {
118132
sortSpansByStart()
119133
basicSpan(it, "parent")
120-
cassandraSpan(it, statement, keyspace, renameService, span(0))
134+
cassandraSpan(it, statement, expectedKeySpace, renameService, span(0))
121135
basicSpan(it, "callbackListener", span(0))
122136
}
123137
}
@@ -126,12 +140,17 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
126140
session.close()
127141

128142
where:
129-
statement | keyspace | renameService
130-
"DROP KEYSPACE IF EXISTS async_test" | null | false
131-
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | true
132-
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )" | "async_test" | false
133-
"INSERT INTO async_test.users (id, name) values (uuid(), 'alice')" | "async_test" | false
134-
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "async_test" | true
143+
statement | keyspace | expectedKeySpace | renameService | extractFromStatement
144+
"DROP KEYSPACE IF EXISTS async_test" | null | null | false | false
145+
"DROP KEYSPACE IF EXISTS a_ks" | null | null | false | false
146+
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | false
147+
"CREATE KEYSPACE a_ks WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | false
148+
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )" | "async_test" | "async_test" | false | false
149+
"INSERT INTO async_test.users (id, name) values (uuid(), 'alice')" | "async_test" | "async_test" | false | false
150+
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "async_test" | "async_test" | false | false
151+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | null | null | false | false
152+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "a_ks" | false | false
153+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | null | "async_test" | true | true
135154
}
136155

137156
String normalize(String statement){

‎dd-java-agent/instrumentation/datastax-cassandra-3/src/main/java/datadog/trace/instrumentation/datastax/cassandra/CassandraClientDecorator.java

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,12 @@
11
package datadog.trace.instrumentation.datastax.cassandra;
22

3+
import static datadog.trace.bootstrap.instrumentation.api.Tags.DB_INSTANCE;
4+
5+
import com.datastax.driver.core.ColumnDefinitions;
36
import com.datastax.driver.core.Host;
47
import com.datastax.driver.core.ResultSet;
58
import com.datastax.driver.core.Session;
9+
import datadog.trace.api.Config;
610
import datadog.trace.api.cache.DDCache;
711
import datadog.trace.api.cache.DDCaches;
812
import datadog.trace.api.naming.SpanNaming;
@@ -11,6 +15,7 @@
1115
import datadog.trace.bootstrap.instrumentation.api.InternalSpanTypes;
1216
import datadog.trace.bootstrap.instrumentation.api.UTF8BytesString;
1317
import datadog.trace.bootstrap.instrumentation.decorator.DBTypeProcessingDatabaseClientDecorator;
18+
import datadog.trace.util.Strings;
1419
import java.util.function.ToIntFunction;
1520

1621
public class CassandraClientDecorator extends DBTypeProcessingDatabaseClientDecorator<Session> {
@@ -82,6 +87,15 @@ public AgentSpan onResponse(final AgentSpan span, final ResultSet result) {
8287
if (result != null) {
8388
final Host host = result.getExecutionInfo().getQueriedHost();
8489
onPeerConnection(span, host.getSocketAddress());
90+
if (Config.get().isCassandraKeyspaceStatementExtractionEnabled()) {
91+
final ColumnDefinitions defs = result.getColumnDefinitions();
92+
if (defs != null && defs.size() > 0) {
93+
final String keySpace = defs.getKeyspace(0);
94+
if (Strings.isNotBlank(keySpace) && !keySpace.equals(span.getTag(DB_INSTANCE))) {
95+
onInstance(span, keySpace);
96+
}
97+
}
98+
}
8599
}
86100
return span;
87101
}

‎dd-java-agent/instrumentation/datastax-cassandra-3/src/test/groovy/CassandraClientTest.groovy

Lines changed: 34 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import static datadog.trace.agent.test.utils.TraceUtils.basicSpan
22
import static datadog.trace.agent.test.utils.TraceUtils.runUnderTrace
3+
import static datadog.trace.api.config.TraceInstrumentationConfig.CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED
34
import static datadog.trace.api.config.TraceInstrumentationConfig.DB_CLIENT_HOST_SPLIT_BY_INSTANCE
45

56
import com.datastax.driver.core.Cluster
@@ -57,8 +58,11 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
5758
def "test sync"() {
5859
setup:
5960

60-
Session session = cluster.connect(keyspace)
61+
Session session = keyspace ? cluster.connect(keyspace) : cluster.connect()
6162
injectSysConfig(DB_CLIENT_HOST_SPLIT_BY_INSTANCE, "$renameService")
63+
if (extractFromStatement) {
64+
injectSysConfig(CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED, "true")
65+
}
6266

6367
when:
6468
session.execute(statement)
@@ -71,27 +75,37 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
7175
}
7276
}
7377
trace(1) {
74-
cassandraSpan(it, statement, keyspace, renameService)
78+
cassandraSpan(it, statement, expectedKeySpace, renameService)
7579
}
7680
}
7781

7882
cleanup:
7983
session.close()
8084

8185
where:
82-
statement | keyspace | renameService
83-
"DROP KEYSPACE IF EXISTS sync_test" | null | false
84-
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | true
85-
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )" | "sync_test" | false
86-
"INSERT INTO sync_test.users (id, name) values (uuid(), 'alice')" | "sync_test" | false
87-
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "sync_test" | true
86+
statement | keyspace | expectedKeySpace | renameService | extractFromStatement
87+
"DROP KEYSPACE IF EXISTS sync_test" | null | null | false | true
88+
"DROP KEYSPACE IF EXISTS a_ks" | null | null | false | true
89+
"CREATE KEYSPACE sync_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | true
90+
"CREATE KEYSPACE a_ks WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | false | true
91+
"CREATE TABLE sync_test.users ( id UUID PRIMARY KEY, name text )" | "sync_test" | "sync_test" | false | true
92+
"INSERT INTO sync_test.users (id, name) values (uuid(), 'alice')" | "sync_test" | "sync_test" | false | true
93+
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "sync_test" | "sync_test" | true | true
94+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "sync_test" | false | true
95+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "sync_test" | true | true
96+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | "sync_test" | false | true
97+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | "sync_test" | true | true
98+
"SELECT * FROM sync_test.users where name = 'alice' ALLOW FILTERING" | null | null | false | false
8899
}
89100

90101
def "test async"() {
91102
setup:
92103

93104
def callbackExecuted = new CountDownLatch(1)
94105
injectSysConfig(DB_CLIENT_HOST_SPLIT_BY_INSTANCE, "$renameService")
106+
if (extractFromStatement) {
107+
injectSysConfig(CASSANDRA_KEYSPACE_STATEMENT_EXTRACTION_ENABLED, "true")
108+
}
95109

96110
when:
97111
Session session = cluster.connect(keyspace)
@@ -117,7 +131,7 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
117131
trace(3) {
118132
sortSpansByStart()
119133
basicSpan(it, "parent")
120-
cassandraSpan(it, statement, keyspace, renameService, span(0))
134+
cassandraSpan(it, statement, expectedKeySpace, renameService, span(0))
121135
basicSpan(it, "callbackListener", span(0))
122136
}
123137
}
@@ -126,12 +140,17 @@ abstract class CassandraClientTest extends VersionedNamingTestBase {
126140
session.close()
127141

128142
where:
129-
statement | keyspace | renameService
130-
"DROP KEYSPACE IF EXISTS async_test" | null | false
131-
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | true
132-
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )" | "async_test" | false
133-
"INSERT INTO async_test.users (id, name) values (uuid(), 'alice')" | "async_test" | false
134-
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "async_test" | true
143+
statement | keyspace | expectedKeySpace | renameService | extractFromStatement
144+
"DROP KEYSPACE IF EXISTS async_test" | null | null | false | false
145+
"DROP KEYSPACE IF EXISTS a_ks" | null | null | false | false
146+
"CREATE KEYSPACE async_test WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | false
147+
"CREATE KEYSPACE a_ks WITH REPLICATION = {'class':'SimpleStrategy', 'replication_factor':3}" | null | null | true | false
148+
"CREATE TABLE async_test.users ( id UUID PRIMARY KEY, name text )" | "async_test" | "async_test" | false | false
149+
"INSERT INTO async_test.users (id, name) values (uuid(), 'alice')" | "async_test" | "async_test" | false | false
150+
"SELECT * FROM users where name = 'alice' ALLOW FILTERING" | "async_test" | "async_test" | false | false
151+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | null | null | false | false
152+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | "a_ks" | "a_ks" | false | false
153+
"SELECT * FROM async_test.users where name = 'alice' ALLOW FILTERING" | null | "async_test" | true | true
135154
}
136155

137156
String normalize(String statement){

0 commit comments

Comments
 (0)