Trace raw SQL statements.

This is done using the ExecutionContext as storage place
(since it is an unique object shared for a single instance).
It is an ideal place as raw statement simply come as a string,
so no way to store the current span there.

We still try to ignore PRAGMA statements, though.

Author: carlosalberto

sqlalchemy_opentracing/__init__.py

@@ -101,12 +101,6 @@ def _can_operation_be_traced(conn, stmt_obj):
 
     return False
 
-def _get_span(obj):
-    '''
-    Get the span of a statement object, if any.
-    '''
-    return getattr(obj, '_span', None)
-
 def _set_traced_with_session(conn, session):
     '''
     Mark a connection to be traced with a session tracing information.
@@ -117,6 +111,11 @@ def _set_traced_with_session(conn, session):
         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):
@@ -125,15 +124,19 @@ def _normalize_stmt(statement):
 def _engine_before_cursor_handler(conn, cursor,
                                        statement, parameters,
                                        context, executemany):
-    if context.compiled is None: # PRAGMA
-        return
+    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.
-    stmt_obj = context.compiled.statement
     if not (g_trace_all 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)
@@ -148,26 +151,23 @@ def _engine_before_cursor_handler(conn, cursor,
     span.set_tag('db.statement', _normalize_stmt(statement))
     span.set_tag('sqlalchemy.dialect', context.dialect.name)
 
-    stmt_obj._span = span
+    context._span = span
 
 def _engine_after_cursor_handler(conn, cursor,
                                       statement, parameters,
                                       context, executemany):
-    if context.compiled is None: # PRAGMA
-        return
-
-    stmt_obj = context.compiled.statement
-    span = _get_span(stmt_obj)
+    span = getattr(context, '_span', None)
     if span is None:
         return
 
     span.finish()
-    clear_traced(stmt_obj)
+
+    if context.compiled is not None:
+        clear_traced(context.compiled.statement)
 
 def _engine_error_handler(exception_context):
     execution_context = exception_context.execution_context
-    stmt_obj = execution_context.compiled.statement
-    span = _get_span(stmt_obj)
+    span = getattr(execution_context, '_span', None)
     if span is None:
         return
 

tests/test_core.py

@@ -78,6 +78,44 @@ def test_traced_error(self):
             'error': 'true',
         })
 
+    def test_trace_text(self):
+        tracer = DummyTracer()
+        creat = CreateTable(self.users_table)
+        self.engine.execute(creat)
+
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        self.engine.execute('SELECT name FROM users')
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual(tracer.spans[0].operation_name, 'textclause')
+        self.assertEqual(tracer.spans[0].is_finished, True)
+        self.assertEqual(tracer.spans[0].tags, {
+            'component': 'sqlalchemy',
+            'db.statement': 'SELECT name FROM users',
+            'db.type': 'sql',
+            'sqlalchemy.dialect': 'sqlite',
+        })
+
+    def test_trace_text_error(self):
+        tracer = DummyTracer()
+
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        try:
+            self.engine.execute('SELECT name FROM users')
+        except OperationalError:
+            pass
+
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual(tracer.spans[0].operation_name, 'textclause')
+        self.assertEqual(tracer.spans[0].is_finished, True)
+        self.assertEqual(tracer.spans[0].tags, {
+            'component': 'sqlalchemy',
+            'db.statement': 'SELECT name FROM users',
+            'db.type': 'sql',
+            'sqlalchemy.dialect': 'sqlite',
+            'sqlalchemy.exception': 'no such table: users',
+            'error': 'true',
+        })
+
     def test_traced_transaction(self):
         tracer = DummyTracer()
         creat = CreateTable(self.users_table)

tests/test_orm.py

@@ -1,7 +1,7 @@
 import unittest
 from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.exc import IntegrityError
+from sqlalchemy.exc import IntegrityError, OperationalError
 from sqlalchemy.orm import sessionmaker
 
 import sqlalchemy_opentracing
@@ -113,6 +113,51 @@ def test_traced_parent(self):
         self.assertEqual(True, all(map(lambda x: x.is_finished, tracer.spans)))
         self.assertEqual(True, all(map(lambda x: x.child_of == parent_span, tracer.spans)))
 
+    def test_traced_text(self):
+        tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+
+        session = self.session
+        span = DummySpan('parent span')
+        sqlalchemy_opentracing.set_parent_span(session, span)
+        session.execute('SELECT name FROM users')
+
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual(tracer.spans[0].operation_name, 'textclause')
+        self.assertEqual(tracer.spans[0].is_finished, True)
+        self.assertEqual(tracer.spans[0].child_of, span)
+        self.assertEqual(tracer.spans[0].tags, {
+            'component': 'sqlalchemy',
+            'db.statement': 'SELECT name FROM users',
+            'db.type': 'sql',
+            'sqlalchemy.dialect': 'sqlite',
+        })
+
+    def test_traced_text_error(self):
+        tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+
+        session = self.session
+        span = DummySpan('parent span')
+        sqlalchemy_opentracing.set_parent_span(session, span)
+        try:
+            session.execute('SELECT zipcode FROM addresses')
+        except OperationalError:
+            pass
+
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual(tracer.spans[0].operation_name, 'textclause')
+        self.assertEqual(tracer.spans[0].is_finished, True)
+        self.assertEqual(tracer.spans[0].child_of, span)
+        self.assertEqual(tracer.spans[0].tags, {
+            'component': 'sqlalchemy',
+            'db.statement': 'SELECT zipcode FROM addresses',
+            'db.type': 'sql',
+            'sqlalchemy.dialect': 'sqlite',
+            'sqlalchemy.exception': 'no such table: addresses',
+            'error': 'true',
+        })
+
     def test_traced_after_commit(self):
         tracer = DummyTracer()
         sqlalchemy_opentracing.init_tracing(tracer)