diff --git a/aws_xray_sdk/ext/sqlalchemy/util/decorators.py b/aws_xray_sdk/ext/sqlalchemy/util/decorators.py index d0fe5e02..d82ef32d 100644 --- a/aws_xray_sdk/ext/sqlalchemy/util/decorators.py +++ b/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: diff --git a/aws_xray_sdk/ext/sqlalchemy_core/patch.py b/aws_xray_sdk/ext/sqlalchemy_core/patch.py index df7a8a67..5dd9ff62 100644 --- a/aws_xray_sdk/ext/sqlalchemy_core/patch.py +++ b/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 _process_request(wrapped, instance, args, kwargs) + + +def _xray_traced_sqlalchemy_session(wrapped, instance, args, kwargs): + return _process_request(wrapped, instance.bind, args, kwargs) + + +def _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') diff --git a/tests/ext/sqlalchemy_core/test_base.py b/tests/ext/sqlalchemy_core/test_base.py new file mode 100644 index 00000000..b21cf08c --- /dev/null +++ b/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() diff --git a/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py b/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py index 05368968..4dbac2b6 100644 --- a/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py +++ b/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'] \ No newline at end of file + assert password not in sql_meta['sanitized_query'] diff --git a/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py b/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py new file mode 100644 index 00000000..fe44d325 --- /dev/null +++ b/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'] diff --git a/tox.ini b/tox.ini index ae7a19e0..8faee34d 100644 --- a/tox.ini +++ b/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