Skip to content

Commit 2b72cb5

Browse files
amarzialiamarziali
authored
Cassandra: allow extracting keyspace from statement result (#8239)
* Cassandra: allow extracting keyspace from statement result * testing * review
1 parent 693c605 commit 2b72cb5

File tree

10 files changed

+166
-56
lines changed

10 files changed

+166
-56
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/agent-bootstrap/src/test/groovy/datadog/trace/bootstrap/instrumentation/decorator/DatabaseClientDecoratorTest.groovy‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,9 @@ class DatabaseClientDecoratorTest extends ClientDecoratorTest {
4747
then:
4848
if (session) {
4949
1 * span.setTag(Tags.DB_USER, session.user)
50-
1 * span.setTag(Tags.DB_INSTANCE, session.instance)
50+
if (session.instance != null) {
51+
1 * span.setTag(Tags.DB_INSTANCE, session.instance)
52+
}
5153
if (session.hostname != null) {
5254
1 * span.setTag(Tags.PEER_HOSTNAME, session.hostname)
5355
}

‎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: 17 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,18 @@ 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+
try {
91+
if (Config.get().isCassandraKeyspaceStatementExtractionEnabled()) {
92+
final ColumnDefinitions defs = result.getColumnDefinitions();
93+
if (defs != null && defs.size() > 0) {
94+
final String keySpace = defs.getKeyspace(0);
95+
if (Strings.isNotBlank(keySpace) && !keySpace.equals(span.getTag(DB_INSTANCE))) {
96+
onInstance(span, keySpace);
97+
}
98+
}
99+
}
100+
} catch (final Throwable ignored) {
101+
}
85102
}
86103
return span;
87104
}

‎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)