From 35a1b14ed6e7b0a315ca940c72ff90ddaf4e204a Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Mon, 7 Nov 2016 18:19:15 +0000 Subject: [PATCH 1/4] pin: add __repr__ --- ddtrace/pin.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ddtrace/pin.py b/ddtrace/pin.py index 690dc73b193..6496361b776 100644 --- a/ddtrace/pin.py +++ b/ddtrace/pin.py @@ -40,3 +40,9 @@ def onto(self, obj): return setattr(obj, '_datadog_pin', self) except AttributeError: log.warn("can't pin onto object", exc_info=True) + + def __repr__(self): + return "Pin(service:%s,app:%s,name:%s)" % ( + self.service, + self.app, + self.name) From 62486b075f92666c30bb64c786d35bf24a094e98 Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Mon, 7 Nov 2016 19:27:52 +0000 Subject: [PATCH 2/4] patch/psycopg2: handle extensions.register_type this function does c-level checks of the type of the argument so it must be a raw psycopg connection passed in. this adds a hook to transparently downgrade. --- ddtrace/contrib/dbapi/__init__.py | 7 ---- ddtrace/contrib/psycopg/__init__.py | 4 +- ddtrace/contrib/psycopg/patch.py | 56 ++++++++++++++++++++------- tests/contrib/__init__.py | 2 + tests/contrib/psycopg/test_psycopg.py | 31 ++++++++------- 5 files changed, 63 insertions(+), 37 deletions(-) diff --git a/ddtrace/contrib/dbapi/__init__.py b/ddtrace/contrib/dbapi/__init__.py index 67affd6f725..311717e7758 100644 --- a/ddtrace/contrib/dbapi/__init__.py +++ b/ddtrace/contrib/dbapi/__init__.py @@ -57,13 +57,6 @@ def __init__(self, conn): name = _get_vendor(conn) self._datadog_pin = Pin(service=name, app=name) - def execute(self, *args, **kwargs): - # this method only exists on some clients, so trigger an attribute - # error if it doesn't. - getattr(self.__wrapped__, 'execute') - # otherwise, keep going. - return self.cursor().execute(*args, **kwargs) - def cursor(self, *args, **kwargs): cursor = self.__wrapped__.cursor(*args, **kwargs) pin = self._datadog_pin diff --git a/ddtrace/contrib/psycopg/__init__.py b/ddtrace/contrib/psycopg/__init__.py index 6e47bca0bdc..3bd8818bde1 100644 --- a/ddtrace/contrib/psycopg/__init__.py +++ b/ddtrace/contrib/psycopg/__init__.py @@ -20,6 +20,6 @@ with require_modules(required_modules) as missing_modules: if not missing_modules: from .connection import connection_factory - from .patch import patch + from .patch import patch, patch_conn - __all__ = ['connection_factory', 'patch'] + __all__ = ['connection_factory', 'patch', 'patch_conn'] diff --git a/ddtrace/contrib/psycopg/patch.py b/ddtrace/contrib/psycopg/patch.py index 9c9a3efd72d..2cb0f300a01 100644 --- a/ddtrace/contrib/psycopg/patch.py +++ b/ddtrace/contrib/psycopg/patch.py @@ -19,20 +19,18 @@ def patch(): """ Patch monkey patches psycopg's connection function so that the connection's functions are traced. """ - setattr(_connect, 'datadog_patched_func', psycopg2.connect) - wrapt.wrap_function_wrapper('psycopg2', 'connect', _connect) + wrapt.wrap_function_wrapper(psycopg2, 'connect', _connect) + _patch_extensions() # do this early just in case -def unpatch(): - """ Unpatch will undo any monkeypatching. """ - connect = getattr(_connect, 'datadog_patched_func', None) - if connect is not None: - psycopg2.connect = connect - -def wrap(conn, service="postgres", tracer=None): +def patch_conn(conn, service="postgres", tracer=None): """ Wrap will patch the instance so that it's queries are traced. Optionally set the service name of the connection. """ + # ensure we've patched extensions (this is idempotent) in + # case we're only tracing some connections. + _patch_extensions() + c = dbapi.TracedConnection(conn) # fetch tags from the dsn @@ -45,15 +43,45 @@ def wrap(conn, service="postgres", tracer=None): "db.application" : dsn.get("application_name"), } - pin = Pin( + Pin( service=service, app="postgres", tracer=tracer, - tags=tags) + tags=tags).onto(c) - pin.onto(c) return c +def _patch_extensions(): + # we must patch extensions all the time (it's pretty harmless) so split + # from global patching of connections. must be idempotent. + for m, f, w in _extensions: + if not hasattr(m, f) or isinstance(getattr(m, f), wrapt.ObjectProxy): + continue + wrapt.wrap_function_wrapper(m, f, w) + + +# +# monkeypatch targets +# + def _connect(connect_func, _, args, kwargs): - db = connect_func(*args, **kwargs) - return wrap(db) + conn = connect_func(*args, **kwargs) + return patch_conn(conn) + +def _extensions_register_type(func, _, args, kwargs): + def _unroll_args(obj, scope=None): + return obj, scope + obj, scope = _unroll_args(*args, **kwargs) + + # register_type performs a c-level check of the object + # type so we must be sure to pass in the actual db connection + if scope and isinstance(scope, wrapt.ObjectProxy): + scope = scope.__wrapped__ + + return func(obj, scope) if scope else func(obj) + + +# extension hooks +_extensions = [ + (psycopg2.extensions, 'register_type', _extensions_register_type), +] diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py index e69de29bb2d..139597f9cb0 100644 --- a/tests/contrib/__init__.py +++ b/tests/contrib/__init__.py @@ -0,0 +1,2 @@ + + diff --git a/tests/contrib/psycopg/test_psycopg.py b/tests/contrib/psycopg/test_psycopg.py index 1c529dd63e3..9a43b23ece5 100644 --- a/tests/contrib/psycopg/test_psycopg.py +++ b/tests/contrib/psycopg/test_psycopg.py @@ -3,6 +3,7 @@ # 3p import psycopg2 +from psycopg2 import extras from nose.tools import eq_ # project @@ -10,8 +11,9 @@ from ddtrace.contrib.psycopg import connection_factory # testing -from ..config import POSTGRES_CONFIG -from ...test_tracer import DummyWriter +from tests.contrib.config import POSTGRES_CONFIG +from tests.test_tracer import get_test_tracer +from ddtrace.contrib.psycopg import patch_conn TEST_PORT = str(POSTGRES_CONFIG['port']) @@ -71,21 +73,22 @@ def assert_conn_is_traced(tracer, db, service): eq_(span.span_type, "sql") def test_manual_wrap(): - from ddtrace.contrib.psycopg.patch import wrap - db = psycopg2.connect(**POSTGRES_CONFIG) - - writer = DummyWriter() - tracer = Tracer() - tracer.writer = writer - wrapped = wrap(db, service="foo", tracer=tracer) + conn = psycopg2.connect(**POSTGRES_CONFIG) + tracer = get_test_tracer() + wrapped = patch_conn(conn, service="foo", tracer=tracer) assert_conn_is_traced(tracer, wrapped, "foo") - +def test_manual_wrap_extension_types(): + conn = psycopg2.connect(**POSTGRES_CONFIG) + tracer = get_test_tracer() + wrapped = patch_conn(conn, service="foo", tracer=tracer) + # NOTE: this will crash if it doesn't work. + # _ext.register_type(_ext.UUID, conn_or_curs) + # TypeError: argument 2 must be a connection, cursor or None + extras.register_uuid(conn_or_curs=wrapped) def test_connect_factory(): - writer = DummyWriter() - tracer = Tracer() - tracer.writer = writer + tracer = get_test_tracer() services = ["db", "another"] for service in services: @@ -94,7 +97,7 @@ def test_connect_factory(): assert_conn_is_traced(tracer, db, service) # ensure we have the service types - services = writer.pop_services() + services = tracer.writer.pop_services() expected = { "db" : {"app":"postgres", "app_type":"db"}, "another" : {"app":"postgres", "app_type":"db"}, From 1f9f9957e592ac6759f1405d4b6be2d629782867 Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Mon, 7 Nov 2016 19:32:35 +0000 Subject: [PATCH 3/4] remove empty file --- tests/contrib/__init__.py | 2 -- 1 file changed, 2 deletions(-) delete mode 100644 tests/contrib/__init__.py diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py deleted file mode 100644 index 139597f9cb0..00000000000 --- a/tests/contrib/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ - - From 3cd2f4f4f7900d986f8f43b203386ce325c2553c Mon Sep 17 00:00:00 2001 From: Matt Perpick Date: Mon, 7 Nov 2016 20:00:59 +0000 Subject: [PATCH 4/4] add file we need --- tests/contrib/__init__.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 tests/contrib/__init__.py diff --git a/tests/contrib/__init__.py b/tests/contrib/__init__.py new file mode 100644 index 00000000000..e69de29bb2d