import edward

database/database.py

@@ -0,0 +1,37 @@
+from contextlib import contextmanager
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker, declarative_base
+
+from log.log_manager import logger
+
+Base = declarative_base()
+
+DATABASE_URL = 'postgresql+psycopg://postgres:K8u3fg0o@47.119.128.161:60001/squirrel'
+engine = create_engine(
+    DATABASE_URL,
+    pool_size=10,
+    max_overflow=20,
+    pool_timeout=30,
+    pool_recycle=1800,  # 防止数据库端连接过期
+    connect_args={
+        'connect_timeout': 15,
+        'keepalives_idle': 60,
+        'keepalives_interval': 10,
+        'keepalives_count': 5
+    }
+)
+Base.metadata.create_all(engine)
+
+@contextmanager
+def get_session():
+    session = sessionmaker(bind=engine)()
+    try:
+        yield session
+        session.commit()  # 自动提交成功的事务
+    except Exception as e:
+        session.rollback()  # 异常时回滚
+        logger.error(f"Database operation failed: {str(e)}")
+        raise  # 重新抛出异常
+    finally:
+        session.close()  # 确保会话关闭