Merge pull request #4 from carlosalberto/trace_all_engines

Add a flag to initialize the tracing for *all* engines/databases.

Author: carlosalberto

README.rst

@@ -32,14 +32,27 @@ By default, only statements marked to be traced are taken into account (explicit
 
 .. code-block:: python
 
-    sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+    sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
     sqlalchemy_opentracing.register_engine(engine)
 
     # this statement will be traced too (without a parent span, though)
     with engine.begin() as conn:
         sel = select([users])
         conn.execute(sel)
 
+It is also possible to have all engines being registered automatically (independently of the `trace_all_queries` flag, which can be enabled or disabled):
+
+.. code-block:: python
+
+    # No need to call register_engine()
+    sqlalchemy_opentracing.init_tracing(tracer, trace_all_engines=True)
+
+    with engine.begin() as conn:
+        sel = select([users])
+        ...
+
+This is equivalent to calling `register_engine` with the `sqlalchemy.engine.Engine` class.
+
 The resulting spans will have an operation name related to the sql statement (such as `create-table` or `insert`), and will include exception information (if any), the dialect/backend (such as sqlite), and a few other hints.
 
 Tracing under a Connection

sqlalchemy_opentracing/__init__.py

@@ -1,23 +1,26 @@
-from sqlalchemy.engine import Connection
+from sqlalchemy.engine import Connection, Engine
 from sqlalchemy.event import contains, listen, remove
 from sqlalchemy.orm import Session
 
 g_tracer = None
-g_trace_all = True
+g_trace_all_queries = False
 
-def init_tracing(tracer, trace_all=False):
+def init_tracing(tracer, trace_all_engines=False, trace_all_queries=False):
     '''
     Set our global tracer.
     Tracer objects from our pyramid/flask/django libraries
     can be passed as well.
     '''
-    global g_tracer, g_trace_all
+    global g_tracer, g_trace_all_queries
     if hasattr(tracer, '_tracer'):
         g_tracer = tracer._tracer
     else:
         g_tracer = tracer
 
-    g_trace_all = trace_all
+    g_trace_all_queries = trace_all_queries
+
+    if trace_all_engines:
+        register_engine(Engine)
 
 def get_traced(obj):
     '''
@@ -130,7 +133,7 @@ def _engine_before_cursor_handler(conn, cursor,
 
     # Don't trace if trace_all is disabled
     # and the connection/statement wasn't marked explicitly.
-    if not (g_trace_all or _can_operation_be_traced(conn, stmt_obj)):
+    if not (g_trace_all_queries or _can_operation_be_traced(conn, stmt_obj)):
         return
 
     # Don't trace PRAGMA statements coming from SQLite

tests/test_api.py

@@ -1,4 +1,5 @@
 import unittest
+from mock import patch
 from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.schema import CreateTable
@@ -18,20 +19,32 @@ def test_init(self):
         tracer = DummyTracer()
         sqlalchemy_opentracing.init_tracing(tracer)
         self.assertEqual(tracer, sqlalchemy_opentracing.g_tracer)
-        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all)
+        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all_queries)
 
     def test_init_subtracer(self):
         tracer = DummyTracer(with_subtracer=True)
         sqlalchemy_opentracing.init_tracing(tracer)
         self.assertEqual(tracer._tracer, sqlalchemy_opentracing.g_tracer)
-        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all)
+        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all_queries)
 
-    def test_init_traceall(self):
-        sqlalchemy_opentracing.init_tracing(DummyTracer(), False)
-        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all)
+    def test_init_trace_all_queries(self):
+        sqlalchemy_opentracing.init_tracing(DummyTracer(), trace_all_queries=False)
+        self.assertEqual(False, sqlalchemy_opentracing.g_trace_all_queries)
 
-        sqlalchemy_opentracing.init_tracing(DummyTracer(), True)
-        self.assertEqual(True, sqlalchemy_opentracing.g_trace_all)
+        sqlalchemy_opentracing.init_tracing(DummyTracer(), trace_all_queries=True)
+        self.assertEqual(True, sqlalchemy_opentracing.g_trace_all_queries)
+
+    @patch('sqlalchemy_opentracing.register_engine')
+    def test_init_trace_all_engines(self, mock_register):
+        tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        self.assertEqual(0, mock_register.call_count)
+
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_engines=False)
+        self.assertEqual(0, mock_register.call_count)
+
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_engines=True)
+        self.assertEqual(1, mock_register.call_count)
 
     def test_traced_property(self):
         stmt_obj = CreateTable(self.users_table)

tests/test_core.py

@@ -1,5 +1,6 @@
 import unittest
 from sqlalchemy import create_engine, MetaData, Table, Column, Integer, String
+from sqlalchemy.engine import Engine
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.schema import CreateTable
 from sqlalchemy.sql import select
@@ -44,15 +45,33 @@ def test_traced_none(self):
 
         self.assertEqual(0, len(tracer.spans))
 
-    def test_traced_all(self):
+    def test_traced_all_queries(self):
         tracer = DummyTracer()
         creat = CreateTable(self.users_table)
 
-        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
         self.engine.execute(creat)
 
         self.assertEqual(1, len(tracer.spans))
 
+    def test_traced_all_engines(self):
+        # Create an engine that is not registered.
+        e = create_engine('sqlite:///:memory:')
+        tracer = DummyTracer()
+        creat = CreateTable(self.users_table)
+
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_engines=True)
+        sqlalchemy_opentracing.set_traced(creat)
+        e.execute(creat)
+
+        # Unregister the main Engine class before doing our assertions,
+        # in case we fail.
+        sqlalchemy_opentracing.unregister_engine(Engine)
+
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual('create_table', tracer.spans[0].operation_name)
+        self.assertEqual(True, tracer.spans[0].is_finished)
+
     def test_traced_error(self):
         tracer = DummyTracer()
         creat = CreateTable(self.users_table)
@@ -82,7 +101,7 @@ def test_traced_error(self):
     def test_trace_text(self):
         tracer = DummyTracer()
 
-        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
         self.engine.execute('CREATE TABLE users (id INTEGER NOT NULL, name VARCHAR, PRIMARY KEY (id))')
         self.assertEqual(1, len(tracer.spans))
         self.assertEqual(tracer.spans[0].operation_name, 'textclause')
@@ -97,7 +116,7 @@ def test_trace_text(self):
     def test_trace_text_error(self):
         tracer = DummyTracer()
 
-        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
         try:
             self.engine.execute('SELECT name FROM users')
         except OperationalError:
@@ -260,7 +279,7 @@ def test_unregister_engine(self):
         tracer = DummyTracer()
         creat = CreateTable(self.users_table)
 
-        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
         self.engine.execute(creat)
         self.assertEqual(1, len(tracer.spans))
 

tests/test_orm.py

@@ -56,7 +56,7 @@ def test_traced_none(self):
     # test mixing insert with select and insert
     def test_traced_all(self):
         tracer = DummyTracer()
-        sqlalchemy_opentracing.init_tracing(tracer, trace_all=True)
+        sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
 
         session = self.session
         session.add(User(name='John Doe'))