🔧 chore(base.py): import contextlib's contextmanager to improve code readability and maintainability
✨ feat(base.py): add session_getter context manager to handle session creation, rollback, and closure in a more robust way 🔧 chore(base.py): update get_session function to use session_getter context manager for session creation
This commit is contained in:
parent
9811b97908
commit
a722fe063b
1 changed files with 16 additions and 2 deletions
|
|
@ -1,3 +1,4 @@
|
||||||
|
from contextlib import contextmanager
|
||||||
from langflow.settings import settings
|
from langflow.settings import settings
|
||||||
from sqlmodel import SQLModel, Session, create_engine
|
from sqlmodel import SQLModel, Session, create_engine
|
||||||
from langflow.utils.logger import logger
|
from langflow.utils.logger import logger
|
||||||
|
|
@ -32,6 +33,19 @@ def create_db_and_tables():
|
||||||
logger.debug("Database and tables created successfully")
|
logger.debug("Database and tables created successfully")
|
||||||
|
|
||||||
|
|
||||||
def get_session():
|
@contextmanager
|
||||||
with Session(engine) as session:
|
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
|
yield session
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue