Report an error when registering an engine and all are traced already.

Author: carlosalberto

sqlalchemy_opentracing/__init__.py

@@ -4,14 +4,15 @@
 
 g_tracer = None
 g_trace_all_queries = False
+g_trace_all_engines = 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_queries
+    global g_tracer, g_trace_all_engines, g_trace_all_queries
 
     if hasattr(tracer, '_tracer'):
         tracer = tracer._tracer
@@ -20,6 +21,7 @@ def init_tracing(tracer, trace_all_engines=False, trace_all_queries=False):
 
     g_tracer = tracer
     g_trace_all_queries = trace_all_queries
+    g_trace_all_engines = trace_all_engines
 
     if trace_all_engines:
         register_engine(Engine)
@@ -83,6 +85,8 @@ def register_engine(obj):
     '''
     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)

tests/test_api.py

@@ -1,6 +1,7 @@
 import unittest
 from mock import patch
 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
 
@@ -72,3 +73,12 @@ def test_register_no_tracer(self):
         with self.assertRaises(RuntimeError):
             sqlalchemy_opentracing.register_engine(engine)
 
+    def test_register_dup(self):
+        engine = create_engine('sqlite:///:memory:')
+        sqlalchemy_opentracing.init_tracing(DummyTracer(), trace_all_engines=True)
+        with self.assertRaises(RuntimeError):
+            sqlalchemy_opentracing.register_engine(engine)
+
+        # Manually clear the Engine from listening events.
+        sqlalchemy_opentracing.unregister_engine(Engine)
+