Fixing broken instrumentation for sqlalchemy >= 1.4.0

aws_xray_sdk/ext/sqlalchemy/util/decorators.py

@@ -1,4 +1,6 @@
 import re
+import types
+
 from aws_xray_sdk.core import xray_recorder
 from aws_xray_sdk.ext.util import strip_url
 from future.standard_library import install_aliases
@@ -13,7 +15,7 @@ def decorator(cls):
             for name, obj in vars(c).items():
                 if name.startswith("_"):
                     continue
-                if callable(obj):
+                if isinstance(obj, types.FunctionType):
                     try:
                         obj = obj.__func__  # unwrap Python 2 unbound method
                     except AttributeError:

aws_xray_sdk/ext/sqlalchemy_core/patch.py

@@ -1,5 +1,6 @@
 import logging
 import sys
+
 if sys.version_info >= (3, 0, 0):
     from urllib.parse import urlparse, uses_netloc
 else:
@@ -13,10 +14,10 @@
 from aws_xray_sdk.ext.util import unwrap
 
 
-def _sql_meta(instance, args):
+def _sql_meta(engine_instance, args):
     try:
         metadata = {}
-        url = urlparse(str(instance.engine.url))
+        url = urlparse(str(engine_instance.engine.url))
         # Add Scheme to uses_netloc or // will be missing from url.
         uses_netloc.append(url.scheme)
         if url.password is None:
@@ -29,17 +30,20 @@ def _sql_meta(instance, args):
             metadata['url'] = parts.geturl()
             name = host_info
         metadata['user'] = url.username
-        metadata['database_type'] = instance.engine.name
+        metadata['database_type'] = engine_instance.engine.name
         try:
-            version = getattr(instance.dialect, '{}_version'.format(instance.engine.driver))
+            version = getattr(engine_instance.dialect, '{}_version'.format(engine_instance.engine.driver))
             version_str = '.'.join(map(str, version))
-            metadata['driver_version'] = "{}-{}".format(instance.engine.driver, version_str)
+            metadata['driver_version'] = "{}-{}".format(engine_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))
+            metadata['driver_version'] = engine_instance.engine.driver
+        if engine_instance.dialect.server_version_info is not None:
+            metadata['database_version'] = '.'.join(map(str, engine_instance.dialect.server_version_info))
         if xray_recorder.stream_sql:
-            metadata['sanitized_query'] = str(args[0])
+            try:
+                metadata['sanitized_query'] = str(args[0])
+            except Exception:
+                logging.getLogger(__name__).exception('Error getting the sanitized query')
     except Exception:
         metadata = None
         name = None
@@ -48,7 +52,15 @@ def _sql_meta(instance, args):
 
 
 def _xray_traced_sqlalchemy_execute(wrapped, instance, args, kwargs):
-    name, sql = _sql_meta(instance, args)
+    return _fetch_sql_metadata_and_process_request(wrapped, instance, args, kwargs)
+
+
+def _xray_traced_sqlalchemy_session(wrapped, instance, args, kwargs):
+    return _fetch_sql_metadata_and_process_request(wrapped, instance.bind, args, kwargs)
+
+
+def _fetch_sql_metadata_and_process_request(wrapped, engine_instance, args, kwargs):
+    name, sql = _sql_meta(engine_instance, args)
     if sql is not None:
         subsegment = xray_recorder.begin_subsegment(name, namespace='remote')
     else:
@@ -75,6 +87,12 @@ def patch():
         _xray_traced_sqlalchemy_execute
     )
 
+    wrapt.wrap_function_wrapper(
+        'sqlalchemy.orm.session',
+        'Session.execute',
+        _xray_traced_sqlalchemy_session
+    )
+
 
 def unpatch():
     """
@@ -84,3 +102,4 @@ def unpatch():
     _PATCHED_MODULES.discard('sqlalchemy_core')
     import sqlalchemy
     unwrap(sqlalchemy.engine.base.Connection, 'execute')
+    unwrap(sqlalchemy.orm.session.Session, 'execute')

tests/ext/sqlalchemy_core/test_base.py

@@ -0,0 +1,51 @@
+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 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()

tests/ext/sqlalchemy_core/test_sqlalchemy_core.py

@@ -1,56 +1,6 @@
-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 .test_base import User, session, engine, connection
 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()
-
+from aws_xray_sdk.core import xray_recorder
 
 def test_all(session):
     """ Test calling all() on get all records.
@@ -63,19 +13,6 @@ def test_all(session):
     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"""
@@ -97,6 +34,7 @@ def test_connection_add(connection):
     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)
@@ -105,4 +43,4 @@ def test_connection_query(connection):
     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']
+    assert password not in sql_meta['sanitized_query']

tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py

@@ -0,0 +1,15 @@
+from .test_base import User, session, engine, connection
+from sqlalchemy.sql.expression import select
+from aws_xray_sdk.core import xray_recorder
+
+# 2.0 style execution test. see https://docs.sqlalchemy.org/en/14/changelog/migration_14.html#orm-query-is-internally
+# -unified-with-select-update-delete-2-0-style-execution-available
+def test_orm_style_select_execution(session):
+    statement = select(User).where(
+        User.name == 'John'
+    )
+    session.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('SELECT')
+    assert 'FROM users' in sql_meta['sanitized_query']

tox.ini

@@ -6,6 +6,8 @@ envlist =
     py{36,37,38,39}-django30
     py{36,37,38,39}-django31
     py{36,37,38,39}-django32
+    py{27,36,37,38,39}-sqlalchemy
+    py{34,35}-sqlalchemy
     coverage-report
 
 skip_missing_interpreters = True
@@ -19,8 +21,8 @@ deps =
     requests
     bottle >= 0.10
     flask >= 0.10
-    sqlalchemy==1.3.*
-    Flask-SQLAlchemy==2.4.*
+    sqlalchemy
+    Flask-SQLAlchemy
     future
     django22: Django==2.2.*
     django30: Django==3.0.*
@@ -51,8 +53,10 @@ deps =
     py{35,36,37,38,39}: aiobotocore >= 0.10.0
 
 commands =
-    py{27,34}-default: coverage run --source aws_xray_sdk -m py.test tests --ignore tests/ext/aiohttp --ignore tests/ext/aiobotocore --ignore tests/ext/django --ignore tests/test_async_local_storage.py --ignore tests/test_async_recorder.py
-    py{35,36,37,38,39}-default: coverage run --source aws_xray_sdk -m py.test --ignore tests/ext/django tests
+    py{27,34}-default: coverage run --source aws_xray_sdk -m py.test tests --ignore tests/ext/aiohttp --ignore tests/ext/aiobotocore --ignore tests/ext/django --ignore tests/test_async_local_storage.py --ignore tests/test_async_recorder.py --ignore tests/ext/sqlalchemy_core
+    py{35,36,37,38,39}-default: coverage run --source aws_xray_sdk -m py.test tests --ignore tests/ext/django --ignore tests/ext/sqlalchemy_core
+    py{27,36,37,38,39}-default: coverage run --source aws_xray_sdk -m py.test tests/ext/sqlalchemy_core
+    py{34,35}-default: coverage run --source aws_xray_sdk -m py.test tests/ext/sqlalchemy_core/ --ignore tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py
     django{22,30,31,32}: coverage run --source aws_xray_sdk -m py.test tests/ext/django
     codecov