Merge pull request #2 from carlosalberto/orm_integration

Implement initial support for ORM tracing.

Author: carlosalberto

examples/child-spans.py examples/core-childspans.pyexamples/child-spans.py renamed to examples/core-childspans.py

@@ -0,0 +1,54 @@
+import os
+from sqlalchemy import MetaData, Table, Integer, String, Column, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+import lightstep
+import sqlalchemy_opentracing
+
+DB_LOCATION = '/tmp/simple.db'
+
+tracer = lightstep.Tracer(
+    component_name='sqlalchemy-orm-childspan',
+    access_token='{your_lightstep_token}'
+)
+
+Base = declarative_base()
+
+class User(Base):
+    __tablename__ = 'users'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+
+if __name__ == '__main__':
+    if os.path.exists(DB_LOCATION):
+        os.remove(DB_LOCATION) # cleanup
+
+    engine = create_engine('sqlite:///%s' % DB_LOCATION)
+    session = sessionmaker(bind=engine)()
+
+    sqlalchemy_opentracing.init_tracing(tracer)
+    sqlalchemy_opentracing.register_connectable(engine)
+
+    User.metadata.create_all(engine)
+
+    # Create a parent span
+    span = tracer.start_span('parent span')
+
+    # Register the session for the current transaction.
+    sqlalchemy_opentracing.set_parent_span(session, span)
+
+    session.add(User(name='John Doe'))
+    session.add(User(name='Jason Bourne'))
+
+    users = session.query(User).all()
+    for user in users:
+        print user.name
+
+    session.add(User(name='Who'))
+
+    # Commit the session and close the parent span.
+    session.commit()
+    span.finish()
+

examples/simple.py examples/core-simple.pyexamples/simple.py renamed to examples/core-simple.py

@@ -0,0 +1,45 @@
+import os
+from sqlalchemy import MetaData, Table, Integer, String, Column, create_engine
+from sqlalchemy.ext.declarative import declarative_base
+from sqlalchemy.orm import sessionmaker
+
+import lightstep
+import sqlalchemy_opentracing
+
+DB_LOCATION = '/tmp/simple.db'
+
+tracer = lightstep.Tracer(
+    component_name='sqlalchemy-orm-simple',
+    access_token='{your_lightstep_token}'
+)
+
+Base = declarative_base()
+
+class User(Base):
+    __tablename__ = 'users'
+
+    id = Column(Integer, primary_key=True)
+    name = Column(String)
+
+if __name__ == '__main__':
+    if os.path.exists(DB_LOCATION):
+        os.remove(DB_LOCATION) # cleanup
+
+    engine = create_engine('sqlite:///%s' % DB_LOCATION)
+    session = sessionmaker(bind=engine)()
+
+    sqlalchemy_opentracing.init_tracing(tracer)
+    sqlalchemy_opentracing.register_connectable(engine)
+
+    User.metadata.create_all(engine)
+
+    # Register the session for the current transaction.
+    sqlalchemy_opentracing.set_traced(session)
+
+    # Insert a new row.
+    session.add(User(name='John Doe'))
+
+    # Finish the commit.
+    # This will stop tracing the session.
+    session.commit()
+

examples/orm-childspan.py

@@ -1,4 +1,5 @@
 from sqlalchemy.event import listen, remove
+from sqlalchemy.orm import Session
 
 g_tracer = None
 g_trace_all = True
@@ -17,92 +18,138 @@ def init_tracing(tracer, trace_all=False):
 
     g_trace_all = trace_all
 
-def get_traced(stmt_obj):
+def get_traced(obj):
     '''
     Gets a bool indicating whether or not this
-    statement is marked for tracing.
+    object is marked for tracing.
     '''
-    return getattr(stmt_obj, '_traced', False)
+    return getattr(obj, '_traced', False)
 
-def set_traced(stmt_obj):
+def set_traced(obj):
     '''
-    Mark a statement to be traced.
+    Mark a statement/session to be traced.
     '''
-    stmt_obj._traced = True
+    obj._traced = True
 
-def get_parent_span(stmt_obj):
+    # Session needs to have its connection/statements
+    # decorated as soon as a connection is acquired.
+    if isinstance(obj, Session):
+        _register_session_connection_event(obj)
+
+def get_parent_span(obj):
     '''
-    Gets a parent span for this statement, if any.
+    Gets a parent span for this object, if any.
     '''
-    return getattr(stmt_obj, '_parent_span', None)
+    return getattr(obj, '_parent_span', None)
 
-def set_parent_span(stmt_obj, parent_span):
+def set_parent_span(obj, parent_span):
     '''
-    Marks a statement as a child of a span.
+    Marks an object as a child of a span.
     It gets marked to be traced if it wasn't before.
     '''
-    stmt_obj._parent_span = parent_span
-    stmt_obj._traced = True
+    obj._parent_span = parent_span
+    set_traced(obj)
 
-def has_parent_span(stmt_obj):
+def has_parent_span(obj):
     '''
     Get whether or not the statement has
     a parent span.
     '''
-    return hasattr(stmt_obj, '_parent_span')
+    return hasattr(obj, '_parent_span')
 
-def get_span(stmt_obj):
+def get_span(obj):
     '''
     Get the span of a statement object, if any.
     '''
-    return getattr(stmt_obj, '_span', None)
+    return getattr(obj, '_span', None)
 
 def register_connectable(obj):
     '''
     Register an object to have its events be traced.
     Any Connectable object is accepted, which
     includes Connection and Engine.
     '''
-    listen(obj, 'before_cursor_execute', _before_cursor_handler)
-    listen(obj, 'after_cursor_execute', _after_cursor_handler)
-    listen(obj, 'handle_error', _error_handler)
+    listen(obj, 'before_cursor_execute', _connectable_before_cursor_handler)
+    listen(obj, 'after_cursor_execute', _connectable_after_cursor_handler)
+    listen(obj, 'handle_error', _connectable_error_handler)
 
 def unregister_connectable(obj):
     '''
     Remove a connectable from having its events being
     traced.
     '''
-    remove(obj, 'before_cursor_execute', _before_cursor_handler)
-    remove(obj, 'after_cursor_execute', _after_cursor_handler)
-    remove(obj, 'handle_error', _error_handler)
+    remove(obj, 'before_cursor_execute', _connectable_before_cursor_handler)
+    remove(obj, 'after_cursor_execute', _connectable_after_cursor_handler)
+    remove(obj, 'handle_error', _connectable_error_handler)
+
+def _can_operation_be_traced(conn, stmt_obj):
+    '''
+    Get whether an operation can be traced, depending on its
+    connection or the statement being executed, having the latter
+    the priority.
+    '''
+    if hasattr(stmt_obj, '_traced'):
+        return stmt_obj._traced
+    if hasattr(conn, '_traced'):
+        return conn._traced
+
+    return False
+
+def _clear_traced(obj):
+    '''
+    Clear an object's decorated tracing fields,
+    to prevent unintended further tracing.
+    '''
+    if hasattr(obj, '_parent_span'):
+        del obj._parent_span
+    if hasattr(obj, '_traced'):
+        del obj._traced
+
+def _set_traced_with_session(conn, session):
+    '''
+    Mark a connection to be traced with a session tracing information.
+    '''
+    conn._traced = True
+    parent_span = get_parent_span(session)
+    if parent_span is not None:
+        conn._parent_span = parent_span
 
 def _get_operation_name(stmt_obj):
     return stmt_obj.__visit_name__
 
 def _normalize_stmt(statement):
     return statement.strip().replace('\n', '').replace('\t', '')
 
-def _before_cursor_handler(conn, cursor, statement, parameters, context, executemany):
+def _connectable_before_cursor_handler(conn, cursor,
+                                       statement, parameters,
+                                       context, executemany):
     if context.compiled is None: # PRAGMA
         return
 
-    # Don't trace if trace_all is disabled and the statement wasn't marked
+    # Don't trace if trace_all is disabled
+    # and the connection/statement wasn't marked explicitly.
     stmt_obj = context.compiled.statement
-    if not (g_trace_all or get_traced(stmt_obj)):
+    if not (g_trace_all or _can_operation_be_traced(conn, stmt_obj)):
         return
 
+    # Retrieve the parent span, if any,
+    # either from the statement or inherited from the connection.
     parent_span = get_parent_span(stmt_obj)
-    operation_name = _get_operation_name(stmt_obj)
+    if parent_span is None:
+        parent_span = get_parent_span(conn)
 
     # Start a new span for this query.
-    span = g_tracer.start_span(operation_name=operation_name, child_of=parent_span)
+    name = _get_operation_name(stmt_obj)
+    span = g_tracer.start_span(operation_name=name, child_of=parent_span)
     span.set_tag('component', 'sqlalchemy')
     span.set_tag('db.type', 'sql')
     span.set_tag('db.statement', _normalize_stmt(statement))
 
     stmt_obj._span = span
 
-def _after_cursor_handler(conn, cursor, statement, parameters, context, executemany):
+def _connectable_after_cursor_handler(conn, cursor,
+                                      statement, parameters,
+                                      context, executemany):
     if context.compiled is None: # PRAGMA
         return
 
@@ -113,7 +160,7 @@ def _after_cursor_handler(conn, cursor, statement, parameters, context, executem
 
     span.finish()
 
-def _error_handler(exception_context):
+def _connectable_error_handler(exception_context):
     execution_context = exception_context.execution_context
     stmt_obj = execution_context.compiled.statement
     span = get_span(stmt_obj)
@@ -123,3 +170,28 @@ def _error_handler(exception_context):
     span.set_tag('error', 'true')
     span.finish()
 
+def _register_session_connection_event(session):
+    listen(session, 'after_begin', _session_after_begin_handler)
+
+def _session_after_begin_handler(session, transaction, conn):
+    # Have the connections inherit the tracing info
+    # from the session (including parent span, if any).
+    _set_traced_with_session(conn, session)
+
+    # Plug post-operation clean up handlers.
+    # The actual session commit/rollback is not traced by us.
+    listen(session, 'before_commit', _session_before_commit_handler, once=True)
+    listen(session, 'after_rollback', _session_rollback_handler, once=True)
+
+# Event handlers can't remove themselves
+# Fine, as long as we only fire one of them *once*.
+def _session_before_commit_handler(session):
+    _clear_traced(session)
+    remove(session, 'after_begin', _session_after_begin_handler)
+    remove(session, 'after_rollback', _session_rollback_handler)
+
+def _session_rollback_handler(session):
+    _clear_traced(session)
+    remove(session, 'after_begin', _session_after_begin_handler)
+    remove(session, 'before_commit', _session_before_commit_handler)
+