Add patch support for pg8000 (Pure python driver) (#115)

* Add patch support for pg8000

* pg8000: add unpatch

aws_xray_sdk/core/patcher.py

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

aws_xray_sdk/ext/pg8000/__init__.py

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

aws_xray_sdk/ext/pg8000/patch.py

@@ -0,0 +1,41 @@
+import pg8000
+import wrapt
+
+from aws_xray_sdk.ext.dbapi2 import XRayTracedConn
+from aws_xray_sdk.core.patcher import _PATCHED_MODULES
+from aws_xray_sdk.ext.util import unwrap
+
+
+def patch():
+
+    wrapt.wrap_function_wrapper(
+        'pg8000',
+        'connect',
+        _xray_traced_connect
+    )
+
+
+def _xray_traced_connect(wrapped, instance, args, kwargs):
+
+    conn = wrapped(*args, **kwargs)
+    meta = {
+        'database_type': 'PostgreSQL',
+        'user': conn.user.decode('utf-8'),
+        'driver_version': 'Pg8000'
+    }
+
+    if hasattr(conn, '_server_version'):
+        version = getattr(conn, '_server_version')
+        if version:
+            meta['database_version'] = str(version)
+
+    return XRayTracedConn(conn, meta)
+
+
+def unpatch():
+    """
+    Unpatch any previously patched modules.
+    This operation is idempotent.
+    """
+    _PATCHED_MODULES.discard('pg8000')
+    unwrap(pg8000, 'connect')

tests/ext/pg8000/test_pg8000.py

@@ -0,0 +1,71 @@
+import pg8000
+
+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
+from aws_xray_sdk.ext.pg8000 import unpatch
+
+
+@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.
+    """
+    patch(('pg8000',))
+    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()
+    unpatch()
+
+
+def test_execute_dsn_kwargs():
+    q = 'SELECT 1'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        conn = pg8000.connect(database=dsn['database'],
+                              user=dsn['user'],
+                              password='',
+                              host=dsn['host'],
+                              port=dsn['port'])
+        cur = conn.cursor()
+        cur.execute(q)
+
+    subsegment = xray_recorder.current_segment().subsegments[-1]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'PostgreSQL'
+    assert sql['user'] == dsn['user']
+    assert sql['database_version']
+
+
+def test_execute_bad_query():
+    q = 'SELECT blarg'
+    with testing.postgresql.Postgresql() as postgresql:
+        dsn = postgresql.dsn()
+        conn = pg8000.connect(database=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[-1]
+    assert subsegment.name == 'execute'
+    sql = subsegment.sql
+    assert sql['database_type'] == 'PostgreSQL'
+    assert sql['user'] == dsn['user']
+    assert sql['database_version']
+
+    exception = subsegment.cause['exceptions'][0]
+    assert exception.type == 'ProgrammingError'

tests/ext/psycopg2/test_psycopg2.py

@@ -157,4 +157,4 @@ def test_register_extensions():
                                 ' port=' + str(dsn['port']) +
                                 ' user=' + dsn['user'])
         assert psycopg2.extras.register_uuid(None, conn)
-        assert psycopg2.extras.register_uuid(None, conn.cursor())
+        assert psycopg2.extras.register_uuid(None, conn.cursor())

tox.ini

@@ -20,6 +20,7 @@ deps =
     django >= 1.10, <2.0
     pynamodb >= 3.3.1
     psycopg2
+    pg8000
     testing.postgresql
 
     # Python2 only deps