my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

#1	import os
#2
#3	from alembic import command
#4	from alembic.config import Config
#5	from sqlalchemy import create_engine
#6	from sqlalchemy.engine.base import Engine
#7	from sqlalchemy.orm import Session as SQLAlchemySession
#8	from sqlalchemy.orm import scoped_session, sessionmaker
#9
#10	from .models import Base
#11
#12
#13	class DatabaseManager:
#14	def __init__(self, echo: bool = False):
#15	self.database_uri = os.environ.get("EMBEDCHAIN_DB_URI")
#16	self.echo = echo
#17	self.engine: Engine = None
#18	self._session_factory = None
#19
#20	def setup_engine(self) -> None:
#21	"""Initializes the database engine and session factory."""
#22	if not self.database_uri:
#23	raise RuntimeError("Database URI is not set. Set the EMBEDCHAIN_DB_URI environment variable.")
#24	connect_args = {}
#25	if self.database_uri.startswith("sqlite"):
#26	connect_args["check_same_thread"] = False
#27	self.engine = create_engine(self.database_uri, echo=self.echo, connect_args=connect_args)
#28	self._session_factory = scoped_session(sessionmaker(bind=self.engine))
#29	Base.metadata.bind = self.engine
#30
#31	def init_db(self) -> None:
#32	"""Creates all tables defined in the Base metadata."""
#33	if not self.engine:
#34	raise RuntimeError("Database engine is not initialized. Call setup_engine() first.")
#35	Base.metadata.create_all(self.engine)
#36
#37	def get_session(self) -> SQLAlchemySession:
#38	"""Provides a session for database operations."""
#39	if not self._session_factory:
#40	raise RuntimeError("Session factory is not initialized. Call setup_engine() first.")
#41	return self._session_factory()
#42
#43	def close_session(self) -> None:
#44	"""Closes the current session."""
#45	if self._session_factory:
#46	self._session_factory.remove()
#47
#48	def execute_transaction(self, transaction_block):
#49	"""Executes a block of code within a database transaction."""
#50	session = self.get_session()
#51	try:
#52	transaction_block(session)
#53	session.commit()
#54	except Exception as e:
#55	session.rollback()
#56	raise e
#57	finally:
#58	self.close_session()
#59
#60
#61	# Singleton pattern to use throughout the application
#62	database_manager = DatabaseManager()
#63
#64
#65	# Convenience functions for backward compatibility and ease of use
#66	def setup_engine(database_uri: str, echo: bool = False) -> None:
#67	database_manager.database_uri = database_uri
#68	database_manager.echo = echo
#69	database_manager.setup_engine()
#70
#71
#72	def alembic_upgrade() -> None:
#73	"""Upgrades the database to the latest version."""
#74	alembic_config_path = os.path.join(os.path.dirname(__file__), "..", "..", "alembic.ini")
#75	alembic_cfg = Config(alembic_config_path)
#76	command.upgrade(alembic_cfg, "head")
#77
#78
#79	def init_db() -> None:
#80	alembic_upgrade()
#81
#82
#83	def get_session() -> SQLAlchemySession:
#84	return database_manager.get_session()
#85
#86
#87	def execute_transaction(transaction_block):
#88	database_manager.execute_transaction(transaction_block)
#89

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public