Add client hooks for psycopg2 (#39)

Author: bitglue

opentracing_instrumentation/client_hooks/__init__.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2015 Uber Technologies, Inc.
+# Copyright (c) 2015-2017 Uber Technologies, Inc.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -34,6 +34,7 @@ def install_all_patches():
     If a specific module is not available on the path, it is ignored.
     """
     from . import mysqldb
+    from . import psycopg2
     from . import strict_redis
     from . import sqlalchemy
     from . import tornado_http
@@ -42,6 +43,7 @@ def install_all_patches():
     from . import requests
 
     mysqldb.install_patches()
+    psycopg2.install_patches()
     strict_redis.install_patches()
     sqlalchemy.install_patches()
     tornado_http.install_patches()

opentracing_instrumentation/client_hooks/_dbapi2.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2015 Uber Technologies, Inc.
+# Copyright (c) 2015-2017 Uber Technologies, Inc.
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
 # of this software and associated documentation files (the "Software"), to deal
@@ -115,6 +115,8 @@ def __call__(self, *args, **kwargs):
 
 
 class ConnectionWrapper(wrapt.ObjectProxy):
+    __slots__ = ('_module_name', '_connect_params')
+
     def __init__(self, connection, module_name, connect_params):
         super(ConnectionWrapper, self).__init__(wrapped=connection)
         self._module_name = module_name
@@ -174,6 +176,8 @@ def __exit__(self, exc, value, tb):
 
 
 class CursorWrapper(wrapt.ObjectProxy):
+    __slots__ = ('_module_name', '_connect_params', '_cursor_params')
+
     def __init__(self, cursor, module_name,
                  connect_params=None, cursor_params=None):
         super(CursorWrapper, self).__init__(wrapped=cursor)

opentracing_instrumentation/client_hooks/psycopg2.py

@@ -0,0 +1,47 @@
+# Copyright (c) 2017 Uber Technologies, Inc.
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+# THE SOFTWARE.
+
+from __future__ import absolute_import
+from ._dbapi2 import ContextManagerConnectionWrapper as ConnectionWrapper
+from ._dbapi2 import ConnectionFactory
+from ._singleton import singleton
+
+# Try to save the original entry points
+try:
+    import psycopg2
+except ImportError:  # pragma: no cover
+    pass
+else:
+    _psycopg2_connect = psycopg2.connect
+
+
+@singleton
+def install_patches():
+    try:
+        import psycopg2
+    except ImportError:  # pragma: no cover
+        return
+
+    factory = ConnectionFactory(connect_func=psycopg2.connect,
+                                module_name='psycopg2',
+                                conn_wrapper_ctor=ConnectionWrapper)
+    setattr(psycopg2, 'connect', factory)
+    if hasattr(psycopg2, 'Connect'):
+        setattr(psycopg2, 'Connect', factory)

opentracing_instrumentation/request_context.py

@@ -181,6 +181,9 @@ def span_in_context(span):
     ## Usage example in WSGI middleware:
 
     .. code-block:: python
+        from opentracing_instrumentation.http_server import WSGIRequestWrapper
+        from opentracing_instrumentation.http_server import before_request
+        from opentracing_instrumentation import request_context
 
         def create_wsgi_tracing_middleware(other_wsgi):