From a722fe063bcd0f352c912656ff5f8341f35b2d3f Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Fri, 21 Jul 2023 09:24:59 -0300
Subject: [PATCH] =?UTF-8?q?=F0=9F=94=A7=20chore(base.py):=20import=20conte?=
 =?UTF-8?q?xtlib's=20contextmanager=20to=20improve=20code=20readability=20?=
 =?UTF-8?q?and=20maintainability=20=E2=9C=A8=20feat(base.py):=20add=20sess?=
 =?UTF-8?q?ion=5Fgetter=20context=20manager=20to=20handle=20session=20crea?=
 =?UTF-8?q?tion,=20rollback,=20and=20closure=20in=20a=20more=20robust=20wa?=
 =?UTF-8?q?y=20=F0=9F=94=A7=20chore(base.py):=20update=20get=5Fsession=20f?=
 =?UTF-8?q?unction=20to=20use=20session=5Fgetter=20context=20manager=20for?=
 =?UTF-8?q?=20session=20creation?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/backend/langflow/database/base.py | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/src/backend/langflow/database/base.py b/src/backend/langflow/database/base.py
index 256434523..338298a6b 100644
--- a/src/backend/langflow/database/base.py
+++ b/src/backend/langflow/database/base.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from langflow.settings import settings
 from sqlmodel import SQLModel, Session, create_engine
 from langflow.utils.logger import logger
@@ -32,6 +33,19 @@ def create_db_and_tables():
         logger.debug("Database and tables created successfully")
 
 
-def get_session():
-    with Session(engine) as session:
+@contextmanager
+def session_getter():
+    try:
+        session = Session(engine)
+        yield session
+    except Exception as e:
+        print("Session rollback because of exception:", e)
+        session.rollback()
+        raise
+    finally:
+        session.close()
+
+
+def get_session():
+    with session_getter() as session:
         yield session