1 files changed, 11 insertions, 0 deletions
diff --git a/osprofiler/sqlalchemy.py b/osprofiler/sqlalchemy.py
index e98c59c..c593684 100644
--- a/osprofiler/sqlalchemy.py
+++ b/osprofiler/sqlalchemy.py
@@ -13,6 +13,8 @@
 #    License for the specific language governing permissions and limitations
 #    under the License.
 
+import contextlib
+
 from osprofiler import profiler
 
 
@@ -42,6 +44,15 @@ def add_tracing(sqlalchemy, engine, name):
                                 _after_cursor_execute())
 
 
+@contextlib.contextmanager
+def wrap_session(sqlalchemy, sess):
+    with sess as s:
+        if not getattr(s.bind, "traced", False):
+            add_tracing(sqlalchemy, s.bind, "db")
+            s.bind.traced = True
+        yield s
+
+
 def _before_cursor_execute(name):
     """Add listener that will send trace info before query is executed."""