chore(config): 添加环境变量示例及.gitignore更新
- 新增 .env.example,包含 Tushare API、钉钉机器人和PostgreSQL数据库配置模板 - 更新.gitignore,忽略本地配置文件如 .env.local 和 config_local.py - 添加对报表文件命名规则的支持,保留示例文件不忽略 - 删除废弃的 chart.py 及相关图表模块代码 - 新增 config/settings.py,实现从环境变量读取配置的统一接口 - 设置数据目录及缓存目录,确保目录存在,提高配置管理规范性
This commit is contained in:
96
core/common/db.py
Normal file
96
core/common/db.py
Normal file
@@ -0,0 +1,96 @@
|
||||
"""
|
||||
数据库配置和连接工具
|
||||
"""
|
||||
|
||||
import psycopg2
|
||||
from psycopg2.extras import RealDictCursor
|
||||
from sqlalchemy import create_engine
|
||||
import pandas as pd
|
||||
from loguru import logger
|
||||
from typing import Optional
|
||||
|
||||
from config.settings import get_db_config
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""数据库管理类"""
|
||||
|
||||
def __init__(self, config: dict = None):
|
||||
self.config = config or get_db_config()
|
||||
self.engine = None
|
||||
|
||||
def get_engine(self):
|
||||
"""获取SQLAlchemy引擎"""
|
||||
if self.engine is None:
|
||||
conn_str = (
|
||||
f"postgresql://{self.config['username']}:{self.config['password']}"
|
||||
f"@{self.config['host']}:{self.config['port']}/{self.config['database']}"
|
||||
)
|
||||
self.engine = create_engine(
|
||||
conn_str,
|
||||
pool_pre_ping=True,
|
||||
pool_recycle=300,
|
||||
echo=False,
|
||||
)
|
||||
return self.engine
|
||||
|
||||
def get_connection(self):
|
||||
"""获取psycopg2连接"""
|
||||
return psycopg2.connect(
|
||||
host=self.config["host"],
|
||||
port=self.config["port"],
|
||||
database=self.config["database"],
|
||||
user=self.config["username"],
|
||||
password=self.config["password"],
|
||||
)
|
||||
|
||||
def test_connection(self) -> bool:
|
||||
"""测试数据库连接"""
|
||||
try:
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor() as cursor:
|
||||
cursor.execute("SELECT 1")
|
||||
result = cursor.fetchone()
|
||||
logger.info("数据库连接测试成功")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"数据库连接测试失败: {e}")
|
||||
return False
|
||||
|
||||
def execute_query(self, query: str, params: tuple = None) -> Optional[list]:
|
||||
"""执行查询并返回结果"""
|
||||
try:
|
||||
with self.get_connection() as conn:
|
||||
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
||||
cursor.execute(query, params)
|
||||
result = cursor.fetchall()
|
||||
return [dict(row) for row in result]
|
||||
except Exception as e:
|
||||
logger.error(f"执行查询失败: {e}")
|
||||
return None
|
||||
|
||||
def insert_dataframe(
|
||||
self, df: pd.DataFrame, table_name: str, if_exists: str = "append"
|
||||
) -> bool:
|
||||
"""将DataFrame插入到数据库表"""
|
||||
try:
|
||||
engine = self.get_engine()
|
||||
df.to_sql(
|
||||
table_name,
|
||||
engine,
|
||||
if_exists=if_exists,
|
||||
index=False,
|
||||
method="multi",
|
||||
chunksize=1000,
|
||||
)
|
||||
logger.info(f"成功插入 {len(df)} 条记录到表 {table_name}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"插入数据到表 {table_name} 失败: {e}")
|
||||
return False
|
||||
|
||||
def close(self):
|
||||
"""关闭连接"""
|
||||
if self.engine:
|
||||
self.engine.dispose()
|
||||
self.engine = None
|
||||
Reference in New Issue
Block a user