Add psycopg2 wrapper

Closes #344
iopipe · kolanos · Aug 16, 2019 · Aug 14, 2019 · Aug 15, 2019 · Aug 15, 2019
commit 8d2505c8b82c2bfc8c3278983bd4d21d37e6e94f
diff --git a/iopipe/contrib/trace/auto_db.py b/iopipe/contrib/trace/auto_db.py
@@ -2,13 +2,18 @@
 import uuid
 import wrapt
 
+from .dbapi import AdapterProxy, ConnectionProxy, CursorProxy
 from .util import ensure_utf8
 
 Request = collections.namedtuple(
     "Request", ["command", "key", "hostname", "port", "connectionName", "db", "table"]
 )
 
 
+def collect_psycopg2_metrics(context, trace, instance):
+    pass
+
+
 def collect_pymongo_metrics(context, trace, instance, response):
     from pymongo.cursor import Cursor
     from pymongo.results import (
@@ -126,6 +131,69 @@ def wrapper(wrapped, instance, args, kwargs):
             )
 
 
+def patch_psycopg2(context):
+    """
+    Monkey patches psycopg2 client, if available. Overloads the
+    execute method to add tracing and metrics collection.
+    """
+
+    class PGCursorProxy(CursorProxy):
+        def execute(self, *args, **kwargs):
+            if not hasattr(context, "iopipe") or not hasattr(
+                context.iopipe, "mark"
+            ):  # pragma: no cover
+                self.__wrapped__.execute(*args, **kwargs)
+
+            id = ensure_utf8(str(uuid.uuid4()))
+            with context.iopipe.mark(id):
+                self.__wrapped__.execute(*args, **kwargs)
+            trace = context.iopipe.mark.measure(id)
+            context.iopipe.mark.delete(id)
+            collect_psycopg2_metrics(context, trace, self)
+
+    class PGConnectionProxy(ConnectionProxy):
+        def cursor(self, *args, **kwargs):
+            cursor = self.__wrapped__.cursor(*args, **kwargs)
+            return PGCursorProxy(cursor, self)
+
+    def adapt_wrapper(wrapped, instance, args, kwargs):
+        adapter = wrapped(*args, **kwargs)
+        return AdapterProxy(adapter) if hasattr(adapter, "prepare") else adapter
+
+    def connect_wrapper(wrapped, instance, args, kwargs):
+        connection = wrapped(*args, **kwargs)
+        return PGConnectionProxy(connection, args, kwargs)
+
+    def register_type_wrapper(wrapped, instance, args, kwargs):
+        def _extract_arguments(obj, scope=None):
+            return obj, scope
+
+        obj, scope = _extract_arguments(*args, **kwargs)
+
+        if scope is not None:
+            if isinstance(scope, wrapt.ObjectProxy):
+                scope = scope.__wrapped__
+            return wrapped(obj, scope)
+
+        return wrapped(obj)
+
+    try:
+        wrapt.wrap_function_wrapper("psycopg2", "connect", connect_wrapper)
+    except Exception:  # pragma: no cover
+        pass
+    else:
+        wrapt.wrap_function_wrapper("psycopg2.extensions", "adapt", adapt_wrapper)
+        wrapt.wrap_function_wrapper(
+            "psycopg2.extensions", "register_type", register_type_wrapper
+        )
+        wrapt.wrap_function_wrapper(
+            "psycopg2._psycopg", "register_type", register_type_wrapper
+        )
+        wrapt.wrap_function_wrapper(
+            "psycopg2._json", "register_type", register_type_wrapper
+        )
+
+
 def patch_redis(context):
     """
     Monkey patches redis client, if available. Overloads the
@@ -179,6 +247,11 @@ def pipeline_wrapper(wrapped, instance, args, kwargs):  # pragma: no cover
             wrapt.wrap_function_wrapper(module_name, class_method, pipeline_wrapper)
 
 
+def restore_psycopg2():
+    """Restores psycopg2"""
+    pass
+
+
 def restore_pymongo():
     """Restores pymongo"""
     try:
@@ -224,10 +297,12 @@ def patch_db_requests(context):
     if not hasattr(context, "iopipe"):
         return
 
+    patch_psycopg2(context)
     patch_pymongo(context)
     patch_redis(context)
 
 
 def restore_db_requests():
+    restore_psycopg2()
     restore_pymongo()
     restore_redis()
diff --git a/iopipe/contrib/trace/dbapi.py b/iopipe/contrib/trace/dbapi.py
@@ -0,0 +1,48 @@
+import wrapt
+
+
+class CursorProxy(wrapt.ObjectProxy):
+    def __init__(self, cursor, connection_proxy):
+        super(CursorProxy, self).__init__(cursor)
+
+        self._self_connection = connection_proxy
+
+    @property
+    def connection_proxy(self):
+        return self._self_connection
+
+    def execute(self, *args, **kwargs):
+        self.__wrapped__.execute(*args, **kwargs)
+
+
+class ConnectionProxy(wrapt.ObjectProxy):
+    def __init__(self, connection, args, kwargs):
+        super(ConnectionProxy, self).__init__(connection)
+
+        self._self_args = args
+        self._self_kwargs = kwargs
+
+    def cursor(self, *args, **kwargs):
+        cursor = self.__wrapped__.cursor(*args, **kwargs)
+        return CursorProxy(cursor, self)
+
+    @property
+    def extract_hostname(self):
+        return self._self_kwargs.get("host", "localhost")
+
+    @property
+    def extract_dbname(self):
+        return self._self_kwargs.get("db", self._self_kwargs.get("database", ""))
+
+
+class AdapterProxy(wrapt.ObjectProxy):
+    def prepare(self, *args, **kwargs):
+        if not args:
+            return self.__wrapped__.prepare(*args, **kwargs)
+
+        connection = args[0]
+
+        if isinstance(connection, wrapt.ObjectProxy):
+            connection = connection.__wrapped__
+
+        return self.__wrapped__.prepare(connection, *args[1:], **kwargs)
diff --git a/setup.py b/setup.py
@@ -21,6 +21,7 @@
     "mock",
     "mongomock==3.17.0",
     "more-itertools<6.0.0",
+    "psycopg2==2.8.3",
     "pymongo==3.8.0",
     "pytest==4.1.0",
     "pytest-benchmark==3.2.0",