Sqlalchemy query

aws_xray_sdk/ext/sqlalchemy/util/decerators.py

@@ -5,7 +5,6 @@
 from urllib.parse import urlparse, uses_netloc
 
 
-
 def decorate_all_functions(function_decorator):
     def decorator(cls):
         for c in cls.__bases__:
@@ -21,6 +20,7 @@ def decorator(cls):
         return cls
     return decorator
 
+
 def xray_on_call(cls, func):
     def wrapper(*args, **kw):
         from ..query import XRayQuery, XRaySession
@@ -40,9 +40,8 @@ def wrapper(*args, **kw):
                 if isinstance(arg, XRayQuery):
                     try:
                         sql = parse_bind(arg.session.bind)
-                        # Commented our for later PR
-                        # sql['sanitized_query'] = str(arg)
-                    except:
+                        sql['sanitized_query'] = str(arg)
+                    except Exception:
                         sql = None
         if sql is not None:
             if getattr(c._local, 'entities', None) is not None:
@@ -52,7 +51,7 @@ def wrapper(*args, **kw):
         res = func(*args, **kw)
         if subsegment is not None:
             subsegment.set_sql(sql)
-            subsegment.put_annotation("sqlalchemy", class_name+'.'+func.__name__ );
+            subsegment.put_annotation("sqlalchemy", class_name+'.'+func.__name__)
             xray_recorder.end_subsegment()
         return res
     return wrapper

tests/ext/flask_sqlalchemy/test_query.py

@@ -41,7 +41,7 @@ def test_all(capsys, session):
     User.query.all()
     subsegment = find_subsegment_by_annotation(xray_recorder.current_segment(), 'sqlalchemy', 'sqlalchemy.orm.query.all')
     assert subsegment['annotations']['sqlalchemy'] == 'sqlalchemy.orm.query.all'
-    # assert subsegment['sql']['sanitized_query']
+    assert subsegment['sql']['sanitized_query']
     assert subsegment['sql']['url']
 
 

tests/ext/sqlalchemy/test_query.py

@@ -42,7 +42,7 @@ def test_all(capsys, session):
     session.query(User).all()
     subsegment = find_subsegment_by_annotation(xray_recorder.current_segment(), 'sqlalchemy', 'sqlalchemy.orm.query.all')
     assert subsegment['annotations']['sqlalchemy'] == 'sqlalchemy.orm.query.all'
-    # assert subsegment['sql']['sanitized_query']
+    assert subsegment['sql']['sanitized_query']
     assert subsegment['sql']['url']
 
 
@@ -57,13 +57,14 @@ def test_add(capsys, session):
     assert subsegment['sql']['url']
 
 
-def test_filter(capsys, session):
-    """ Test calling all() on get all records.
+def test_filter_first(capsys, session):
+    """ Test calling filter().first() on get first filtered records.
     Verify we run the query and return the SQL as metdata"""
     # with capsys.disabled():
-    session.query(User).filter(User.password=="mypassword!")
-    subsegment = find_subsegment_by_annotation(xray_recorder.current_segment(), 'sqlalchemy', 'sqlalchemy.orm.query.filter')
-    assert subsegment['annotations']['sqlalchemy'] == 'sqlalchemy.orm.query.filter'
-    # assert subsegment['sql']['sanitized_query']
-    # assert "mypassword!" not in subsegment['sql']['sanitized_query']
+    session.query(User).filter(User.password=="mypassword!").first()
+    subsegment = find_subsegment_by_annotation(xray_recorder.current_segment(), 'sqlalchemy', 'sqlalchemy.orm.query.first')
+    assert subsegment['annotations']['sqlalchemy'] == 'sqlalchemy.orm.query.first'
+    assert subsegment['sql']['sanitized_query']
+    assert "mypassword!" not in subsegment['sql']['sanitized_query']
+    assert "users.password = ?" in  subsegment['sql']['sanitized_query']
     assert subsegment['sql']['url']