Add Tracing for SQLAlchemy and Flask-SQLAlcemy

README.md

@@ -175,7 +175,61 @@ app.router.add_get("/", handler)
 
 web.run_app(app)
 ```
+**Add Flask middleware**
+
+```python
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.ext.flask.middleware import XRayMiddleware
+
+app = Flask(__name__)
+
+xray_recorder.configure(service='fallback_name', dynamic_naming='*mysite.com*')
+XRayMiddleware(app, xray_recorder)
+```
+
+
+**Use SQLAlchemy ORM**
+The SQLAlchemy integration requires you to override the Session and Query Classes for SQL Alchemy
+
+SQLAlchemy integration uses subsegments so you need to have a segment started before you make a query.
+
+```python
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import *
+from sqlalchemy.orm import sessionmaker
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+from aws_xray_sdk.ext.sqlalchemy.query import XRaySessionMaker
+
+xray_recorder.configure(service='test', sampling=False, context=Context())
+xray_recorder.clear_trace_entities()
+xray_recorder.begin_segment('SQLAlchemyTest')
+
+Session = XRaySessionMaker(bind=engine)
+session = Session()
 
+xray_recorder.end_segment()
+app = Flask(__name__)
+
+xray_recorder.configure(service='fallback_name', dynamic_naming='*mysite.com*')
+XRayMiddleware(app, xray_recorder)
+```
+
+**Add Flask-SQLAlchemy**
+
+```python
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.ext.flask.middleware import XRayMiddleware
+from aws_xray_sdk.ext.flask_sqlalchemy.query import XRayFlaskSqlAlchemy
+
+app = Flask(__name__)
+app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
+
+xray_recorder.configure(service='fallback_name', dynamic_naming='*mysite.com*')
+XRayMiddleware(app, xray_recorder)
+db = XRayFlaskSqlAlchemy(app)
+
+```
 ## License
 
 The AWS X-Ray SDK for Python is licensed under the Apache 2.0 License. See LICENSE and NOTICE.txt for more information.

aws_xray_sdk/ext/flask_sqlalchemy/query.py

@@ -0,0 +1,95 @@
+from __future__ import print_function
+from builtins import super
+import aws_xray_sdk
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+from sqlalchemy.orm.base import _generative
+from sqlalchemy.orm.query import Query
+from flask_sqlalchemy.model import Model
+from sqlalchemy.orm.session import Session, sessionmaker
+from functools import wraps
+from flask_sqlalchemy import SQLAlchemy, BaseQuery, _SessionSignalEvents, get_state
+from aws_xray_sdk.ext.sqlalchemy.query import XRaySession, XRayQuery
+
+def decorate_all_functions(function_decorator):
+    def decorator(cls):
+        for c in cls.__bases__:
+            for name, obj in vars(c).items():
+                if name.startswith("_"):
+                    continue
+                if callable(obj):
+                    try:
+                        obj = obj.__func__  # unwrap Python 2 unbound method
+                    except AttributeError:
+                        pass  # not needed in Python 3
+                    setattr(c, name, function_decorator(c, obj))
+        return cls
+    return decorator
+
+def xray_on_call(cls, func):
+    def wrapper(*args, **kw):
+        class_name = str(cls.__module__)
+        c = xray_recorder._context
+        if getattr(c._local, 'entities', None) is  not None:
+             trace = xray_recorder.begin_subsegment(class_name+'.'+func.__name__)
+        else:
+            trace = None
+        res = func(*args, **kw)
+        if trace is not None:
+            if class_name == 'sqlalchemy.orm.query':
+                for arg in args:
+                    if isinstance(arg, aws_xray_sdk.ext.sqlalchemy.query.XRayQuery):
+                        trace.put_metadata("sql", str(arg));
+            xray_recorder.end_subsegment()
+        return res
+    return wrapper
+
+@decorate_all_functions(xray_on_call)   
+class XRayBaseQuery(BaseQuery):
+    BaseQuery.__bases__ = (XRayQuery,)
+
+class XRaySignallingSession(XRaySession):
+    """The signalling session is the default session that Flask-SQLAlchemy
+    uses.  It extends the default session system with bind selection and
+    modification tracking.
+    If you want to use a different session you can override the
+    :meth:`SQLAlchemy.create_session` function.
+    .. versionadded:: 2.0
+    .. versionadded:: 2.1
+        The `binds` option was added, which allows a session to be joined
+        to an external transaction.
+    """
+
+    def __init__(self, db, autocommit=False, autoflush=True, **options):
+        #: The application that this session belongs to.
+        self.app = app = db.get_app()
+        track_modifications = app.config['SQLALCHEMY_TRACK_MODIFICATIONS']
+        bind = options.pop('bind', None) or db.engine
+        binds = options.pop('binds', db.get_binds(app))
+
+        if track_modifications is None or track_modifications:
+            _SessionSignalEvents.register(self)
+
+        XRaySession.__init__(
+            self, autocommit=autocommit, autoflush=autoflush,
+            bind=bind, binds=binds, **options
+        )
+
+    def get_bind(self, mapper=None, clause=None):
+        # mapper is None if someone tries to just get a connection
+        if mapper is not None:
+            info = getattr(mapper.mapped_table, 'info', {})
+            bind_key = info.get('bind_key')
+            if bind_key is not None:
+                state = get_state(self.app)
+                return state.db.get_engine(self.app, bind=bind_key)
+        return XRaySession.get_bind(self, mapper, clause)
+
+class XRayFlaskSqlAlchemy(SQLAlchemy):
+    def __init__(self, app=None, use_native_unicode=True, session_options=None,
+                 metadata=None, query_class=XRayBaseQuery, model_class=Model):
+        super().__init__(app, use_native_unicode, session_options,
+              metadata, query_class, model_class)
+
+    def create_session(self, options):
+         return sessionmaker(class_=XRaySignallingSession, db=self, **options)

aws_xray_sdk/ext/sqlalchemy/query.py

@@ -0,0 +1,61 @@
+from __future__ import print_function
+from builtins import super
+import aws_xray_sdk
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+from sqlalchemy.orm.base import _generative
+from sqlalchemy.orm.query import Query
+from sqlalchemy.orm.session import Session, sessionmaker
+from functools import wraps
+
+
+def decorate_all_functions(function_decorator):
+    def decorator(cls):
+        for c in cls.__bases__:
+            for name, obj in vars(c).items():
+                if name.startswith("_"):
+                    continue
+                if callable(obj):
+                    try:
+                        obj = obj.__func__  # unwrap Python 2 unbound method
+                    except AttributeError:
+                        pass  # not needed in Python 3
+                    setattr(c, name, function_decorator(c, obj))
+        return cls
+    return decorator
+
+def xray_on_call(cls, func):
+    def wrapper(*args, **kw):
+        class_name = str(cls.__module__)
+        c = xray_recorder._context
+        if getattr(c._local, 'entities', None) is  not None:
+             trace = xray_recorder.begin_subsegment(class_name+'.'+func.__name__)
+        else:
+            trace = None
+        res = func(*args, **kw)
+        if trace is not None:
+            if class_name == 'sqlalchemy.orm.query':
+                for arg in args:
+                    if isinstance(arg, aws_xray_sdk.ext.sqlalchemy.query.XRayQuery):
+                        trace.put_metadata("sql", str(arg));
+            xray_recorder.end_subsegment()
+        return res
+    return wrapper
+
+@decorate_all_functions(xray_on_call)   
+class XRaySession(Session):
+    pass
+
+@decorate_all_functions(xray_on_call)   
+class XRayQuery(Query):
+    pass
+
+@decorate_all_functions(xray_on_call)
+class XRaySessionMaker(sessionmaker):
+    def __init__(self, bind=None, class_=XRaySession, autoflush=True,
+                 autocommit=False,
+                 expire_on_commit=True,
+                 info=None, **kw):
+        kw['query_cls'] = XRayQuery
+        super().__init__(bind, class_, autoflush, autocommit, expire_on_commit,
+                         info, **kw)

setup.py

@@ -35,7 +35,7 @@
         'Programming Language :: Python :: 3.6',
     ],
 
-    install_requires=['jsonpickle', 'wrapt', 'requests'],
+    install_requires=['jsonpickle', 'wrapt', 'requests', 'future'],
 
     keywords='aws xray sdk',
 

tests/ext/flask_sqlalchemy/test_query.py

@@ -0,0 +1,79 @@
+from __future__ import absolute_import
+import json
+import jsonpickle
+import pytest
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+from aws_xray_sdk.ext.sqlalchemy.query import XRayQuery, XRaySession
+from aws_xray_sdk.ext.flask_sqlalchemy.query import XRayFlaskSqlAlchemy
+from flask import Flask
+from flask_sqlalchemy import SQLAlchemy
+
+
+app = Flask(__name__)
+app.config["SQLALCHEMY_TRACK_MODIFICATIONS"] = False
+app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
+db = XRayFlaskSqlAlchemy(app)
+class User(db.Model):
+    __tablename__ = "users"
+
+    id = db.Column(db.Integer, primary_key=True)
+    name = db.Column(db.String(255), nullable=False, unique=True)
+    fullname = db.Column(db.String(255), nullable=False)
+    password = db.Column(db.String(255), nullable=False)
+
+
+def _search_entity(entity, name):
+    """Helper function to that recursivly looks at subentities
+    Returns a serialized entity that matches the name given or None"""
+    if 'name' in entity:
+        my_name = entity['name']
+        if my_name == name:
+            return entity
+        else:
+            if "subsegments" in entity:
+                for s in entity['subsegments']:
+                    result = _search_entity(s,name)
+                    if result != None:
+                        return result
+    return None
+def find_sub(segment, name):
+    """Helper function to find a subsegment by name in the entity tree"""
+    segment = jsonpickle.encode(segment, unpicklable=False)
+    segment = json.loads(segment)
+    for entity in segment['subsegments']:
+        result = _search_entity(entity, name)
+        if result != None:
+            return result
+    return None
+
+@pytest.fixture()
+def session():
+    """Test Fixture to Create DataBase Tables and start a trace segment"""
+    xray_recorder.configure(service='test', sampling=False, context=Context())
+    xray_recorder.clear_trace_entities()
+    xray_recorder.begin_segment('SQLAlchemyTest')
+    db.create_all()
+    yield
+    xray_recorder.end_segment()
+    xray_recorder.clear_trace_entities()
+
+def test_all(capsys, session):
+    """ Test calling all() on get all records.
+    Verify that we capture trace of query and return the SQL as metdata"""
+    # with capsys.disabled():
+    User.query.all()
+    subsegment = find_sub(xray_recorder.current_segment(),'sqlalchemy.orm.query.all')
+    assert subsegment['name'] == 'sqlalchemy.orm.query.all'
+    assert subsegment['metadata']['default']['sql']
+
+def test_add(capsys, session):
+    """ Test calling add() on insert a row.
+    Verify we that we capture trace for the add"""
+    # with capsys.disabled():
+    john = User(name='John', fullname = "John Doe", password="password")
+    db.session.add(john)
+    subsegment = find_sub(xray_recorder.current_segment(),'sqlalchemy.orm.session.add')
+    assert subsegment['name'] == 'sqlalchemy.orm.session.add'
+
+

tests/ext/sqlalchemy/test_query.py

@@ -0,0 +1,82 @@
+from __future__ import absolute_import
+import pytest
+import json
+import jsonpickle
+from aws_xray_sdk.core import xray_recorder
+from aws_xray_sdk.core.context import Context
+from aws_xray_sdk.ext.sqlalchemy.query import XRaySessionMaker
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy import *
+# from sqlalchemy.orm import sessionmaker
+
+
+
+Base = declarative_base()
+
+class User(Base):
+        __tablename__ = 'users'
+
+        id = Column(Integer, primary_key=True)
+        name = Column(String)
+        fullname = Column(String)
+        password = Column(String)
+
+
+def _search_entity(entity, name):
+    """Helper function to that recursivly looks at subentities
+    Returns a serialized entity that matches the name given or None"""
+    if 'name' in entity:
+        my_name = entity['name']
+        if my_name == name:
+            return entity
+        else:
+            if "subsegments" in entity:
+                for s in entity['subsegments']:
+                    result = _search_entity(s,name)
+                    if result != None:
+                        return result
+    return None
+def find_sub(segment, name):
+    """Helper function to find a subsegment by name in the entity tree"""
+    segment = jsonpickle.encode(segment, unpicklable=False)
+    segment = json.loads(segment)
+    for entity in segment['subsegments']:
+        result = _search_entity(entity, name)
+        if result != None:
+            return result
+    return None
+
+@pytest.fixture()
+def session():
+    """Test Fixture to Create DataBase Tables and start a trace segment"""
+    engine = create_engine('sqlite:///:memory:')
+    xray_recorder.configure(service='test', sampling=False, context=Context())
+    xray_recorder.clear_trace_entities()
+    xray_recorder.begin_segment('SQLAlchemyTest')
+    Session = XRaySessionMaker(bind=engine)
+    Base.metadata.create_all(engine)
+    session = Session()
+    yield session
+    xray_recorder.end_segment()
+    xray_recorder.clear_trace_entities()
+
+def test_all(capsys, session):
+    """ Test calling all() on get all records.
+    Verify we run the query and return the SQL as metdata"""
+    # with capsys.disabled():
+    session.query(User).all()
+    sub = find_sub(xray_recorder.current_segment(), 'sqlalchemy.orm.query.all')
+    assert sub['name'] == 'sqlalchemy.orm.query.all'
+    assert sub['metadata']['default']['sql']
+
+
+def test_add(capsys, session):
+    """ Test calling add() on insert a row.
+    Verify we that we capture trace for the add"""
+    # with capsys.disabled():
+    john = User(name='John', fullname = "John Doe", password="password")
+    session.add(john)
+    sub = find_sub(xray_recorder.current_segment(), 'sqlalchemy.orm.session.add')
+    assert sub['name'] == 'sqlalchemy.orm.session.add'
+
+