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

Author: antdking

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