Merge pull request #5 from carlosalberto/extra_clean_checks

Extra clean checks

Author: carlosalberto

NOTES.txt

@@ -70,7 +70,14 @@ exists for both Connection and Session objects. Sadly, such dictionary
 is 'moved' around for every Connection of the parent Engine, thus is
 not unique, thus we can't use it for thread scenarios.
 
-One alternative was also to have thread local objects - specifically
+A second alternative is to use a weak dictionary - so we keep
+the statements/connection/session around, and even if they are
+used/ignored/forgotten, we don't keep them around forever. This is
+similar to what the Django/Flask/Pyramid connectors do, with the
+exception of them having a well defined lifetime for the decorated
+objects (the request one).
+
+A third alternative was also to have thread local objects - specifically
 to hold the current active span. While this may have worked just fine
 for us, this is something that has been done and re-done in other
 OT libraries/language. Thus we will wait to see how the eventual

sqlalchemy_opentracing/__init__.py

@@ -4,20 +4,24 @@
 
 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'):
-        g_tracer = tracer._tracer
+        tracer = tracer._tracer
     else:
-        g_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)
@@ -79,6 +83,11 @@ 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)
@@ -91,6 +100,13 @@ def unregister_engine(obj):
     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

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
 
@@ -15,6 +16,9 @@ def setUp(self):
             Column('name', String),
         )
 
+    def tearDown(self):
+        sqlalchemy_opentracing._clear_tracer()
+
     def test_init(self):
         tracer = DummyTracer()
         sqlalchemy_opentracing.init_tracing(tracer)
@@ -64,3 +68,17 @@ def test_has_parent_none(self):
         self.assertEqual(False, sqlalchemy_opentracing.has_parent_span(stmt))
         self.assertEqual(None, sqlalchemy_opentracing.get_parent_span(stmt))
 
+    def test_register_no_tracer(self):
+        engine = create_engine('sqlite:///:memory:')
+        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)
+

tests/test_core.py

@@ -15,13 +15,15 @@ def setUp(self):
             Column('id', Integer, primary_key=True),
             Column('name', String),
         )
-        sqlalchemy_opentracing.register_engine(self.engine)
+    def tearDown(self):
+        sqlalchemy_opentracing._clear_tracer()
 
     def test_traced(self):
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
+        creat = CreateTable(self.users_table)
         sqlalchemy_opentracing.set_traced(creat)
         self.engine.execute(creat)
 
@@ -38,31 +40,32 @@ def test_traced(self):
 
     def test_traced_none(self):
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
+        creat = CreateTable(self.users_table)
         self.engine.execute(creat)
 
         self.assertEqual(0, len(tracer.spans))
 
     def test_traced_all_queries(self):
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
+        creat = CreateTable(self.users_table)
         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:')
+        # Don't register the engine explicitly.
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer, trace_all_engines=True)
+
+        creat = CreateTable(self.users_table)
         sqlalchemy_opentracing.set_traced(creat)
-        e.execute(creat)
+        self.engine.execute(creat)
 
         # Unregister the main Engine class before doing our assertions,
         # in case we fail.
@@ -74,9 +77,10 @@ def test_traced_all_engines(self):
 
     def test_traced_error(self):
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
+        creat = CreateTable(self.users_table)
         self.engine.execute(creat)
         self.assertEqual(0, len(tracer.spans))
 
@@ -100,8 +104,9 @@ def test_traced_error(self):
 
     def test_trace_text(self):
         tracer = DummyTracer()
-
         sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         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')
@@ -115,8 +120,9 @@ def test_trace_text(self):
 
     def test_trace_text_error(self):
         tracer = DummyTracer()
-
         sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         try:
             self.engine.execute('SELECT name FROM users')
         except OperationalError:
@@ -136,11 +142,13 @@ def test_trace_text_error(self):
 
     def test_traced_transaction(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
         ins = self.users_table.insert().values(name='John Doe')
         sel = select([self.users_table])
 
-        sqlalchemy_opentracing.init_tracing(tracer)
         parent_span = DummySpan('parent span')
         conn = self.engine.connect()
         with conn.begin() as trans:
@@ -157,11 +165,13 @@ def test_traced_transaction(self):
 
     def test_traced_transaction_nested(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
         ins = self.users_table.insert().values(name='John Doe')
         sel = select([self.users_table])
 
-        sqlalchemy_opentracing.init_tracing(tracer)
         parent_span = DummySpan('parent span')
         conn = self.engine.connect()
 
@@ -181,13 +191,15 @@ def test_traced_transaction_nested(self):
 
     def test_traced_rollback(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
         ins = self.users_table.insert().values(name='John Doe')
 
         # Don't trace this.
         self.engine.execute(creat)
 
-        sqlalchemy_opentracing.init_tracing(tracer)
         parent_span = DummySpan('parent span')
         conn = self.engine.connect()
         try:
@@ -208,9 +220,11 @@ def test_traced_rollback(self):
 
     def test_traced_after_transaction(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
 
-        sqlalchemy_opentracing.init_tracing(tracer)
         conn = self.engine.connect()
         with conn.begin() as tx:
             sqlalchemy_opentracing.set_traced(conn)
@@ -229,10 +243,12 @@ def test_traced_after_transaction(self):
 
     def test_traced_after_rollback(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
 
         # Create a table, but don't trace it
-        sqlalchemy_opentracing.init_tracing(tracer)
         conn = self.engine.connect()
         with conn.begin() as tx:
             conn.execute(creat)
@@ -257,11 +273,12 @@ def test_traced_after_rollback(self):
 
     def test_traced_clear_connection(self):
         tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
         creat = CreateTable(self.users_table)
         ins = self.users_table.insert().values(name='John Doe')
 
-        sqlalchemy_opentracing.init_tracing(tracer)
-
         conn = self.engine.connect()
         with conn.begin() as tx:
             sqlalchemy_opentracing.set_traced(conn)
@@ -277,9 +294,10 @@ def test_traced_clear_connection(self):
 
     def test_unregister_engine(self):
         tracer = DummyTracer()
-        creat = CreateTable(self.users_table)
-
         sqlalchemy_opentracing.init_tracing(tracer, trace_all_queries=True)
+        sqlalchemy_opentracing.register_engine(self.engine)
+
+        creat = CreateTable(self.users_table)
         self.engine.execute(creat)
         self.assertEqual(1, len(tracer.spans))