From 6915b7c3d8bef51c3fda0d900e9c5727d0e9840b Mon Sep 17 00:00:00 2001 From: Benjamin Fernandes Date: Mon, 27 Jun 2016 17:35:04 +0200 Subject: [PATCH] Use unified structure for contrib imports --- ddtrace/contrib/elasticsearch/__init__.py | 10 +- ddtrace/contrib/elasticsearch/transport.py | 5 +- ddtrace/contrib/flask/__init__.py | 137 +-------------------- ddtrace/contrib/flask/middleware.py | 134 ++++++++++++++++++++ ddtrace/contrib/psycopg/__init__.py | 125 +------------------ ddtrace/contrib/psycopg/connection.py | 113 +++++++++++++++++ ddtrace/contrib/pylons/__init__.py | 56 +-------- ddtrace/contrib/pylons/middleware.py | 55 +++++++++ ddtrace/contrib/sqlite3/__init__.py | 58 +-------- ddtrace/contrib/sqlite3/connection.py | 56 +++++++++ ddtrace/contrib/util.py | 19 +++ tests/contrib/elasticsearch/test.py | 19 ++- tests/contrib/flask/test_flask.py | 6 + tests/contrib/psycopg/test_psycopg.py | 15 +-- tests/contrib/pylons/test_pylons.py | 1 - 15 files changed, 424 insertions(+), 385 deletions(-) create mode 100644 ddtrace/contrib/flask/middleware.py create mode 100644 ddtrace/contrib/psycopg/connection.py create mode 100644 ddtrace/contrib/pylons/middleware.py create mode 100644 ddtrace/contrib/sqlite3/connection.py create mode 100644 ddtrace/contrib/util.py diff --git a/ddtrace/contrib/elasticsearch/__init__.py b/ddtrace/contrib/elasticsearch/__init__.py index c20ca59237a..a493a531274 100644 --- a/ddtrace/contrib/elasticsearch/__init__.py +++ b/ddtrace/contrib/elasticsearch/__init__.py @@ -1,3 +1,9 @@ -from .transport import get_traced_transport +from ..util import require_modules -__all__ = ['get_traced_transport'] +required_modules = ['elasticsearch'] + +with require_modules(required_modules) as missing_modules: + if not missing_modules: + from .transport import get_traced_transport + + __all__ = ['get_traced_transport'] diff --git a/ddtrace/contrib/elasticsearch/transport.py b/ddtrace/contrib/elasticsearch/transport.py index 72e98167132..cd7b7c3dd13 100644 --- a/ddtrace/contrib/elasticsearch/transport.py +++ b/ddtrace/contrib/elasticsearch/transport.py @@ -1,7 +1,4 @@ -try: - from elasticsearch import Transport -except ImportError: - Transport = object +from elasticsearch import Transport from .quantize import quantize from . import metadata diff --git a/ddtrace/contrib/flask/__init__.py b/ddtrace/contrib/flask/__init__.py index c4f0d41d575..6d2b7a9cca0 100644 --- a/ddtrace/contrib/flask/__init__.py +++ b/ddtrace/contrib/flask/__init__.py @@ -1,134 +1,9 @@ -""" -Datadog trace code for flask. +from ..util import require_modules -Requires a modern version of flask and the `blinker` library (which is a -dependency of flask signals). -""" +required_modules = ['flask'] -# stdlib -import time -import logging +with require_modules(required_modules) as missing_modules: + if not missing_modules: + from .middleware import TraceMiddleware -# project -from ...ext import http, errors - -# 3p -from flask import g, request, signals - - -class TraceMiddleware(object): - - def __init__(self, app, tracer, service="flask", use_signals=True): - self.app = app - self.app.logger.info("initializing trace middleware") - - # save our traces. - self._tracer = tracer - self._service = service - - self.use_signals = use_signals - - if self.use_signals and signals.signals_available: - # if we're using signals, and things are correctly installed, use - # signal hooks to track the responses. - self.app.logger.info("connecting trace signals") - signals.request_started.connect(self._request_started, sender=self.app) - signals.request_finished.connect(self._request_finished, sender=self.app) - signals.got_request_exception.connect(self._request_exception, sender=self.app) - signals.before_render_template.connect(self._template_started, sender=self.app) - signals.template_rendered.connect(self._template_done, sender=self.app) - else: - if self.use_signals: # warn the user that signals lib isn't installed - self.app.logger.info(_blinker_not_installed_msg) - - # Fallback to using after request hook. Unfortunately, this won't - # handle exceptions. - self.app.before_request(self._before_request) - self.app.after_request(self._after_request) - - # common methods - - def _start_span(self): - try: - g.flask_datadog_span = self._tracer.trace( - "flask.request", - service=self._service, - span_type=http.TYPE, - ) - except Exception: - self.app.logger.exception("error tracing request") - - def _finish_span(self, response=None, exception=None): - """ Close and finsh the active span if it exists. """ - span = getattr(g, 'flask_datadog_span', None) - if span: - error = 0 - code = response.status_code if response else None - - # if we didn't get a response, but we did get an exception, set - # codes accordingly. - if not response and exception: - error = 1 - code = 500 - span.set_tag(errors.ERROR_TYPE, type(exception)) - span.set_tag(errors.ERROR_MSG, exception) - - span.resource = str(request.endpoint or "").lower() - span.set_tag(http.URL, str(request.base_url or "")) - span.set_tag(http.STATUS_CODE, code) - span.error = error - span.finish() - # Clear our span just in case. - g.flask_datadog_span = None - - # Request hook methods - - def _before_request(self): - """ Starts tracing the current request and stores it in the global - request object. - """ - self._start_span() - - def _after_request(self, response): - """ handles a successful response. """ - try: - self._finish_span(response=response) - except Exception: - self.app.logger.exception("error finishing trace") - finally: - return response - - # signal handling methods - - def _request_started(self, sender): - self._start_span() - - def _request_finished(self, sender, response, **kwargs): - try: - self._finish_span(response=response) - except Exception: - self.app.logger.exception("error finishing trace") - return response - - def _request_exception(self, *args, **kwargs): - """ handles an error response. """ - exception = kwargs.pop("exception", None) - try: - self._finish_span(exception=exception) - except Exception: - self.app.logger.exception("error tracing error") - - def _template_started(self, sender, template, *args, **kwargs): - span = self._tracer.trace('flask.template') - try: - span.span_type = http.TEMPLATE - span.set_tag("flask.template", template.name or "string") - finally: - g.flask_datadog_tmpl_span = span - - def _template_done(self, *arg, **kwargs): - span = getattr(g, 'flask_datadog_tmpl_span', None) - if span: - span.finish() - -_blinker_not_installed_msg = "please install blinker to use flask signals. http://flask.pocoo.org/docs/0.11/signals/" + __all__ = ['TraceMiddleware'] diff --git a/ddtrace/contrib/flask/middleware.py b/ddtrace/contrib/flask/middleware.py new file mode 100644 index 00000000000..c4f0d41d575 --- /dev/null +++ b/ddtrace/contrib/flask/middleware.py @@ -0,0 +1,134 @@ +""" +Datadog trace code for flask. + +Requires a modern version of flask and the `blinker` library (which is a +dependency of flask signals). +""" + +# stdlib +import time +import logging + +# project +from ...ext import http, errors + +# 3p +from flask import g, request, signals + + +class TraceMiddleware(object): + + def __init__(self, app, tracer, service="flask", use_signals=True): + self.app = app + self.app.logger.info("initializing trace middleware") + + # save our traces. + self._tracer = tracer + self._service = service + + self.use_signals = use_signals + + if self.use_signals and signals.signals_available: + # if we're using signals, and things are correctly installed, use + # signal hooks to track the responses. + self.app.logger.info("connecting trace signals") + signals.request_started.connect(self._request_started, sender=self.app) + signals.request_finished.connect(self._request_finished, sender=self.app) + signals.got_request_exception.connect(self._request_exception, sender=self.app) + signals.before_render_template.connect(self._template_started, sender=self.app) + signals.template_rendered.connect(self._template_done, sender=self.app) + else: + if self.use_signals: # warn the user that signals lib isn't installed + self.app.logger.info(_blinker_not_installed_msg) + + # Fallback to using after request hook. Unfortunately, this won't + # handle exceptions. + self.app.before_request(self._before_request) + self.app.after_request(self._after_request) + + # common methods + + def _start_span(self): + try: + g.flask_datadog_span = self._tracer.trace( + "flask.request", + service=self._service, + span_type=http.TYPE, + ) + except Exception: + self.app.logger.exception("error tracing request") + + def _finish_span(self, response=None, exception=None): + """ Close and finsh the active span if it exists. """ + span = getattr(g, 'flask_datadog_span', None) + if span: + error = 0 + code = response.status_code if response else None + + # if we didn't get a response, but we did get an exception, set + # codes accordingly. + if not response and exception: + error = 1 + code = 500 + span.set_tag(errors.ERROR_TYPE, type(exception)) + span.set_tag(errors.ERROR_MSG, exception) + + span.resource = str(request.endpoint or "").lower() + span.set_tag(http.URL, str(request.base_url or "")) + span.set_tag(http.STATUS_CODE, code) + span.error = error + span.finish() + # Clear our span just in case. + g.flask_datadog_span = None + + # Request hook methods + + def _before_request(self): + """ Starts tracing the current request and stores it in the global + request object. + """ + self._start_span() + + def _after_request(self, response): + """ handles a successful response. """ + try: + self._finish_span(response=response) + except Exception: + self.app.logger.exception("error finishing trace") + finally: + return response + + # signal handling methods + + def _request_started(self, sender): + self._start_span() + + def _request_finished(self, sender, response, **kwargs): + try: + self._finish_span(response=response) + except Exception: + self.app.logger.exception("error finishing trace") + return response + + def _request_exception(self, *args, **kwargs): + """ handles an error response. """ + exception = kwargs.pop("exception", None) + try: + self._finish_span(exception=exception) + except Exception: + self.app.logger.exception("error tracing error") + + def _template_started(self, sender, template, *args, **kwargs): + span = self._tracer.trace('flask.template') + try: + span.span_type = http.TEMPLATE + span.set_tag("flask.template", template.name or "string") + finally: + g.flask_datadog_tmpl_span = span + + def _template_done(self, *arg, **kwargs): + span = getattr(g, 'flask_datadog_tmpl_span', None) + if span: + span.finish() + +_blinker_not_installed_msg = "please install blinker to use flask signals. http://flask.pocoo.org/docs/0.11/signals/" diff --git a/ddtrace/contrib/psycopg/__init__.py b/ddtrace/contrib/psycopg/__init__.py index be44f475009..70fad3d752a 100644 --- a/ddtrace/contrib/psycopg/__init__.py +++ b/ddtrace/contrib/psycopg/__init__.py @@ -1,122 +1,9 @@ -""" -Tracing utilities for the psycopg potgres client library. -""" - -# stdlib -import functools -import logging - -from ...ext import net -from ...ext import sql as sqlx - -# 3p -_installed = False -try: - from psycopg2.extensions import connection, cursor - _installed = True -except ImportError: - connection, cursor = object, object - - -log = logging.getLogger(__name__) - - -def connection_factory(tracer, service="postgres"): - """ Return a connection factory class that will can be used to trace - sqlite queries. - - >>> factory = connection_factor(my_tracer, service="my_db_service") - >>> conn = pyscopg2.connect(..., connection_factory=factory) - """ - if not _installed: - log.info("missing psycopg import") - return None - - return functools.partial(TracedConnection, - datadog_tracer=tracer, - datadog_service=service, - ) - - - -class TracedCursor(cursor): - """Wrapper around cursor creating one span per query""" - - def __init__(self, *args, **kwargs): - self._datadog_tracer = kwargs.pop("datadog_tracer", None) - self._datadog_service = kwargs.pop("datadog_service", None) - self._datadog_tags = kwargs.pop("datadog_tags", None) - super(TracedCursor, self).__init__(*args, **kwargs) - - def execute(self, query, vars=None): - """ just wrap the cursor execution in a span """ - if not self._datadog_tracer: - return cursor.execute(self, query, vars) - - with self._datadog_tracer.trace("postgres.query") as s: - s.resource = query - s.service = self._datadog_service - s.span_type = sqlx.TYPE - s.set_tag(sqlx.QUERY, query) - s.set_tags(self._datadog_tags) - try: - return super(TracedCursor, self).execute(query, vars) - finally: - s.set_tag("db.rowcount", self.rowcount) - - def callproc(self, procname, vars=None): - """ just wrap the execution in a span """ - return cursor.callproc(self, procname, vars) - - -class TracedConnection(connection): - """Wrapper around psycopg2 for tracing""" - - def __init__(self, *args, **kwargs): - - self._datadog_tracer = kwargs.pop("datadog_tracer", None) - self._datadog_service = kwargs.pop("datadog_service", None) - - super(TracedConnection, self).__init__(*args, **kwargs) - - # add metadata (from the connection, string, etc) - dsn = _parse_dsn(self.dsn) - self._datadog_tags = { - net.TARGET_HOST: dsn.get("host"), - net.TARGET_PORT: dsn.get("port"), - "db.name": dsn.get("dbname"), - "db.user": dsn.get("user"), - "db.application" : dsn.get("application_name"), - } - - self._datadog_cursor_class = functools.partial(TracedCursor, - datadog_tracer=self._datadog_tracer, - datadog_service=self._datadog_service, - datadog_tags=self._datadog_tags, - ) - - # DogTrace.register_service( - # service=self._dogtrace_service, - # app="postgres", - # app_type="sql", - # ) - - def cursor(self, *args, **kwargs): - """ register our custom cursor factory """ - kwargs.setdefault('cursor_factory', self._datadog_cursor_class) - return super(TracedConnection, self).cursor(*args, **kwargs) - - -def _parse_dsn(dsn): - """ - Return a diciontary of the components of a postgres DSN. - - >>> _parse_dsn('user=dog port=1543 dbname=dogdata') - {"user":"dog", "port":"1543", "dbname":"dogdata"} - """ - # FIXME: replace by psycopg2.extensions.parse_dsn when available - # https://github.com/psycopg/psycopg2/pull/321 - return {chunk.split("=")[0]: chunk.split("=")[1] for chunk in dsn.split() if "=" in chunk} +from ..util import require_modules +required_modules = ['psycopg2'] +with require_modules(required_modules) as missing_modules: + if not missing_modules: + from .connection import connection_factory + __all__ = ['connection_factory'] diff --git a/ddtrace/contrib/psycopg/connection.py b/ddtrace/contrib/psycopg/connection.py new file mode 100644 index 00000000000..18f35c6b38b --- /dev/null +++ b/ddtrace/contrib/psycopg/connection.py @@ -0,0 +1,113 @@ +""" +Tracing utilities for the psycopg potgres client library. +""" + +# stdlib +import functools +import logging + +from ...ext import net +from ...ext import sql as sqlx + +# 3p +from psycopg2.extensions import connection, cursor + + +log = logging.getLogger(__name__) + + +def connection_factory(tracer, service="postgres"): + """ Return a connection factory class that will can be used to trace + sqlite queries. + + >>> factory = connection_factor(my_tracer, service="my_db_service") + >>> conn = pyscopg2.connect(..., connection_factory=factory) + """ + return functools.partial(TracedConnection, + datadog_tracer=tracer, + datadog_service=service, + ) + + + +class TracedCursor(cursor): + """Wrapper around cursor creating one span per query""" + + def __init__(self, *args, **kwargs): + self._datadog_tracer = kwargs.pop("datadog_tracer", None) + self._datadog_service = kwargs.pop("datadog_service", None) + self._datadog_tags = kwargs.pop("datadog_tags", None) + super(TracedCursor, self).__init__(*args, **kwargs) + + def execute(self, query, vars=None): + """ just wrap the cursor execution in a span """ + if not self._datadog_tracer: + return cursor.execute(self, query, vars) + + with self._datadog_tracer.trace("postgres.query") as s: + s.resource = query + s.service = self._datadog_service + s.span_type = sqlx.TYPE + s.set_tag(sqlx.QUERY, query) + s.set_tags(self._datadog_tags) + try: + return super(TracedCursor, self).execute(query, vars) + finally: + s.set_tag("db.rowcount", self.rowcount) + + def callproc(self, procname, vars=None): + """ just wrap the execution in a span """ + return cursor.callproc(self, procname, vars) + + +class TracedConnection(connection): + """Wrapper around psycopg2 for tracing""" + + def __init__(self, *args, **kwargs): + + self._datadog_tracer = kwargs.pop("datadog_tracer", None) + self._datadog_service = kwargs.pop("datadog_service", None) + + super(TracedConnection, self).__init__(*args, **kwargs) + + # add metadata (from the connection, string, etc) + dsn = _parse_dsn(self.dsn) + self._datadog_tags = { + net.TARGET_HOST: dsn.get("host"), + net.TARGET_PORT: dsn.get("port"), + "db.name": dsn.get("dbname"), + "db.user": dsn.get("user"), + "db.application" : dsn.get("application_name"), + } + + self._datadog_cursor_class = functools.partial(TracedCursor, + datadog_tracer=self._datadog_tracer, + datadog_service=self._datadog_service, + datadog_tags=self._datadog_tags, + ) + + # DogTrace.register_service( + # service=self._dogtrace_service, + # app="postgres", + # app_type="sql", + # ) + + def cursor(self, *args, **kwargs): + """ register our custom cursor factory """ + kwargs.setdefault('cursor_factory', self._datadog_cursor_class) + return super(TracedConnection, self).cursor(*args, **kwargs) + + +def _parse_dsn(dsn): + """ + Return a diciontary of the components of a postgres DSN. + + >>> _parse_dsn('user=dog port=1543 dbname=dogdata') + {"user":"dog", "port":"1543", "dbname":"dogdata"} + """ + # FIXME: replace by psycopg2.extensions.parse_dsn when available + # https://github.com/psycopg/psycopg2/pull/321 + return {chunk.split("=")[0]: chunk.split("=")[1] for chunk in dsn.split() if "=" in chunk} + + + diff --git a/ddtrace/contrib/pylons/__init__.py b/ddtrace/contrib/pylons/__init__.py index 8ba9ec51a13..0c35fcb59f3 100644 --- a/ddtrace/contrib/pylons/__init__.py +++ b/ddtrace/contrib/pylons/__init__.py @@ -1,55 +1,3 @@ -import logging +from .middleware import PylonsTraceMiddleware -from ...ext import http - -log = logging.getLogger(__name__) - - -class PylonsTraceMiddleware(object): - - def __init__(self, app, tracer, service="pylons"): - self.app = app - self._service = service - self._tracer = tracer - - def __call__(self, environ, start_response): - span = None - try: - span = self._tracer.trace("pylons.request", service=self._service, span_type=http.TYPE) - log.debug("Initialize new trace %d", span.trace_id) - - def _start_response(status, *args, **kwargs): - """ a patched response callback which will pluck some metadata. """ - span.span_type = http.TYPE - http_code = int(status.split()[0]) - span.set_tag(http.STATUS_CODE, http_code) - if http_code >= 500: - span.error = 1 - return start_response(status, *args, **kwargs) - except Exception: - log.exception("error starting span") - - try: - return self.app(environ, _start_response) - except Exception: - if span: - span.set_traceback() - raise - finally: - if not span: - return - try: - controller = environ.get('pylons.routes_dict', {}).get('controller') - action = environ.get('pylons.routes_dict', {}).get('action') - span.resource = "%s.%s" % (controller, action) - - span.set_tags({ - http.METHOD: environ.get('REQUEST_METHOD'), - http.URL: environ.get('PATH_INFO'), - "pylons.user": environ.get('REMOTE_USER', ''), - "pylons.route.controller": controller, - "pylons.route.action": action, - }) - span.finish() - except Exception: - log.exception("Error finishing trace") +__all__ = ['PylonsTraceMiddleware'] diff --git a/ddtrace/contrib/pylons/middleware.py b/ddtrace/contrib/pylons/middleware.py new file mode 100644 index 00000000000..8ba9ec51a13 --- /dev/null +++ b/ddtrace/contrib/pylons/middleware.py @@ -0,0 +1,55 @@ +import logging + +from ...ext import http + +log = logging.getLogger(__name__) + + +class PylonsTraceMiddleware(object): + + def __init__(self, app, tracer, service="pylons"): + self.app = app + self._service = service + self._tracer = tracer + + def __call__(self, environ, start_response): + span = None + try: + span = self._tracer.trace("pylons.request", service=self._service, span_type=http.TYPE) + log.debug("Initialize new trace %d", span.trace_id) + + def _start_response(status, *args, **kwargs): + """ a patched response callback which will pluck some metadata. """ + span.span_type = http.TYPE + http_code = int(status.split()[0]) + span.set_tag(http.STATUS_CODE, http_code) + if http_code >= 500: + span.error = 1 + return start_response(status, *args, **kwargs) + except Exception: + log.exception("error starting span") + + try: + return self.app(environ, _start_response) + except Exception: + if span: + span.set_traceback() + raise + finally: + if not span: + return + try: + controller = environ.get('pylons.routes_dict', {}).get('controller') + action = environ.get('pylons.routes_dict', {}).get('action') + span.resource = "%s.%s" % (controller, action) + + span.set_tags({ + http.METHOD: environ.get('REQUEST_METHOD'), + http.URL: environ.get('PATH_INFO'), + "pylons.user": environ.get('REMOTE_USER', ''), + "pylons.route.controller": controller, + "pylons.route.action": action, + }) + span.finish() + except Exception: + log.exception("Error finishing trace") diff --git a/ddtrace/contrib/sqlite3/__init__.py b/ddtrace/contrib/sqlite3/__init__.py index 9f6722f30fa..996d7eb107f 100644 --- a/ddtrace/contrib/sqlite3/__init__.py +++ b/ddtrace/contrib/sqlite3/__init__.py @@ -1,57 +1,3 @@ +from .connection import connection_factory -import functools - -from sqlite3 import Connection, Cursor -from ...ext import sql as sqlx - - -def connection_factory(tracer, service="sqlite3"): - """ Return a connection factory class that will can be used to trace - sqlite queries. - - >>> factory = connection_factor(my_tracer, service="my_db_service") - >>> conn = sqlite3.connect(":memory:", factory=factory) - """ - return functools.partial(TracedConnection, - datadog_tracer=tracer, - datadog_service=service, - ) - - -class TracedCursor(Cursor): - """ A cursor base class that will trace sql queries. """ - - def __init__(self, *args, **kwargs): - self._datadog_tracer = kwargs.pop("datadog_tracer", None) - self._datadog_service = kwargs.pop("datadog_service", None) - Cursor.__init__(self, *args, **kwargs) - - def execute(self, sql, *args, **kwargs): - if not self._datadog_tracer: - return Cursor.execute(self, sql, *args, **kwargs) - - with self._datadog_tracer.trace("sqlite3.query", span_type=sqlx.TYPE) as s: - s.set_tag(sqlx.QUERY, sql) - s.service = self._datadog_service - s.resource = sql # will be normalized - return Cursor.execute(self, sql, *args, **kwargs) - - -class TracedConnection(Connection): - """ A cursor base class that will trace sql queries. """ - - def __init__(self, *args, **kwargs): - self._datadog_tracer = kwargs.pop("datadog_tracer", None) - self._datadog_service = kwargs.pop("datadog_service", None) - Connection.__init__(self, *args, **kwargs) - - self._datadog_cursor_class = functools.partial(TracedCursor, - datadog_tracer=self._datadog_tracer, - datadog_service=self._datadog_service, - ) - - def cursor(self, *args, **kwargs): - if self._datadog_tracer: - kwargs.setdefault('factory', self._datadog_cursor_class) - return Connection.cursor(self, *args, **kwargs) - +__all__ = ['connection_factory'] diff --git a/ddtrace/contrib/sqlite3/connection.py b/ddtrace/contrib/sqlite3/connection.py new file mode 100644 index 00000000000..11626ca31d0 --- /dev/null +++ b/ddtrace/contrib/sqlite3/connection.py @@ -0,0 +1,56 @@ +import functools + +from sqlite3 import Connection, Cursor +from ...ext import sql as sqlx + + +def connection_factory(tracer, service="sqlite3"): + """ Return a connection factory class that will can be used to trace + sqlite queries. + + >>> factory = connection_factor(my_tracer, service="my_db_service") + >>> conn = sqlite3.connect(":memory:", factory=factory) + """ + return functools.partial(TracedConnection, + datadog_tracer=tracer, + datadog_service=service, + ) + + +class TracedCursor(Cursor): + """ A cursor base class that will trace sql queries. """ + + def __init__(self, *args, **kwargs): + self._datadog_tracer = kwargs.pop("datadog_tracer", None) + self._datadog_service = kwargs.pop("datadog_service", None) + Cursor.__init__(self, *args, **kwargs) + + def execute(self, sql, *args, **kwargs): + if not self._datadog_tracer: + return Cursor.execute(self, sql, *args, **kwargs) + + with self._datadog_tracer.trace("sqlite3.query", span_type=sqlx.TYPE) as s: + s.set_tag(sqlx.QUERY, sql) + s.service = self._datadog_service + s.resource = sql # will be normalized + return Cursor.execute(self, sql, *args, **kwargs) + + +class TracedConnection(Connection): + """ A cursor base class that will trace sql queries. """ + + def __init__(self, *args, **kwargs): + self._datadog_tracer = kwargs.pop("datadog_tracer", None) + self._datadog_service = kwargs.pop("datadog_service", None) + Connection.__init__(self, *args, **kwargs) + + self._datadog_cursor_class = functools.partial(TracedCursor, + datadog_tracer=self._datadog_tracer, + datadog_service=self._datadog_service, + ) + + def cursor(self, *args, **kwargs): + if self._datadog_tracer: + kwargs.setdefault('factory', self._datadog_cursor_class) + return Connection.cursor(self, *args, **kwargs) + diff --git a/ddtrace/contrib/util.py b/ddtrace/contrib/util.py new file mode 100644 index 00000000000..adedc8f55da --- /dev/null +++ b/ddtrace/contrib/util.py @@ -0,0 +1,19 @@ +from importlib import import_module + + +class require_modules(object): + """Context manager to check the availability of required modules""" + + def __init__(self, modules): + self._missing_modules = [] + for module in modules: + try: + import_module(module) + except ImportError: + self._missing_modules.append(module) + + def __enter__(self): + return self._missing_modules + + def __exit__(self, exc_type, exc_value, traceback): + return False diff --git a/tests/contrib/elasticsearch/test.py b/tests/contrib/elasticsearch/test.py index ad420017c3e..ebe37ee1df9 100644 --- a/tests/contrib/elasticsearch/test.py +++ b/tests/contrib/elasticsearch/test.py @@ -1,15 +1,15 @@ import unittest -from nose.tools import eq_ -# We should probably be smarter than that -try: - import elasticsearch -except ImportError: - elasticsearch = None +from ddtrace.contrib.elasticsearch import missing_modules + +if missing_modules: + raise unittest.SkipTest("Missing dependencies %s" % missing_modules) + +import elasticsearch +from nose.tools import eq_ -from ddtrace.contrib.elasticsearch import metadata -from ddtrace.contrib.elasticsearch import get_traced_transport from ddtrace.tracer import Tracer +from ddtrace.contrib.elasticsearch import get_traced_transport, metadata from ...test_tracer import DummyWriter @@ -27,9 +27,6 @@ class ElasticsearchTest(unittest.TestCase): def setUp(self): """Prepare ES""" - if not elasticsearch: - raise unittest.SkipTest("elasticsearch module isn't available") - es = elasticsearch.Elasticsearch() es.indices.delete(index=self.ES_INDEX, ignore=[400, 404]) diff --git a/tests/contrib/flask/test_flask.py b/tests/contrib/flask/test_flask.py index 8b4050940b8..dda653eaddb 100644 --- a/tests/contrib/flask/test_flask.py +++ b/tests/contrib/flask/test_flask.py @@ -1,3 +1,9 @@ +import unittest + +from ddtrace.contrib.flask import missing_modules + +if missing_modules: + raise unittest.SkipTest("Missing dependencies %s" % missing_modules) import time import logging diff --git a/tests/contrib/psycopg/test_psycopg.py b/tests/contrib/psycopg/test_psycopg.py index eae7349691d..466c732d052 100644 --- a/tests/contrib/psycopg/test_psycopg.py +++ b/tests/contrib/psycopg/test_psycopg.py @@ -1,7 +1,14 @@ +import unittest + +from ddtrace.contrib.flask import missing_modules + +if missing_modules: + raise unittest.SkipTest("Missing dependencies %s" % missing_modules) + import time +import psycopg2 from nose.tools import eq_ -from nose.plugins.skip import SkipTest from ddtrace import Tracer from ddtrace.contrib.psycopg import connection_factory @@ -10,12 +17,6 @@ def test_wrap(): - - try: - import psycopg2 - except ImportError: - raise SkipTest("missing psycopg") - writer = DummyWriter() tracer = Tracer(writer=writer) diff --git a/tests/contrib/pylons/test_pylons.py b/tests/contrib/pylons/test_pylons.py index 5df6e83d928..a66316ad620 100644 --- a/tests/contrib/pylons/test_pylons.py +++ b/tests/contrib/pylons/test_pylons.py @@ -1,4 +1,3 @@ - import time from nose.tools import eq_