From a6ab1936198f0e7cde1d9af86cac820b697ed942 Mon Sep 17 00:00:00 2001 From: Anthony King Date: Wed, 28 Apr 2021 16:56:24 +0100 Subject: [PATCH 1/5] chore(sqla_core): setup an area to run postgres tests --- tests/ext/sqlalchemy_core/test_postgres.py | 58 ++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 tests/ext/sqlalchemy_core/test_postgres.py diff --git a/tests/ext/sqlalchemy_core/test_postgres.py b/tests/ext/sqlalchemy_core/test_postgres.py new file mode 100644 index 00000000..0dbfc773 --- /dev/null +++ b/tests/ext/sqlalchemy_core/test_postgres.py @@ -0,0 +1,58 @@ +import pytest + +from .test_base import connection, User, session, Base + +from sqlalchemy import create_engine + +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 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'], + ) + + +@pytest.fixture() +def engine(postgres_db): + """ + 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(postgres_db.url()) + 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() + + +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') From b6aac11637bca83afc88e879781f4b6930173141 Mon Sep 17 00:00:00 2001 From: Anthony King Date: Wed, 28 Apr 2021 16:57:25 +0100 Subject: [PATCH 2/5] fix: run a supported version of postgres in travis by default it runs 9.4. We need to use a feature added in 9.5 --- .travis.yml | 3 +++ 1 file changed, 3 insertions(+) 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 From ef5056fd462468f6eab9f8cd25362f9639ea037a Mon Sep 17 00:00:00 2001 From: Anthony King Date: Wed, 28 Apr 2021 17:01:17 +0100 Subject: [PATCH 3/5] feat(sqla-core): Add support for rendering Database Specific queries This makes the sqla-core plugin aware of how to compile a database specific query. It seems in some cases, we see strings come through (mostly pragma calls when doing 'create_all'), so do a check to see if the object looks like a SQL Expression object. Fixes: #290 --- aws_xray_sdk/ext/sqlalchemy_core/patch.py | 7 ++++++- tests/ext/sqlalchemy_core/test_postgres.py | 14 ++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) 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_postgres.py b/tests/ext/sqlalchemy_core/test_postgres.py index 0dbfc773..1b8f2b4f 100644 --- a/tests/ext/sqlalchemy_core/test_postgres.py +++ b/tests/ext/sqlalchemy_core/test_postgres.py @@ -3,6 +3,7 @@ from .test_base import connection, User, session, Base 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 @@ -56,3 +57,16 @@ def test_all(session, sanitized_db_url): 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(self, 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'] From 90e1d339e0748ec1835b33196db961d898693ad5 Mon Sep 17 00:00:00 2001 From: Anthony King Date: Wed, 28 Apr 2021 17:09:04 +0100 Subject: [PATCH 4/5] fix: remove a spurious 'self' in the fixtures --- tests/ext/sqlalchemy_core/test_postgres.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/ext/sqlalchemy_core/test_postgres.py b/tests/ext/sqlalchemy_core/test_postgres.py index 1b8f2b4f..0a4a5609 100644 --- a/tests/ext/sqlalchemy_core/test_postgres.py +++ b/tests/ext/sqlalchemy_core/test_postgres.py @@ -59,7 +59,7 @@ def test_all(session, sanitized_db_url): assert sql_meta['sanitized_query'].endswith('FROM users') -def test_insert_on_conflict_renders(self, connection): +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() From 7347cd6bb9898fe730bc26ef2b6d0734897d0679 Mon Sep 17 00:00:00 2001 From: Anthony King Date: Wed, 28 Apr 2021 23:01:29 +0100 Subject: [PATCH 5/5] chore: make the DB URL injectable this reduces duplication of the Engine fixture. --- tests/ext/sqlalchemy_core/test_base.py | 9 +++++-- tests/ext/sqlalchemy_core/test_postgres.py | 27 +++++-------------- .../sqlalchemy_core/test_sqlalchemy_core.py | 2 +- .../sqlalchemy_core/test_sqlalchemy_core_2.py | 2 +- 4 files changed, 15 insertions(+), 25 deletions(-) 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 index 0a4a5609..5dde9a8e 100644 --- a/tests/ext/sqlalchemy_core/test_postgres.py +++ b/tests/ext/sqlalchemy_core/test_postgres.py @@ -1,6 +1,6 @@ import pytest -from .test_base import connection, User, session, Base +from .test_base import connection, engine, session, User from sqlalchemy import create_engine from sqlalchemy.dialects.postgresql import insert as pg_insert @@ -17,6 +17,11 @@ def postgres_db(): 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() @@ -28,26 +33,6 @@ def sanitized_db_url(postgres_db): ) -@pytest.fixture() -def engine(postgres_db): - """ - 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(postgres_db.url()) - 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() - - 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""" 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