📦 chore(manager.py): add DatabaseManager class to handle database operations
📦 chore(utils.py): add initialize_database function and session_getter context manager to handle database initialization and session management
This commit is contained in:
parent
8036f74d8b
commit
8436c66aa7
2 changed files with 32 additions and 25 deletions
|
|
@ -1,4 +1,3 @@
|
||||||
from contextlib import contextmanager
|
|
||||||
from pathlib import Path
|
from pathlib import Path
|
||||||
from langflow.services.base import Service
|
from langflow.services.base import Service
|
||||||
from sqlmodel import SQLModel, Session, create_engine
|
from sqlmodel import SQLModel, Session, create_engine
|
||||||
|
|
@ -12,7 +11,7 @@ class DatabaseManager(Service):
|
||||||
|
|
||||||
def __init__(self, database_url: str):
|
def __init__(self, database_url: str):
|
||||||
self.database_url = database_url
|
self.database_url = database_url
|
||||||
# This file is in langflow.services.database.base.py
|
# This file is in langflow.services.database.manager.py
|
||||||
# the ini is in langflow
|
# the ini is in langflow
|
||||||
langflow_dir = Path(__file__).parent.parent.parent
|
langflow_dir = Path(__file__).parent.parent.parent
|
||||||
self.script_location = langflow_dir / "alembic"
|
self.script_location = langflow_dir / "alembic"
|
||||||
|
|
@ -65,26 +64,3 @@ class DatabaseManager(Service):
|
||||||
raise RuntimeError("Something went wrong creating the database and tables.")
|
raise RuntimeError("Something went wrong creating the database and tables.")
|
||||||
else:
|
else:
|
||||||
logger.debug("Database and tables created successfully")
|
logger.debug("Database and tables created successfully")
|
||||||
|
|
||||||
|
|
||||||
@contextmanager
|
|
||||||
def session_getter(db_manager: DatabaseManager):
|
|
||||||
try:
|
|
||||||
session = Session(db_manager.engine)
|
|
||||||
yield session
|
|
||||||
except Exception as e:
|
|
||||||
print("Session rollback because of exception:", e)
|
|
||||||
session.rollback()
|
|
||||||
raise
|
|
||||||
finally:
|
|
||||||
session.close()
|
|
||||||
|
|
||||||
|
|
||||||
def initialize_database():
|
|
||||||
logger.debug("Initializing database")
|
|
||||||
from langflow.services import service_manager, ServiceType
|
|
||||||
|
|
||||||
database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
|
|
||||||
database_manager.run_migrations()
|
|
||||||
database_manager.create_db_and_tables()
|
|
||||||
logger.debug("Database initialized")
|
|
||||||
31
src/backend/langflow/services/database/utils.py
Normal file
31
src/backend/langflow/services/database/utils.py
Normal file
|
|
@ -0,0 +1,31 @@
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
from langflow.utils.logger import logger
|
||||||
|
from contextlib import contextmanager
|
||||||
|
|
||||||
|
from sqlmodel import Session
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from langflow.services.database.manager import DatabaseManager
|
||||||
|
|
||||||
|
|
||||||
|
def initialize_database():
|
||||||
|
logger.debug("Initializing database")
|
||||||
|
from langflow.services import service_manager, ServiceType
|
||||||
|
|
||||||
|
database_manager = service_manager.get(ServiceType.DATABASE_MANAGER)
|
||||||
|
database_manager.run_migrations()
|
||||||
|
database_manager.create_db_and_tables()
|
||||||
|
logger.debug("Database initialized")
|
||||||
|
|
||||||
|
|
||||||
|
@contextmanager
|
||||||
|
def session_getter(db_manager: "DatabaseManager"):
|
||||||
|
try:
|
||||||
|
session = Session(db_manager.engine)
|
||||||
|
yield session
|
||||||
|
except Exception as e:
|
||||||
|
print("Session rollback because of exception:", e)
|
||||||
|
session.rollback()
|
||||||
|
raise
|
||||||
|
finally:
|
||||||
|
session.close()
|
||||||
Loading…
Add table
Add a link
Reference in a new issue