from sqlalchemy.engine import Connection, Engine
from sqlalchemy.event import contains, listen, remove
from sqlalchemy.orm import Session

g_tracer = None
g_trace_all_queries = False
g_trace_all_engines = False

def init_tracing(tracer, trace_all_engines=True, trace_all_queries=True):
    '''
    Set our global tracer.
    Tracer objects from our pyramid/flask/django libraries
    can be passed as well.
    '''
    global g_tracer, g_trace_all_engines, g_trace_all_queries

    if hasattr(tracer, '_tracer'):
        tracer = tracer._tracer

    g_tracer = tracer
    g_trace_all_queries = trace_all_queries
    g_trace_all_engines = trace_all_engines

    if trace_all_engines:
        register_engine(Engine)

def get_traced(obj):
    '''
    Gets a bool indicating whether or not this
    object is marked for tracing.
    '''
    return getattr(obj, '_traced', False)

def set_traced(obj):
    '''
    Mark a statement/session to be traced.
    '''
    obj._traced = True

    if isinstance(obj, Session):
        # Session needs to have its connection/statements
        # decorated as soon as a connection is acquired.
        _register_session_events(obj)
    elif isinstance(obj, Connection):
        # Connection simply needs to be cleaned up
        # after commit/rollback.
        _register_connection_events(obj)

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 get_parent_span(obj):
    '''
    Gets a parent span for this object, if any.
    '''
    return getattr(obj, '_parent_span', None)

def set_parent_span(obj, parent_span):
    '''
    Marks an object as a child of a span.
    It gets marked to be traced if it wasn't before.
    '''
    obj._parent_span = parent_span
    set_traced(obj)

def has_parent_span(obj):
    '''
    Get whether or not the statement has
    a parent span.
    '''
    return hasattr(obj, '_parent_span')

def register_engine(obj):
    '''
    Register an engine to have its events be traced.
    '''
    if g_tracer is None:
        raise RuntimeError('The tracer is not properly set')
    if g_trace_all_engines and obj != Engine:
        raise RuntimeError('Tracing all engines already')

    listen(obj, 'before_cursor_execute', _engine_before_cursor_handler)
    listen(obj, 'after_cursor_execute', _engine_after_cursor_handler)
    listen(obj, 'handle_error', _engine_error_handler)

def unregister_engine(obj):
    '''
    Remove an engine from having its events being traced.
    '''
    remove(obj, 'before_cursor_execute', _engine_before_cursor_handler)
    remove(obj, 'after_cursor_execute', _engine_after_cursor_handler)
    remove(obj, 'handle_error', _engine_error_handler)

def _clear_tracer():
    '''
    Set the tracer to None. For test cases usage.
    '''
    global g_tracer
    g_tracer = None

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 _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):
    if stmt_obj is None:
        # Match what the ORM shows when raw SQL
        # statements are invoked.
        return 'textclause'

    return stmt_obj.__visit_name__

def _normalize_stmt(statement):
    return statement.strip().replace('\n', '').replace('\t', '')

def _engine_before_cursor_handler(conn, cursor,
                                       statement, parameters,
                                       context, executemany):
    stmt_obj = None
    if context.compiled is not None:
        stmt_obj = context.compiled.statement

    # Don't trace if trace_all is disabled
    # and the connection/statement wasn't marked explicitly.
    if not (g_trace_all_queries or _can_operation_be_traced(conn, stmt_obj)):
        return

    # Don't trace PRAGMA statements coming from SQLite
    if stmt_obj is None and statement.startswith('PRAGMA'):
        return

    # Retrieve the parent span, if any,
    # either from the statement or inherited from the connection.
    parent_span = get_parent_span(stmt_obj)
    if parent_span is None:
        parent_span = get_parent_span(conn)

    # Start a new span for this query.
    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))
    span.set_tag('sqlalchemy.dialect', context.dialect.name)

    context._span = span

def _engine_after_cursor_handler(conn, cursor,
                                      statement, parameters,
                                      context, executemany):
    span = getattr(context, '_span', None)
    if span is None:
        return

    span.finish()

    if context.compiled is not None:
        clear_traced(context.compiled.statement)

def _engine_error_handler(exception_context):
    execution_context = exception_context.execution_context
    span = getattr(execution_context, '_span', None)
    if span is None:
        return

    exc = exception_context.original_exception
    span.set_tag('sqlalchemy.exception', str(exc))
    span.set_tag('error', 'true')
    span.finish()

    if execution_context.compiled is not None:
        clear_traced(execution_context.compiled.statement)

def _register_connection_events(conn):
    '''
    Register clean up events for our
    connection only once, as adding/removing them
    seems an expensive operation.
    '''

    # Use 'commit' as a mark to guess the events being handled.
    if contains(conn, 'commit', _connection_cleanup_handler):
        return

    # Plug post-operation clean up handlers.
    listen(conn, 'commit', _connection_cleanup_handler)
    listen(conn, 'rollback', _connection_cleanup_handler)

def _register_session_events(session):
    '''
    Register connection/transaction and clean up events
    for our session only once, as adding/removing them
    seems an expensive operation.
    '''

    # Use 'after_being' as a mark to guess the events being handled.
    if contains(session, 'after_begin', _session_after_begin_handler):
        return

    # Have the connections inherit the tracing info
    # from the session (including parent span, if any).
    listen(session, 'after_begin', _session_after_begin_handler)

    # Plug post-operation clean up handlers.
    # The actual session commit/rollback is not traced by us.
    listen(session, 'after_commit', _session_cleanup_handler)
    listen(session, 'after_rollback', _session_cleanup_handler)

def _connection_cleanup_handler(conn):
    clear_traced(conn)

def _session_after_begin_handler(session, transaction, conn):
    # It's only needed to pass down tracing information
    # if it was explicitly set.
    if get_traced(session):
        _set_traced_with_session(conn, session)

def _session_cleanup_handler(session):
    clear_traced(session)


from sqlalchemy.engine import Connection, Engine
from sqlalchemy.event import contains, listen, remove
from sqlalchemy.orm import Session

g_tracer = None
g_trace_all_queries = False
g_trace_all_engines = False

def init_tracing(tracer, trace_all_engines=True, trace_all_queries=True):
    '''
    Set our global tracer.
    Tracer objects from our pyramid/flask/django libraries
    can be passed as well.
    '''
    global g_tracer, g_trace_all_engines, g_trace_all_queries

    if hasattr(tracer, '_tracer'):
        tracer = tracer._tracer

    g_tracer = tracer
    g_trace_all_queries = trace_all_queries
    g_trace_all_engines = trace_all_engines

    if trace_all_engines:
        register_engine(Engine)

def get_traced(obj):
    '''
    Gets a bool indicating whether or not this
    object is marked for tracing.
    '''
    return getattr(obj, '_traced', False)

def set_traced(obj):
    '''
    Mark a statement/session to be traced.
    '''
    obj._traced = True

    if isinstance(obj, Session):
        # Session needs to have its connection/statements
        # decorated as soon as a connection is acquired.
        _register_session_events(obj)
    elif isinstance(obj, Connection):
        # Connection simply needs to be cleaned up
        # after commit/rollback.
        _register_connection_events(obj)

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 get_parent_span(obj):
    '''
    Gets a parent span for this object, if any.
    '''
    return getattr(obj, '_parent_span', None)

def set_parent_span(obj, parent_span):
    '''
    Marks an object as a child of a span.
    It gets marked to be traced if it wasn't before.
    '''
    obj._parent_span = parent_span
    set_traced(obj)

def has_parent_span(obj):
    '''
    Get whether or not the statement has
    a parent span.
    '''
    return hasattr(obj, '_parent_span')

def register_engine(obj):
    '''
    Register an engine to have its events be traced.
    '''
    if g_tracer is None:
        raise RuntimeError('The tracer is not properly set')
    if g_trace_all_engines and obj != Engine:
        raise RuntimeError('Tracing all engines already')

    listen(obj, 'before_cursor_execute', _engine_before_cursor_handler)
    listen(obj, 'after_cursor_execute', _engine_after_cursor_handler)
    listen(obj, 'handle_error', _engine_error_handler)

def unregister_engine(obj):
    '''
    Remove an engine from having its events being traced.
    '''
    remove(obj, 'before_cursor_execute', _engine_before_cursor_handler)
    remove(obj, 'after_cursor_execute', _engine_after_cursor_handler)
    remove(obj, 'handle_error', _engine_error_handler)

def _clear_tracer():
    '''
    Set the tracer to None. For test cases usage.
    '''
    global g_tracer
    g_tracer = None

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 _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):
    if stmt_obj is None:
        # Match what the ORM shows when raw SQL
        # statements are invoked.
        return 'textclause'

    return stmt_obj.__visit_name__

def _normalize_stmt(statement):
    return statement.strip().replace('\n', '').replace('\t', '')

def _engine_before_cursor_handler(conn, cursor,
                                       statement, parameters,
                                       context, executemany):
    stmt_obj = None
    if context.compiled is not None:
        stmt_obj = context.compiled.statement

    # Don't trace if trace_all is disabled
    # and the connection/statement wasn't marked explicitly.
    if not (g_trace_all_queries or _can_operation_be_traced(conn, stmt_obj)):
        return

    # Don't trace PRAGMA statements coming from SQLite
    if stmt_obj is None and statement.startswith('PRAGMA'):
        return

    # Retrieve the parent span, if any,
    # either from the statement or inherited from the connection.
    parent_span = get_parent_span(stmt_obj)
    if parent_span is None:
        parent_span = get_parent_span(conn)

    # Start a new span for this query.
    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))
    span.set_tag('sqlalchemy.dialect', context.dialect.name)

    context._span = span

def _engine_after_cursor_handler(conn, cursor,
                                      statement, parameters,
                                      context, executemany):
    span = getattr(context, '_span', None)
    if span is None:
        return

    span.finish()

    if context.compiled is not None:
        clear_traced(context.compiled.statement)

def _engine_error_handler(exception_context):
    execution_context = exception_context.execution_context
    span = getattr(execution_context, '_span', None)
    if span is None:
        return

    exc = exception_context.original_exception
    span.set_tag('sqlalchemy.exception', str(exc))
    span.set_tag('error', 'true')
    span.finish()

    if execution_context.compiled is not None:
        clear_traced(execution_context.compiled.statement)

def _register_connection_events(conn):
    '''
    Register clean up events for our
    connection only once, as adding/removing them
    seems an expensive operation.
    '''

    # Use 'commit' as a mark to guess the events being handled.
    if contains(conn, 'commit', _connection_cleanup_handler):
        return

    # Plug post-operation clean up handlers.
    listen(conn, 'commit', _connection_cleanup_handler)
    listen(conn, 'rollback', _connection_cleanup_handler)

def _register_session_events(session):
    '''
    Register connection/transaction and clean up events
    for our session only once, as adding/removing them
    seems an expensive operation.
    '''

    # Use 'after_being' as a mark to guess the events being handled.
    if contains(session, 'after_begin', _session_after_begin_handler):
        return

    # Have the connections inherit the tracing info
    # from the session (including parent span, if any).
    listen(session, 'after_begin', _session_after_begin_handler)

    # Plug post-operation clean up handlers.
    # The actual session commit/rollback is not traced by us.
    listen(session, 'after_commit', _session_cleanup_handler)
    listen(session, 'after_rollback', _session_cleanup_handler)

def _connection_cleanup_handler(conn):
    clear_traced(conn)

def _session_after_begin_handler(session, transaction, conn):
    # It's only needed to pass down tracing information
    # if it was explicitly set.
    if get_traced(session):
        _set_traced_with_session(conn, session)

def _session_cleanup_handler(session):
    clear_traced(session)


Set our global tracer.     Tracer objects from our pyramid/flask/django libraries     can be passed as well. Gets a bool indicating whether or not this     object is marked for tracing. Mark a statement/session to be traced. # Session needs to have its connection/statements # decorated as soon as a connection is acquired. # Connection simply needs to be cleaned up # after commit/rollback. Clear an object's decorated tracing fields,     to prevent unintended further tracing. Gets a parent span for this object, if any. Marks an object as a child of a span.     It gets marked to be traced if it wasn't before. Get whether or not the statement has     a parent span. Register an engine to have its events be traced. Remove an engine from having its events being traced. Set the tracer to None. For test cases usage. Get whether an operation can be traced, depending on its     connection or the statement being executed, having the latter     the priority. Mark a connection to be traced with a session tracing information. # Match what the ORM shows when raw SQL # statements are invoked. # Don't trace if trace_all is disabled # and the connection/statement wasn't marked explicitly. # Don't trace PRAGMA statements coming from SQLite # Retrieve the parent span, if any, # either from the statement or inherited from the connection. # Start a new span for this query. Register clean up events for our     connection only once, as adding/removing them     seems an expensive operation. # Use 'commit' as a mark to guess the events being handled. # Plug post-operation clean up handlers. Register connection/transaction and clean up events     for our session only once, as adding/removing them     seems an expensive operation. # Use 'after_being' as a mark to guess the events being handled. # Have the connections inherit the tracing info # from the session (including parent span, if any). # Plug post-operation clean up handlers. # The actual session commit/rollback is not traced by us. # It's only needed to pass down tracing information # if it was explicitly set.