feat(sqla-core): Add support for rendering Database Specific queries

.travis.yml

@@ -9,6 +9,9 @@ python:
   - "3.8"
   - "3.9"
 
+addons:
+  postgresql: "9.6"
+
 install:
   - pip install tox tox-travis
 

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:

tests/ext/sqlalchemy_core/test_base.py

@@ -21,15 +21,20 @@ 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
     it cleans up context storage again.
     """
     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)

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']

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
 

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