Merge pull request #264 from jonathangreen/sqlalchemy_core

Add support for SQLAlchemy Core

Author: willarmiros

aws_xray_sdk/core/patcher.py

@@ -24,6 +24,7 @@
     'pymysql',
     'psycopg2',
     'pg8000',
+    'sqlalchemy_core',
 )
 
 NO_DOUBLE_PATCH = (
@@ -37,6 +38,7 @@
     'pymysql',
     'psycopg2',
     'pg8000',
+    'sqlalchemy_core',
 )
 
 _PATCHED_MODULES = set()

aws_xray_sdk/ext/sqlalchemy_core/__init__.py

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

aws_xray_sdk/ext/sqlalchemy_core/patch.py

@@ -0,0 +1,86 @@
+import logging
+import sys
+if sys.version_info >= (3, 0, 0):
+    from urllib.parse import urlparse, uses_netloc
+else:
+    from urlparse import urlparse, uses_netloc
+
+import wrapt
+
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.patcher import _PATCHED_MODULES
+from aws_xray_sdk.core.utils import stacktrace
+from aws_xray_sdk.ext.util import unwrap
+
+
+def _sql_meta(instance, args):
+    try:
+        metadata = {}
+        url = urlparse(str(instance.engine.url))
+        # Add Scheme to uses_netloc or // will be missing from url.
+        uses_netloc.append(url.scheme)
+        if url.password is None:
+            metadata['url'] = url.geturl()
+            name = url.netloc
+        else:
+            # Strip password from URL
+            host_info = url.netloc.rpartition('@')[-1]
+            parts = url._replace(netloc='{}@{}'.format(url.username, host_info))
+            metadata['url'] = parts.geturl()
+            name = host_info
+        metadata['user'] = url.username
+        metadata['database_type'] = instance.engine.name
+        try:
+            version = getattr(instance.dialect, '{}_version'.format(instance.engine.driver))
+            version_str = '.'.join(map(str, version))
+            metadata['driver_version'] = "{}-{}".format(instance.engine.driver, version_str)
+        except AttributeError:
+            metadata['driver_version'] = instance.engine.driver
+        if instance.dialect.server_version_info is not None:
+            metadata['database_version'] = '.'.join(map(str, instance.dialect.server_version_info))
+        if xray_recorder.stream_sql:
+            metadata['sanitized_query'] = str(args[0])
+    except Exception:
+        metadata = None
+        name = None
+        logging.getLogger(__name__).exception('Error parsing sql metadata.')
+    return name, metadata
+
+
+def _xray_traced_sqlalchemy_execute(wrapped, instance, args, kwargs):
+    name, sql = _sql_meta(instance, args)
+    if sql is not None:
+        subsegment = xray_recorder.begin_subsegment(name, namespace='remote')
+    else:
+        subsegment = None
+    try:
+        res = wrapped(*args, **kwargs)
+    except Exception:
+        if subsegment is not None:
+            exception = sys.exc_info()[1]
+            stack = stacktrace.get_stacktrace(limit=xray_recorder._max_trace_back)
+            subsegment.add_exception(exception, stack)
+        raise
+    finally:
+        if subsegment is not None:
+            subsegment.set_sql(sql)
+            xray_recorder.end_subsegment()
+    return res
+
+
+def patch():
+    wrapt.wrap_function_wrapper(
+        'sqlalchemy.engine.base',
+        'Connection.execute',
+        _xray_traced_sqlalchemy_execute
+    )
+
+
+def unpatch():
+    """
+    Unpatch any previously patched modules.
+    This operation is idempotent.
+    """
+    _PATCHED_MODULES.discard('sqlalchemy_core')
+    import sqlalchemy
+    unwrap(sqlalchemy.engine.base.Connection, 'execute')

tests/ext/sqlalchemy_core/__init__.py

@@ -0,0 +1,108 @@
+from __future__ import absolute_import
+
+import pytest
+from sqlalchemy import create_engine, Column, Integer, String
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+from sqlalchemy.sql.expression import Insert, Delete
+
+from aws_xray_sdk.core import xray_recorder, patch
+from aws_xray_sdk.core.context import Context
+
+Base = declarative_base()
+
+
+class User(Base):
+        __tablename__ = 'users'
+
+        id = Column(Integer, primary_key=True)
+        name = Column(String)
+        fullname = Column(String)
+        password = Column(String)
+
+
+@pytest.fixture()
+def engine():
+    """
+    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.
+    """
+    from aws_xray_sdk.ext.sqlalchemy_core import unpatch
+    patch(('sqlalchemy_core',))
+    engine = create_engine('sqlite:///:memory:')
+    xray_recorder.configure(service='test', sampling=False, context=Context())
+    xray_recorder.begin_segment('name')
+    Base.metadata.create_all(engine)
+    xray_recorder.clear_trace_entities()
+    xray_recorder.begin_segment('name')
+    yield engine
+    xray_recorder.clear_trace_entities()
+    unpatch()
+
+
+@pytest.fixture()
+def connection(engine):
+    return engine.connect()
+
+
+@pytest.fixture()
+def session(engine):
+    Session = sessionmaker(bind=engine)
+    return Session()
+
+
+def test_all(session):
+    """ Test calling all() on get all records.
+    Verify we run the query and return the SQL as metdata"""
+    session.query(User).all()
+    assert len(xray_recorder.current_segment().subsegments) == 1
+    sql_meta = xray_recorder.current_segment().subsegments[0].sql
+    assert sql_meta['url'] == 'sqlite:///:memory:'
+    assert sql_meta['sanitized_query'].startswith('SELECT')
+    assert sql_meta['sanitized_query'].endswith('FROM users')
+
+
+def test_add(session):
+    """ Test calling add() on insert a row.
+    Verify we that we capture trace for the add"""
+    password = "123456"
+    john = User(name='John', fullname="John Doe", password=password)
+    session.add(john)
+    session.commit()
+    assert len(xray_recorder.current_segment().subsegments) == 1
+    sql_meta = xray_recorder.current_segment().subsegments[0].sql
+    assert sql_meta['sanitized_query'].startswith('INSERT INTO users')
+    assert password not in sql_meta['sanitized_query']
+
+
+def test_filter_first(session):
+    """ Test calling filter().first() on get first filtered records.
+    Verify we run the query and return the SQL as metdata"""
+    session.query(User).filter(User.password=="mypassword!").first()
+    assert len(xray_recorder.current_segment().subsegments) == 1
+    sql_meta = xray_recorder.current_segment().subsegments[0].sql
+    assert sql_meta['sanitized_query'].startswith('SELECT')
+    assert 'FROM users' in sql_meta['sanitized_query']
+    assert "mypassword!" not in sql_meta['sanitized_query']
+
+
+def test_connection_add(connection):
+    password = "123456"
+    statement = Insert(User).values(name='John', fullname="John Doe", password=password)
+    connection.execute(statement)
+    assert len(xray_recorder.current_segment().subsegments) == 1
+    sql_meta = xray_recorder.current_segment().subsegments[0].sql
+    assert sql_meta['sanitized_query'].startswith('INSERT INTO users')
+    assert sql_meta['url'] == 'sqlite:///:memory:'
+    assert password not in sql_meta['sanitized_query']
+
+def test_connection_query(connection):
+    password = "123456"
+    statement = Delete(User).where(User.name == 'John').where(User.password == password)
+    connection.execute(statement)
+    assert len(xray_recorder.current_segment().subsegments) == 1
+    sql_meta = xray_recorder.current_segment().subsegments[0].sql
+    assert sql_meta['sanitized_query'].startswith('DELETE FROM users')
+    assert sql_meta['url'] == 'sqlite:///:memory:'
+    assert password not in sql_meta['sanitized_query']