Make clear_traced() public.

This function clears any tracing info on
statements/connection/session objects, and it
helps for the cases where no commit/rollback
happened.

This is specially handy also for doing
bulk updates/insertions as those operation
don't trigger any commit/rollback events at all.

Author: carlosalberto

sqlalchemy_opentracing/__init__.py

@@ -41,6 +41,16 @@ def set_traced(obj):
         # 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.
@@ -94,16 +104,6 @@ def _can_operation_be_traced(conn, stmt_obj):
 
     return False
 
-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_span(obj):
     '''
     Get the span of a statement object, if any.
@@ -212,7 +212,7 @@ def _register_session_connection_event(session):
     listen(session, 'after_rollback', _session_rollback_handler)
 
 def _connection_cleanup_handler(conn):
-    _clear_traced(conn)
+    clear_traced(conn)
 
 def _session_after_begin_handler(session, transaction, conn):
     # It's only needed to pass down tracing information
@@ -221,8 +221,8 @@ def _session_after_begin_handler(session, transaction, conn):
         _set_traced_with_session(conn, session)
 
 def _session_before_commit_handler(session):
-    _clear_traced(session)
+    clear_traced(session)
 
 def _session_rollback_handler(session):
-    _clear_traced(session)
+    clear_traced(session)
 

tests/test_core.py

@@ -173,6 +173,26 @@ def test_traced_after_rollback(self):
 
         self.assertEqual(0, len(tracer.spans))
 
+    def test_traced_clear_connection(self):
+        tracer = DummyTracer()
+        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)
+            conn.execute(creat)
+
+            # Stop tracing from this point.
+            sqlalchemy_opentracing.clear_traced(conn)
+
+            conn.execute(ins)
+
+        self.assertEqual(1, len(tracer.spans))
+        self.assertEqual('create_table', tracer.spans[0].operation_name)
+
     def test_unregister_connectable(self):
         tracer = DummyTracer()
         creat = CreateTable(self.users_table)

tests/test_orm.py

@@ -216,3 +216,18 @@ def test_traced_bulk_insert(self):
             'sqlalchemy.dialect': 'sqlite',
         })
 
+    def test_traced_clear_session(self):
+        tracer = DummyTracer()
+        sqlalchemy_opentracing.init_tracing(tracer)
+
+        session = self.session
+        sqlalchemy_opentracing.set_traced(session)
+        session.add(User(name='John Doe'))
+        session.add(User(name='Jason Bourne'))
+
+        # Clear the tracing info right before committing.
+        sqlalchemy_opentracing.clear_traced(session)
+        session.commit()
+
+        self.assertEqual(0, len(tracer.spans))
+