feat(dbapi): support Connection context management usage (DataDog#1762)

* fix(psycopg2): add test for connection used with contextmanager

* dbapi: handle Connection __enter__

* psycopg2 < 2.5 connection doesn't support context management

* add dbapi to wordlist

* add tests for mysqldb

* add dbapi tests

* Update ddtrace/contrib/dbapi/__init__.py

Co-authored-by: Julien Danjou <julien@danjou.info>

* add some doc strings

* add test cases for other dbapi libraries

* add test for no context manager

* update comment

Co-authored-by: Kyle Verhoog <kyle.verhoog@datadoghq.com>
Co-authored-by: Julien Danjou <julien@danjou.info>
Co-authored-by: Kyle Verhoog <kyle@verhoog.ca>

ddtrace/contrib/dbapi/__init__.py

@@ -9,7 +9,7 @@
 from ...settings import config
 from ...utils.formats import asbool, get_env
 from ...vendor import wrapt
-from ..trace_utils import ext_service
+from ..trace_utils import ext_service, iswrapped
 
 
 log = get_logger(__name__)
@@ -178,6 +178,49 @@ def __init__(self, conn, pin=None, cfg=None, cursor_cls=None):
         self._self_cursor_cls = cursor_cls
         self._self_config = cfg or config.dbapi2
 
+    def __enter__(self):
+        """Context management is not defined by the dbapi spec.
+
+        This means unfortunately that the database clients each define their own
+        implementations.
+
+        The ones we know about are:
+
+        - mysqlclient<2.0 which returns a cursor instance. >=2.0 returns a
+          connection instance.
+        - psycopg returns a connection.
+        - pyodbc returns a connection.
+        - pymysql doesn't implement it.
+        - sqlite3 returns the connection.
+        """
+        r = self.__wrapped__.__enter__()
+
+        if hasattr(r, "cursor"):
+            # r is Connection-like.
+            if r is self.__wrapped__:
+                # Return the reference to this proxy object. Returning r would
+                # return the untraced reference.
+                return self
+            else:
+                # r is a different connection object.
+                # This should not happen in practice but play it safe so that
+                # the original functionality is maintained.
+                return r
+        elif hasattr(r, "execute"):
+            # r is Cursor-like.
+            if iswrapped(r):
+                return r
+            else:
+                pin = Pin.get_from(self)
+                cfg = _get_config(self._self_config)
+                if not pin:
+                    return r
+                return self._self_cursor_cls(r, pin, cfg)
+        else:
+            # Otherwise r is some other object, so maintain the functionality
+            # of the original.
+            return r
+
     def _trace_method(self, method, name, extra_tags, *args, **kwargs):
         pin = Pin.get_from(self)
         if not pin or not pin.enabled():

docs/spelling_wordlist.txt

@@ -0,0 +1,100 @@
+CPython
+INfo
+MySQL
+OpenTracing
+aiobotocore
+aiohttp
+aiopg
+algolia
+algoliasearch
+analytics
+api
+app
+asgi
+autodetected
+autopatching
+backend
+bikeshedding
+boto
+botocore
+config
+coroutine
+coroutines
+datadog
+datadoghq
+datastore
+dbapi
+ddtrace
+django
+dogstatsd
+elasticsearch
+enqueue
+entrypoint
+entrypoints
+gRPC
+gevent
+greenlet
+greenlets
+grpc
+hostname
+http
+httplib
+https
+iPython
+integration
+integrations
+jinja
+kombu
+kubernetes
+kwarg
+lifecycle
+mako
+memcached
+metadata
+microservices
+middleware
+mongoengine
+mysql
+mysqlclient
+mysqldb
+namespace
+opentracer
+opentracing
+plugin
+posix
+postgres
+prepended
+profiler
+psycopg
+py
+pylibmc
+pymemcache
+pymongo
+pymysql
+pynamodb
+pyodbc
+quickstart
+redis
+rediscluster
+renderers
+repo
+runnable
+runtime
+sanic
+sqlalchemy
+sqlite
+starlette
+stringable
+subdomains
+submodules
+timestamp
+tweens
+uWSGI
+unix
+unregister
+url
+urls
+username
+uvicorn
+vertica
+whitelist

releasenotes/notes/dbapi-ctx-manager-008915d487d9f50d.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    dbapi: add support for connection context manager usage

tests/contrib/dbapi/test_dbapi.py

@@ -1,5 +1,7 @@
 import mock
 
+import pytest
+
 from ddtrace import Pin
 from ddtrace.constants import ANALYTICS_SAMPLE_RATE_KEY
 from ddtrace.contrib.dbapi import FetchTracedCursor, TracedCursor, TracedConnection
@@ -589,3 +591,179 @@ def test_connection_analytics_with_rate(self):
             traced_connection.commit()
             span = tracer.writer.pop()[0]
             self.assertIsNone(span.get_metric(ANALYTICS_SAMPLE_RATE_KEY))
+
+    def test_connection_context_manager(self):
+
+        class Cursor(object):
+            rowcount = 0
+
+            def execute(self, *args, **kwargs):
+                pass
+
+            def fetchall(self, *args, **kwargs):
+                pass
+
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *exc):
+                return False
+
+            def commit(self, *args, **kwargs):
+                pass
+
+        # When a connection is returned from a context manager the object proxy
+        # should be returned so that tracing works.
+
+        class ConnectionConnection(object):
+            def __enter__(self):
+                return self
+
+            def __exit__(self, *exc):
+                return False
+
+            def cursor(self):
+                return Cursor()
+
+            def commit(self):
+                pass
+
+        pin = Pin("pin", tracer=self.tracer)
+        conn = TracedConnection(ConnectionConnection(), pin)
+        with conn as conn2:
+            conn2.commit()
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 1
+
+        with conn as conn2:
+            with conn2.cursor() as cursor:
+                cursor.execute("query")
+                cursor.fetchall()
+
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 1
+
+        # If a cursor is returned from the context manager
+        # then it should be instrumented.
+
+        class ConnectionCursor(object):
+            def __enter__(self):
+                return Cursor()
+
+            def __exit__(self, *exc):
+                return False
+
+            def commit(self):
+                pass
+
+        with TracedConnection(ConnectionCursor(), pin) as cursor:
+            cursor.execute("query")
+            cursor.fetchall()
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 1
+
+        # If a traced cursor is returned then it should not
+        # be double instrumented.
+
+        class ConnectionTracedCursor(object):
+            def __enter__(self):
+                return self.cursor()
+
+            def __exit__(self, *exc):
+                return False
+
+            def cursor(self):
+                return TracedCursor(Cursor(), pin, {})
+
+            def commit(self):
+                pass
+
+        with TracedConnection(ConnectionTracedCursor(), pin) as cursor:
+            cursor.execute("query")
+            cursor.fetchall()
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 1
+
+        # Check when a different connection object is returned
+        # from a connection context manager.
+        # No traces should be produced.
+
+        other_conn = ConnectionConnection()
+
+        class ConnectionDifferentConnection(object):
+            def __enter__(self):
+                return other_conn
+
+            def __exit__(self, *exc):
+                return False
+
+            def cursor(self):
+                return Cursor()
+
+            def commit(self):
+                pass
+
+        conn = TracedConnection(ConnectionDifferentConnection(), pin)
+        with conn as conn2:
+            conn2.commit()
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 0
+
+        with conn as conn2:
+            with conn2.cursor() as cursor:
+                cursor.execute("query")
+                cursor.fetchall()
+
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 0
+
+        # When some unexpected value is returned from the context manager
+        # it should be handled gracefully.
+
+        class ConnectionUnknown(object):
+            def __enter__(self):
+                return 123456
+
+            def __exit__(self, *exc):
+                return False
+
+            def cursor(self):
+                return Cursor()
+
+            def commit(self):
+                pass
+
+        conn = TracedConnection(ConnectionDifferentConnection(), pin)
+        with conn as conn2:
+            conn2.commit()
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 0
+
+        with conn as conn2:
+            with conn2.cursor() as cursor:
+                cursor.execute("query")
+                cursor.fetchall()
+
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 0
+
+        # Errors should be the same when no context management is defined.
+
+        class ConnectionNoCtx(object):
+            def cursor(self):
+                return Cursor()
+
+            def commit(self):
+                pass
+
+        conn = TracedConnection(ConnectionNoCtx(), pin)
+        with pytest.raises(AttributeError):
+            with conn:
+                pass
+
+        with pytest.raises(AttributeError):
+            with conn as conn2:
+                pass
+
+        spans = self.tracer.writer.pop()
+        assert len(spans) == 0