Repository: DataDog/dd-trace-py

Issue:
cassandra: bound/batchstatement aren't normalized properly.
need to handle this case.

```
<BoundStatement query="SELECT * FROM my_table WHERE id = ? LIMIT ?", values=[12356L], consistency=LOCAL_ONE>
```

cc @jhgg
cassandra: bound/batchstatement aren't normalized properly.
need to handle this case.

```
<BoundStatement query="SELECT * FROM my_table WHERE id = ? LIMIT ?", values=[12356L], consistency=LOCAL_ONE>
```

cc @jhgg


Discussion:

Resolving patch:
diff --git a/ddtrace/contrib/cassandra/session.py b/ddtrace/contrib/cassandra/session.py
--- a/ddtrace/contrib/cassandra/session.py
+++ b/ddtrace/contrib/cassandra/session.py
@@ -10,7 +10,6 @@
 from ddtrace.compat import stringify
 from ...util import deep_getattr, deprecated
 from ...ext import net, cassandra as cassx
-from ...ext import AppTypes
 
 
 RESOURCE_MAX_LENGTH = 5000
@@ -46,7 +45,7 @@ def traced_execute(func, instance, args, kwargs):
     query = kwargs.get("kwargs") or args[0]
 
     with tracer.trace("cassandra.query", service=service, span_type=cassx.TYPE) as span:
-        span.resource = _sanitize_query(query)
+        _sanitize_query(span, query)
         span.set_tags(_extract_session_metas(instance))     # FIXME[matt] do once?
         span.set_tags(_extract_cluster_metas(cluster))
         result = None
@@ -58,62 +57,6 @@ def traced_execute(func, instance, args, kwargs):
                 span.set_tags(_extract_result_metas(result))
 
 
-@deprecated(message='Use patching instead (see the docs).', version='0.6.0')
-def get_traced_cassandra(tracer, service=SERVICE, meta=None):
-    return _get_traced_cluster(cassandra.cluster, tracer, service, meta)
-
-
-def _get_traced_cluster(cassandra, tracer, service="cassandra", meta=None):
-    """ Trace synchronous cassandra commands by patching the Session class """
-
-    tracer.set_service_info(
-        service=service,
-        app="cassandra",
-        app_type=AppTypes.db,
-    )
-
-    class TracedSession(cassandra.Session):
-        _dd_tracer = tracer
-        _dd_service = service
-        _dd_tags = meta
-
-        def __init__(self, *args, **kwargs):
-            super(TracedSession, self).__init__(*args, **kwargs)
-
-        def execute(self, query, *args, **options):
-            if not self._dd_tracer:
-                return super(TracedSession, self).execute(query, *args, **options)
-
-            with self._dd_tracer.trace("cassandra.query", service=self._dd_service) as span:
-                query_string = _sanitize_query(query)
-                span.resource = query_string
-                span.span_type = cassx.TYPE
-
-                span.set_tags(_extract_session_metas(self))
-                cluster = getattr(self, "cluster", None)
-                span.set_tags(_extract_cluster_metas(cluster))
-
-                result = None
-                try:
-                    result = super(TracedSession, self).execute(query, *args, **options)
-                    return result
-                finally:
-                    span.set_tags(_extract_result_metas(result))
-
-    class TracedCluster(cassandra.Cluster):
-
-        def connect(self, *args, **kwargs):
-            orig = cassandra.Session
-            cassandra.Session = TracedSession
-            traced_session = super(TracedCluster, self).connect(*args, **kwargs)
-
-            # unpatch the Session class so we don't wrap already traced sessions
-            cassandra.Session = orig
-
-            return traced_session
-
-    return TracedCluster
-
 def _extract_session_metas(session):
     metas = {}
 
@@ -130,10 +73,6 @@ def _extract_cluster_metas(cluster):
         metas[cassx.CLUSTER] = cluster.metadata.cluster_name
     if getattr(cluster, "port", None):
         metas[net.TARGET_PORT] = cluster.port
-    if getattr(cluster, "compression", None):
-        metas[cassx.COMPRESSION] = cluster.compression
-    if getattr(cluster, "cql_version", None):
-        metas[cassx.CQL_VERSION] = cluster.cql_version
 
     return metas
 
@@ -141,17 +80,23 @@ def _extract_result_metas(result):
     metas = {}
     if not result:
         return metas
+
     future = getattr(result, "response_future", None)
+
     if future:
+        # get the host
         host = getattr(future, "coordinator_host", None)
         if host:
             metas[net.TARGET_HOST] = host
+        elif hasattr(future, '_current_host'):
+            address = deep_getattr(future, '_current_host.address')
+            if address:
+                metas[net.TARGET_HOST] = address
+
         query = getattr(future, "query", None)
         if getattr(query, "consistency_level", None):
             metas[cassx.CONSISTENCY_LEVEL] = query.consistency_level
         if getattr(query, "keyspace", None):
-            # Overrides session.keyspace if the query has been prepared against a particular
-            # keyspace
             metas[cassx.KEYSPACE] = query.keyspace.lower()
 
     if hasattr(result, "has_more_pages"):
@@ -166,14 +111,43 @@ def _extract_result_metas(result):
 
     return metas
 
-def _sanitize_query(query):
-    """ Sanitize the query to something ready for the agent receiver
-    - Cast to unicode
-    - truncate if needed
-    """
+def _sanitize_query(span, query):
     # TODO (aaditya): fix this hacky type check. we need it to avoid circular imports
-    if type(query).__name__ in ('SimpleStatement', 'PreparedStatement'):
+    t = type(query).__name__
+
+    resource = None
+    if t in ('SimpleStatement', 'PreparedStatement'):
         # reset query if a string is available
-        query = getattr(query, "query_string", query)
+        resource = getattr(query, "query_string", query)
+    elif t == 'BatchStatement':
+        resource = 'BatchStatement'
+        q = "; ".join(q[1] for q in query._statements_and_parameters[:2])
+        span.set_tag("cassandra.query", q)
+        span.set_metric("cassandra.batch_size", len(query._statements_and_parameters))
+    elif t == 'BoundStatement':
+        ps = getattr(query, 'prepared_statement', None)
+        if ps:
+            resource = getattr(ps, 'query_string', None)
+    elif t == 'str':
+        resource = query
+    else:
+        resource = 'unknown-query-type' # FIXME[matt] what else do to here?
+
+    span.resource = stringify(resource)[:RESOURCE_MAX_LENGTH]
+
+
+#
+# DEPRECATED
+#
+
+@deprecated(message='Use patching instead (see the docs).', version='0.6.0')
+def get_traced_cassandra(tracer, service=SERVICE, meta=None):
+    return _get_traced_cluster(cassandra.cluster, tracer, service, meta)
+
+
+def _get_traced_cluster(cassandra, tracer, service="cassandra", meta=None):
+    """ Trace synchronous cassandra commands by patching the Session class """
+    return cassandra.Cluster
+
+
 
-    return stringify(query)[:RESOURCE_MAX_LENGTH]
diff --git a/ddtrace/ext/cassandra.py b/ddtrace/ext/cassandra.py
--- a/ddtrace/ext/cassandra.py
+++ b/ddtrace/ext/cassandra.py
@@ -8,7 +8,3 @@
 CONSISTENCY_LEVEL = "cassandra.consistency_level"
 PAGINATED = "cassandra.paginated"
 ROW_COUNT = "cassandra.row_count"
-COMPRESSION = "cassandra.compression"
-CONTACT_POINTS = "cassandra.contact_points"
-CQL_VERSION = "cassandra.cql_version"
-


Test patch:
diff --git a/tests/contrib/cassandra/test.py b/tests/contrib/cassandra/test.py
--- a/tests/contrib/cassandra/test.py
+++ b/tests/contrib/cassandra/test.py
@@ -1,10 +1,13 @@
 
 # stdlib
+import logging
 import unittest
 
 # 3p
 from nose.tools import eq_
+from nose.plugins.attrib import attr
 from cassandra.cluster import Cluster
+from cassandra.query import BatchStatement, SimpleStatement
 
 # project
 from tests.contrib.config import CASSANDRA_CONFIG
@@ -15,6 +18,9 @@
 from ddtrace import Pin
 
 
+logging.getLogger('cassandra').setLevel(logging.INFO)
+
+
 class CassandraBase(object):
     """
     Needs a running Cassandra
@@ -37,12 +43,17 @@ def setUp(self):
 
         self.cluster = Cluster(port=CASSANDRA_CONFIG['port'])
         session = self.cluster.connect()
-        session.execute("""CREATE KEYSPACE if not exists test WITH REPLICATION = {
-            'class' : 'SimpleStrategy',
-            'replication_factor': 1
-        }""")
-        session.execute("CREATE TABLE if not exists test.person (name text PRIMARY KEY, age int, description text)")
-        session.execute("""INSERT INTO test.person (name, age, description) VALUES ('Cassandra', 100, 'A cruel mistress')""")
+        sqls = [
+            """CREATE KEYSPACE if not exists test WITH REPLICATION = {
+                'class' : 'SimpleStrategy',
+                'replication_factor': 1
+            }""",
+            "DROP TABLE IF EXISTS test.person",
+            "CREATE TABLE if not exists test.person (name text PRIMARY KEY, age int, description text)",
+            "INSERT INTO test.person (name, age, description) VALUES ('Cassandra', 100, 'A cruel mistress')",
+        ]
+        for sql in sqls:
+            session.execute(sql)
 
     def _assert_result_correct(self, result):
         eq_(len(result.current_rows), 1)
@@ -97,17 +108,38 @@ def test_trace_error(self):
         for k in (errors.ERROR_MSG, errors.ERROR_TYPE, errors.ERROR_STACK):
             assert query.get_tag(k)
 
+    @attr('bound')
+    def test_bound_statement(self):
+        session, writer = self._traced_session()
 
-class TestOldSchool(CassandraBase):
-    """Test Cassandra instrumentation with the legacy interface"""
+        query = "INSERT INTO test.person (name, age, description) VALUES (?, ?, ?)"
+        prepared = session.prepare(query)
+        session.execute(prepared, ("matt", 34, "can"))
 
-    TEST_SERVICE = 'test-cassandra-legacy'
+        prepared = session.prepare(query)
+        bound_stmt = prepared.bind(("leo", 16, "fr"))
+        session.execute(bound_stmt)
 
-    def _traced_session(self):
-        tracer = get_dummy_tracer()
-        tracer_cluster = get_traced_cassandra(tracer, service=self.TEST_SERVICE)
-        session = tracer_cluster(port=CASSANDRA_CONFIG['port']).connect(self.TEST_KEYSPACE)
-        return session, tracer.writer
+        spans = writer.pop()
+        eq_(len(spans), 2)
+        for s in spans:
+            eq_(s.resource, query)
+
+
+    def test_batch_statement(self):
+        session, writer = self._traced_session()
+
+        batch = BatchStatement()
+        batch.add(SimpleStatement("INSERT INTO test.person (name, age, description) VALUES (%s, %s, %s)"), ("Joe", 1, "a"))
+        batch.add(SimpleStatement("INSERT INTO test.person (name, age, description) VALUES (%s, %s, %s)"), ("Jane", 2, "b"))
+        session.execute(batch)
+
+        spans = writer.pop()
+        eq_(len(spans), 1)
+        s = spans[0]
+        eq_(s.resource, 'BatchStatement')
+        eq_(s.get_metric('cassandra.batch_size'), 2)
+        assert 'test.person' in s.get_tag('cassandra.query')
 
 
 class TestCassPatchDefault(CassandraBase):