diff --git a/.travis.yml b/.travis.yml index 3e64f01f..2970fcde 100644 --- a/.travis.yml +++ b/.travis.yml @@ -9,6 +9,9 @@ python: - "3.8" - "3.9" +addons: + postgresql: "9.6" + install: - pip install tox tox-travis diff --git a/aws_xray_sdk/ext/sqlalchemy_core/patch.py b/aws_xray_sdk/ext/sqlalchemy_core/patch.py index 5dd9ff62..acab1fd4 100644 --- a/aws_xray_sdk/ext/sqlalchemy_core/patch.py +++ b/aws_xray_sdk/ext/sqlalchemy_core/patch.py @@ -13,6 +13,8 @@ from aws_xray_sdk.core.utils import stacktrace from aws_xray_sdk.ext.util import unwrap +from sqlalchemy.sql.expression import ClauseElement + def _sql_meta(engine_instance, args): try: @@ -41,7 +43,10 @@ def _sql_meta(engine_instance, args): metadata['database_version'] = '.'.join(map(str, engine_instance.dialect.server_version_info)) if xray_recorder.stream_sql: try: - metadata['sanitized_query'] = str(args[0]) + if isinstance(args[0], ClauseElement): + metadata['sanitized_query'] = str(args[0].compile(engine_instance.engine)) + else: + metadata['sanitized_query'] = str(args[0]) except Exception: logging.getLogger(__name__).exception('Error getting the sanitized query') except Exception: diff --git a/tests/ext/sqlalchemy_core/test_base.py b/tests/ext/sqlalchemy_core/test_base.py index b21cf08c..74bbe976 100644 --- a/tests/ext/sqlalchemy_core/test_base.py +++ b/tests/ext/sqlalchemy_core/test_base.py @@ -21,7 +21,12 @@ class User(Base): @pytest.fixture() -def engine(): +def db_url(): + return 'sqlite:///:memory:' + + +@pytest.fixture() +def engine(db_url): """ Clean up context storage on each test run and begin a segment so that later subsegment can be attached. After each test run @@ -29,7 +34,7 @@ def engine(): """ from aws_xray_sdk.ext.sqlalchemy_core import unpatch patch(('sqlalchemy_core',)) - engine = create_engine('sqlite:///:memory:') + engine = create_engine(db_url) xray_recorder.configure(service='test', sampling=False, context=Context()) xray_recorder.begin_segment('name') Base.metadata.create_all(engine) diff --git a/tests/ext/sqlalchemy_core/test_postgres.py b/tests/ext/sqlalchemy_core/test_postgres.py new file mode 100644 index 00000000..5dde9a8e --- /dev/null +++ b/tests/ext/sqlalchemy_core/test_postgres.py @@ -0,0 +1,57 @@ +import pytest + +from .test_base import connection, engine, session, User + +from sqlalchemy import create_engine +from sqlalchemy.dialects.postgresql import insert as pg_insert + +from aws_xray_sdk.core import xray_recorder, patch +from aws_xray_sdk.core.context import Context + +import testing.postgresql + + +@pytest.fixture() +def postgres_db(): + with testing.postgresql.Postgresql() as postgresql: + yield postgresql + + +@pytest.fixture() +def db_url(postgres_db): + return postgres_db.url() + + +@pytest.fixture() +def sanitized_db_url(postgres_db): + dsn = postgres_db.dsn() + return 'postgresql://{user}@{host}:{port}/{db}'.format( + user=dsn['user'], + host=dsn['host'], + port=dsn['port'], + db=dsn['database'], + ) + + +def test_all(session, sanitized_db_url): + """ 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'] == sanitized_db_url + assert sql_meta['sanitized_query'].startswith('SELECT') + assert sql_meta['sanitized_query'].endswith('FROM users') + + +def test_insert_on_conflict_renders(connection): + statement = pg_insert(User).values(name='John', fullname="John Doe", password='123456') + statement = statement.on_conflict_do_nothing() + + 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 'ON CONFLICT DO NOTHING' in sql_meta['sanitized_query'] diff --git a/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py b/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py index 4dbac2b6..10cfaca7 100644 --- a/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py +++ b/tests/ext/sqlalchemy_core/test_sqlalchemy_core.py @@ -1,4 +1,4 @@ -from .test_base import User, session, engine, connection +from .test_base import User, session, db_url, engine, connection from sqlalchemy.sql.expression import Insert, Delete from aws_xray_sdk.core import xray_recorder diff --git a/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py b/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py index fe44d325..3f3f0985 100644 --- a/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py +++ b/tests/ext/sqlalchemy_core/test_sqlalchemy_core_2.py @@ -1,4 +1,4 @@ -from .test_base import User, session, engine, connection +from .test_base import User, session, db_url, engine, connection from sqlalchemy.sql.expression import select from aws_xray_sdk.core import xray_recorder