AcidMap/app/database.py

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy.ext.declarative import declarative_base
import os
from dotenv import load_dotenv
import logging
from sqlalchemy.exc import SQLAlchemyError

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 创建Base类
Base = declarative_base()

# 加载环境变量
load_dotenv("config.env")

# 从环境变量获取数据库连接信息
DB_USER = os.getenv("DB_USER", "postgres")
DB_PASSWORD = os.getenv("DB_PASSWORD", "root")
DB_HOST = os.getenv("DB_HOST", "localhost")
DB_PORT = os.getenv("DB_PORT", "5432")
DB_NAME = os.getenv("DB_NAME", "testdb")
print(DB_USER, DB_PASSWORD, DB_HOST, DB_PORT, DB_NAME)

# 构建数据库连接URL
SQLALCHEMY_DATABASE_URL = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}"

def create_database_engine():
    """创建并配置数据库引擎"""
    try:
        engine = create_engine(
            SQLALCHEMY_DATABASE_URL,
            pool_size=5,
            max_overflow=10,
            pool_timeout=30,
            pool_recycle=1800
        )
        return engine
    except Exception as e:
        logger.error(f"创建数据库引擎失败: {str(e)}")
        raise

def test_database_connection(engine):
    """测试数据库连接"""
    try:
        with engine.connect() as conn:
            logger.info("数据库连接测试成功")
            return True
    except SQLAlchemyError as e:
        logger.error(f"数据库连接测试失败: {str(e)}")
        return False

# 创建数据库引擎
engine = create_database_engine()

# 测试数据库连接
if not test_database_connection(engine):
    raise Exception("无法连接到数据库，请检查数据库配置")

# 创建会话工厂
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

def get_db():
    """获取数据库会话"""
    db = SessionLocal()
    try:
        yield db
    except SQLAlchemyError as e:
        logger.error(f"数据库操作错误: {str(e)}")
        raise
    finally:
        db.close()

def execute_sql(sql_statement):
    """
    执行原始SQL语句
    
    Args:
        sql_statement: 要执行的SQL语句
    
    Returns:
        执行结果
    """
    try:
        with engine.begin() as connection:
            result = connection.execute(sql_statement)
            return result
    except SQLAlchemyError as e:
        logger.error(f"执行SQL语句失败: {str(e)}")
        raise