ISSUE-80: Add patch support for psycopg2

aws_xray_sdk/core/patcher.py

@@ -10,6 +10,7 @@
     'mysql',
     'httplib',
     'pymongo',
+    'psycopg2',
 )
 
 NO_DOUBLE_PATCH = (
@@ -18,6 +19,7 @@
     'sqlite3',
     'mysql',
     '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,31 @@
+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)
+    host = kwargs['host'] if 'host' in kwargs else re.search(r'host=(\S+)\b', args[0]).groups()[0]
+    dbname = kwargs['dbname'] if 'dbname' in kwargs else re.search(r'dbname=(\S+)\b', args[0]).groups()[0]
+    port = kwargs['port'] if 'port' in kwargs else re.search(r'port=(\S+)\b', args[0]).groups()[0]
+    user = kwargs['user'] if 'user' in kwargs else re.search(r'user=(\S+)\b', args[0]).groups()[0]
+    meta = {
+        'database_type': 'PostgreSQL',
+        'url': 'postgresql://{}@{}:{}/{}'.format(user, host, port, dbname),
+        'user': user,
+        'database_version': str(conn.server_version),
+        'driver_version': 'Psycopg 2'
+    }
+
+    return XRayTracedConn(conn, meta)

tests/ext/psycopg2/test_psycopg2.py

@@ -0,0 +1,120 @@
+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:
+        url = postgresql.url()
+        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['user'] == dsn['user']
+    assert sql['url'] == url
+    assert sql['database_version']
+
+
+def test_execute_dsn_string():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        url = postgresql.url()
+        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['user'] == dsn['user']
+    assert sql['url'] == url
+    assert sql['database_version']
+
+
+def test_execute_in_pool():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        url = postgresql.url()
+        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['user'] == dsn['user']
+    assert sql['url'] == url
+    assert sql['database_version']
+
+
+def test_execute_bad_query():
+    q = 'SELECT blarg'
+    with testing.postgresql.Postgresql() as postgresql:
+        url = postgresql.url()
+        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['user'] == dsn['user']
+    assert sql['url'] == url
+    assert sql['database_version']
+
+    exception = subsegment.cause['exceptions'][0]
+    assert exception.type == 'ProgrammingError'

tox.ini

@@ -19,6 +19,8 @@ deps =
     # the sdk doesn't support earlier version of django
     django >= 1.10, <2.0
     pynamodb
+    psycopg2
+    testing.postgresql
 
     # Python2 only deps
     py{27}: enum34