Move DD code into its own directory (#6)

reference/ddtrace/contrib/cassandra/__init__.py

@@ -0,0 +1,35 @@
+"""Instrument Cassandra to report Cassandra queries.
+
+``patch_all`` will automatically patch your Cluster instance to make it work.
+::
+
+    from ddtrace import Pin, patch
+    from cassandra.cluster import Cluster
+
+    # If not patched yet, you can patch cassandra specifically
+    patch(cassandra=True)
+
+    # This will report spans with the default instrumentation
+    cluster = Cluster(contact_points=["127.0.0.1"], port=9042)
+    session = cluster.connect("my_keyspace")
+    # Example of instrumented query
+    session.execute("select id from my_table limit 10;")
+
+    # Use a pin to specify metadata related to this cluster
+    cluster = Cluster(contact_points=['10.1.1.3', '10.1.1.4', '10.1.1.5'], port=9042)
+    Pin.override(cluster, service='cassandra-backend')
+    session = cluster.connect("my_keyspace")
+    session.execute("select id from my_table limit 10;")
+"""
+from ...utils.importlib import require_modules
+
+
+required_modules = ['cassandra.cluster']
+
+with require_modules(required_modules) as missing_modules:
+    if not missing_modules:
+        from .session import get_traced_cassandra, patch
+        __all__ = [
+            'get_traced_cassandra',
+            'patch',
+        ]

reference/ddtrace/contrib/cassandra/patch.py

@@ -0,0 +1,3 @@
+from .session import patch, unpatch
+
+__all__ = ['patch', 'unpatch']

reference/ddtrace/contrib/cassandra/session.py

@@ -0,0 +1,297 @@
+"""
+Trace queries along a session to a cassandra cluster
+"""
+import sys
+
+# 3p
+import cassandra.cluster
+
+# project
+from ...compat import stringify
+from ...constants import ANALYTICS_SAMPLE_RATE_KEY
+from ...ext import SpanTypes, net, cassandra as cassx, errors
+from ...internal.logger import get_logger
+from ...pin import Pin
+from ...settings import config
+from ...utils.deprecation import deprecated
+from ...utils.formats import deep_getattr
+from ...vendor import wrapt
+
+log = get_logger(__name__)
+
+RESOURCE_MAX_LENGTH = 5000
+SERVICE = 'cassandra'
+CURRENT_SPAN = '_ddtrace_current_span'
+PAGE_NUMBER = '_ddtrace_page_number'
+
+# Original connect connect function
+_connect = cassandra.cluster.Cluster.connect
+
+
+def patch():
+    """ patch will add tracing to the cassandra library. """
+    setattr(cassandra.cluster.Cluster, 'connect',
+            wrapt.FunctionWrapper(_connect, traced_connect))
+    Pin(service=SERVICE, app=SERVICE).onto(cassandra.cluster.Cluster)
+
+
+def unpatch():
+    cassandra.cluster.Cluster.connect = _connect
+
+
+def traced_connect(func, instance, args, kwargs):
+    session = func(*args, **kwargs)
+    if not isinstance(session.execute, wrapt.FunctionWrapper):
+        # FIXME[matt] this should probably be private.
+        setattr(session, 'execute_async', wrapt.FunctionWrapper(session.execute_async, traced_execute_async))
+    return session
+
+
+def _close_span_on_success(result, future):
+    span = getattr(future, CURRENT_SPAN, None)
+    if not span:
+        log.debug('traced_set_final_result was not able to get the current span from the ResponseFuture')
+        return
+    try:
+        span.set_tags(_extract_result_metas(cassandra.cluster.ResultSet(future, result)))
+    except Exception:
+        log.debug('an exception occured while setting tags', exc_info=True)
+    finally:
+        span.finish()
+        delattr(future, CURRENT_SPAN)
+
+
+def traced_set_final_result(func, instance, args, kwargs):
+    result = args[0]
+    _close_span_on_success(result, instance)
+    return func(*args, **kwargs)
+
+
+def _close_span_on_error(exc, future):
+    span = getattr(future, CURRENT_SPAN, None)
+    if not span:
+        log.debug('traced_set_final_exception was not able to get the current span from the ResponseFuture')
+        return
+    try:
+        # handling the exception manually because we
+        # don't have an ongoing exception here
+        span.error = 1
+        span.set_tag(errors.ERROR_MSG, exc.args[0])
+        span.set_tag(errors.ERROR_TYPE, exc.__class__.__name__)
+    except Exception:
+        log.debug('traced_set_final_exception was not able to set the error, failed with error', exc_info=True)
+    finally:
+        span.finish()
+        delattr(future, CURRENT_SPAN)
+
+
+def traced_set_final_exception(func, instance, args, kwargs):
+    exc = args[0]
+    _close_span_on_error(exc, instance)
+    return func(*args, **kwargs)
+
+
+def traced_start_fetching_next_page(func, instance, args, kwargs):
+    has_more_pages = getattr(instance, 'has_more_pages', True)
+    if not has_more_pages:
+        return func(*args, **kwargs)
+    session = getattr(instance, 'session', None)
+    cluster = getattr(session, 'cluster', None)
+    pin = Pin.get_from(cluster)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    # In case the current span is not finished we make sure to finish it
+    old_span = getattr(instance, CURRENT_SPAN, None)
+    if old_span:
+        log.debug('previous span was not finished before fetching next page')
+        old_span.finish()
+
+    query = getattr(instance, 'query', None)
+
+    span = _start_span_and_set_tags(pin, query, session, cluster)
+
+    page_number = getattr(instance, PAGE_NUMBER, 1) + 1
+    setattr(instance, PAGE_NUMBER, page_number)
+    setattr(instance, CURRENT_SPAN, span)
+    try:
+        return func(*args, **kwargs)
+    except Exception:
+        with span:
+            span.set_exc_info(*sys.exc_info())
+        raise
+
+
+def traced_execute_async(func, instance, args, kwargs):
+    cluster = getattr(instance, 'cluster', None)
+    pin = Pin.get_from(cluster)
+    if not pin or not pin.enabled():
+        return func(*args, **kwargs)
+
+    query = kwargs.get('query') or args[0]
+
+    span = _start_span_and_set_tags(pin, query, instance, cluster)
+
+    try:
+        result = func(*args, **kwargs)
+        setattr(result, CURRENT_SPAN, span)
+        setattr(result, PAGE_NUMBER, 1)
+        setattr(
+            result,
+            '_set_final_result',
+            wrapt.FunctionWrapper(
+                result._set_final_result,
+                traced_set_final_result
+            )
+        )
+        setattr(
+            result,
+            '_set_final_exception',
+            wrapt.FunctionWrapper(
+                result._set_final_exception,
+                traced_set_final_exception
+            )
+        )
+        setattr(
+            result,
+            'start_fetching_next_page',
+            wrapt.FunctionWrapper(
+                result.start_fetching_next_page,
+                traced_start_fetching_next_page
+            )
+        )
+        # Since we cannot be sure that the previous methods were overwritten
+        # before the call ended, we add callbacks that will be run
+        # synchronously if the call already returned and we remove them right
+        # after.
+        result.add_callbacks(
+            _close_span_on_success,
+            _close_span_on_error,
+            callback_args=(result,),
+            errback_args=(result,)
+        )
+        result.clear_callbacks()
+        return result
+    except Exception:
+        with span:
+            span.set_exc_info(*sys.exc_info())
+        raise
+
+
+def _start_span_and_set_tags(pin, query, session, cluster):
+    service = pin.service
+    tracer = pin.tracer
+    span = tracer.trace('cassandra.query', service=service, span_type=SpanTypes.CASSANDRA)
+    _sanitize_query(span, query)
+    span.set_tags(_extract_session_metas(session))     # FIXME[matt] do once?
+    span.set_tags(_extract_cluster_metas(cluster))
+    # set analytics sample rate if enabled
+    span.set_tag(
+        ANALYTICS_SAMPLE_RATE_KEY,
+        config.cassandra.get_analytics_sample_rate()
+    )
+    return span
+
+
+def _extract_session_metas(session):
+    metas = {}
+
+    if getattr(session, 'keyspace', None):
+        # FIXME the keyspace can be overridden explicitly in the query itself
+        # e.g. 'select * from trace.hash_to_resource'
+        metas[cassx.KEYSPACE] = session.keyspace.lower()
+
+    return metas
+
+
+def _extract_cluster_metas(cluster):
+    metas = {}
+    if deep_getattr(cluster, 'metadata.cluster_name'):
+        metas[cassx.CLUSTER] = cluster.metadata.cluster_name
+    if getattr(cluster, 'port', None):
+        metas[net.TARGET_PORT] = cluster.port
+
+    return metas
+
+
+def _extract_result_metas(result):
+    metas = {}
+    if result is None:
+        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):
+            metas[cassx.KEYSPACE] = query.keyspace.lower()
+
+        page_number = getattr(future, PAGE_NUMBER, 1)
+        has_more_pages = getattr(future, 'has_more_pages')
+        is_paginated = has_more_pages or page_number > 1
+        metas[cassx.PAGINATED] = is_paginated
+        if is_paginated:
+            metas[cassx.PAGE_NUMBER] = page_number
+
+    if hasattr(result, 'current_rows'):
+        result_rows = result.current_rows or []
+        metas[cassx.ROW_COUNT] = len(result_rows)
+
+    return metas
+
+
+def _sanitize_query(span, query):
+    # TODO (aaditya): fix this hacky type check. we need it to avoid circular imports
+    t = type(query).__name__
+
+    resource = None
+    if t in ('SimpleStatement', 'PreparedStatement'):
+        # reset query if a string is available
+        resource = getattr(query, 'query_string', query)
+    elif t == 'BatchStatement':
+        resource = 'BatchStatement'
+        # Each element in `_statements_and_parameters` is:
+        #   (is_prepared, statement, parameters)
+        #  ref:https://github.com/datastax/python-driver/blob/13d6d72be74f40fcef5ec0f2b3e98538b3b87459/cassandra/query.py#L844
+        #
+        # For prepared statements, the `statement` value is just the query_id
+        #   which is not a statement and when trying to join with other strings
+        #   raises an error in python3 around joining bytes to unicode, so this
+        #   just filters out prepared statements from this tag value
+        q = '; '.join(q[1] for q in query._statements_and_parameters[:2] if not q[0])
+        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='1.0.0')
+def get_traced_cassandra(*args, **kwargs):
+    return _get_traced_cluster(*args, **kwargs)
+
+
+def _get_traced_cluster(*args, **kwargs):
+    return cassandra.cluster.Cluster