-
Notifications
You must be signed in to change notification settings - Fork 144
Add Tracing for SQLAlchemy and Flask-SQLAlcemy #14
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from 10 commits
6d6c99d
8c613d0
1c75aa6
0729f16
31070bd
8f8468f
b9a00b8
d576942
ea672bb
c332d3a
941097a
730bef0
21195c1
4c26818
c3646f1
a9f1b0c
d494d15
e3f171c
84c5ae0
bf8d99a
087dd1f
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,59 @@ | ||
from builtins import super | ||
from flask_sqlalchemy.model import Model | ||
from sqlalchemy.orm.session import sessionmaker | ||
from flask_sqlalchemy import SQLAlchemy, BaseQuery, _SessionSignalEvents, get_state | ||
from aws_xray_sdk.ext.sqlalchemy.query import XRaySession, XRayQuery | ||
from aws_xray_sdk.ext.sqlalchemy.util.decerators import xray_on_call, decorate_all_functions | ||
|
||
|
||
@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) |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,25 @@ | ||
from builtins import super | ||
from sqlalchemy.orm.query import Query | ||
from sqlalchemy.orm.session import Session, sessionmaker | ||
from .util.decerators import xray_on_call, decorate_all_functions | ||
|
||
|
||
@decorate_all_functions(xray_on_call) | ||
class XRaySession(Session): | ||
pass | ||
|
||
|
||
@decorate_all_functions(xray_on_call) | ||
class XRayQuery(Query): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Any reason you decorate functions from Query class? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Query Class called by flask-sqlalchemyt, so this gets us complete traces |
||
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) |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
import aws_xray_sdk | ||
from aws_xray_sdk.core import xray_recorder | ||
|
||
|
||
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 |
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -35,7 +35,7 @@ | |
'Programming Language :: Python :: 3.6', | ||
], | ||
|
||
install_requires=['jsonpickle', 'wrapt', 'requests'], | ||
install_requires=['jsonpickle', 'wrapt', 'requests', 'future'], | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Do you need future if you remove all print statement? There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No future will still be needed for the super() builtin |
||
|
||
keywords='aws xray sdk', | ||
|
||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,82 @@ | ||
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.flask_sqlalchemy.query import XRayFlaskSqlAlchemy | ||
from flask import Flask | ||
|
||
|
||
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): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This along with find_sub could also be a helper class shared between your two query_test.py. And since these two functions are very general they could be under tests/ root so any other unit test can utilize your contribution. |
||
"""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 is not None: | ||
return result | ||
return None | ||
|
||
|
||
def find_sub(segment, name): | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I would recommend rename it to find_subsegment. Abbreviate on functions is not a good practice in general if it is not absolutely needed. |
||
"""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 is not 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' |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,84 @@ | ||
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 create_engine, Column, Integer, String | ||
# 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 is not 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 is not 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' |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
I think for a general SQLAlchemy guide there is no need to involve the web framework.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
You have to initialize the app from Flask in order to Initialize Flask-SQLAlchemy
db = XRayFlaskSqlAlchemy(app)
was just trying to show complete example..
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Sounds reasonable. I'm fine with this.