ISSUE-80: Add patch support for psycopg2

aws_xray_sdk/core/patcher.py

@@ -12,6 +12,7 @@
     'mysql',
     'httplib',
     'pymongo',
+    'psycopg2'
 )
 
 _PATCHED_MODULES = set()

aws_xray_sdk/ext/psycopg2/__init__.py

@@ -0,0 +1,4 @@
+from .patch import patch
+
+
+__all__ = ['patch']

aws_xray_sdk/ext/psycopg2/patch.py

@@ -0,0 +1,26 @@
+import re
+import wrapt
+
+from aws_xray_sdk.ext.dbapi2 import XRayTracedConn
+
+
+def patch():
+
+    wrapt.wrap_function_wrapper(
+        'psycopg2',
+        'connect',
+        _xray_traced_connect
+    )
+
+
+def _xray_traced_connect(wrapped, instance, args, kwargs):
+
+    conn = wrapped(*args, **kwargs)
+    meta = {
+        'database_type': 'postgresql',
+        'database_host': kwargs['host'] if 'host' in kwargs else re.search(r'host=(\S+)\b', args[0]).groups()[0],
+        'database_name': kwargs['dbname'] if 'dbname' in kwargs else re.search(r'dbname=(\S+)\b', args[0]).groups()[0],
+        'database_user': kwargs['user'] if 'user' in kwargs else re.search(r'user=(\S+)\b', args[0]).groups()[0],
+    }
+
+    return XRayTracedConn(conn, meta)

tests/ext/psycopg2/test_psycopg2.py

@@ -0,0 +1,117 @@
+import psycopg2
+import psycopg2.pool
+
+import pytest
+import testing.postgresql
+
+from aws_xray_sdk.core import patch
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+
+patch(('psycopg2',))
+
+
+@pytest.fixture(autouse=True)
+def construct_ctx():
+    """
+    Clean up context storage on each test run and begin a segment
+    so that later subsegment can be attached. After each test run
+    it cleans up context storage again.
+    """
+    xray_recorder.configure(service='test', sampling=False, context=Context())
+    xray_recorder.clear_trace_entities()
+    xray_recorder.begin_segment('name')
+    yield
+    xray_recorder.clear_trace_entities()
+
+
+def test_execute_dsn_kwargs():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        conn = psycopg2.connect(dbname=dsn['database'],
+                                user=dsn['user'],
+                                password='',
+                                host=dsn['host'],
+                                port=dsn['port'])
+        cur = conn.cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'postgresql'
+    assert sql['database_host'] == dsn['host']
+    assert sql['database_name'] == dsn['database']
+    assert sql['database_user'] == dsn['user']
+
+
+def test_execute_dsn_string():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        conn = psycopg2.connect('dbname=' + dsn['database'] +
+                                ' password=mypassword' +
+                                ' host=' + dsn['host'] +
+                                ' port=' + str(dsn['port']) +
+                                ' user=' + dsn['user'])
+        cur = conn.cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'postgresql'
+    assert sql['database_host'] == dsn['host']
+    assert sql['database_name'] == dsn['database']
+    assert sql['database_user'] == dsn['user']
+
+
+
+def test_execute_in_pool():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        pool = psycopg2.pool.SimpleConnectionPool(1, 1,
+                                                  dbname=dsn['database'],
+                                                  user=dsn['user'],
+                                                  password='',
+                                                  host=dsn['host'],
+                                                  port=dsn['port'])
+        cur = pool.getconn(key=dsn['user']).cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'postgresql'
+    assert sql['database_host'] == dsn['host']
+    assert sql['database_name'] == dsn['database']
+    assert sql['database_user'] == dsn['user']
+
+
+def test_execute_bad_query():
+    q = 'SELECT blarg'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        conn = psycopg2.connect(dbname=dsn['database'],
+                                user=dsn['user'],
+                                password='',
+                                host=dsn['host'],
+                                port=dsn['port'])
+        cur = conn.cursor()
+        try:
+            cur.execute(q)
+        except Exception:
+            pass
+
+    subsegment = xray_recorder.current_segment().subsegments[0]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'postgresql'
+    assert sql['database_host'] == dsn['host']
+    assert sql['database_name'] == dsn['database']
+    assert sql['database_user'] == dsn['user']
+
+    exception = subsegment.cause['exceptions'][0]
+    assert exception.type == 'ProgrammingError'

tox.ini

@@ -20,6 +20,8 @@ deps =
     # the sdk doesn't support earlier version of django
     django >= 1.10, <2.0
     pynamodb
+    psycopg2
+    testing.postgresql
     # Python3.5+ only deps
     py{35,36}: aiohttp >= 3.0.0
     py{35,36}: pytest-aiohttp