100 lines
3.0 KiB
Python
100 lines
3.0 KiB
Python
|
import os
|
|||
|
|
import pandas as pd
|
|||
|
|
from sqlalchemy import create_engine, text
|
|||
|
|
from utils import load_config
|
|||
|
|
|
|||
|
|
|
|||
|
|
class DatabaseManager:
|
|||
|
|
"""
|
|||
|
|
数据库管理类，负责数据的存储、查询和管理
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
def __init__(self):
|
|||
|
|
self.config = load_config()
|
|||
|
|
self._engine = None
|
|||
|
|
|
|||
|
|
def get_engine(self):
|
|||
|
|
"""获取SQLite数据库引擎，如果不存在则创建"""
|
|||
|
|
if self._engine is not None:
|
|||
|
|
return self._engine
|
|||
|
|
|
|||
|
|
db_path = self.config['sqlite']['path']
|
|||
|
|
# 确保数据库目录存在
|
|||
|
|
os.makedirs(os.path.dirname(db_path), exist_ok=True)
|
|||
|
|
# 创建SQLite数据库引擎
|
|||
|
|
self._engine = create_engine(f'sqlite:///{db_path}', echo=False)
|
|||
|
|
return self._engine
|
|||
|
|
|
|||
|
|
def get_table_name(self, key):
|
|||
|
|
"""根据表键名获取实际表名"""
|
|||
|
|
return self.config['sqlite']['table_name'].get(key, key)
|
|||
|
|
|
|||
|
|
def get_existing_trade_dates(self, table_key):
|
|||
|
|
"""
|
|||
|
|
从数据库中获取已有的交易日期
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
table_key (str): 数据表键名
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
set: 已存在于数据库中的交易日期集合
|
|||
|
|
"""
|
|||
|
|
table_name = self.get_table_name(table_key)
|
|||
|
|
engine = self.get_engine()
|
|||
|
|
query = f"SELECT DISTINCT trade_date FROM {table_name}"
|
|||
|
|
try:
|
|||
|
|
with engine.connect() as connection:
|
|||
|
|
result = connection.execute(text(query))
|
|||
|
|
return {row[0] for row in result}
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"获取已存在交易日期时出错: {e}")
|
|||
|
|
return set()
|
|||
|
|
|
|||
|
|
def load_df_from_db(self, table_key, conditions=None):
|
|||
|
|
"""
|
|||
|
|
从数据库中加载数据
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
table_key (str): 表键名
|
|||
|
|
conditions (str): 过滤条件，如 "trade_date > '20230101'"
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
pandas.DataFrame: 查询结果
|
|||
|
|
"""
|
|||
|
|
table_name = self.get_table_name(table_key)
|
|||
|
|
engine = self.get_engine()
|
|||
|
|
query = f"SELECT * FROM {table_name}"
|
|||
|
|
if conditions:
|
|||
|
|
query += f" WHERE {conditions}"
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
return pd.read_sql(query, engine)
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"从数据库加载数据时出错: {e}")
|
|||
|
|
return pd.DataFrame()
|
|||
|
|
|
|||
|
|
def save_df_to_db(self, df, table_key, if_exists='append'):
|
|||
|
|
"""
|
|||
|
|
保存DataFrame到数据库
|
|||
|
|
|
|||
|
|
参数:
|
|||
|
|
df (pandas.DataFrame): 要保存的数据
|
|||
|
|
table_key (str): 表键名
|
|||
|
|
if_exists (str): 如果表存在时的操作: 'fail', 'replace', 或 'append'
|
|||
|
|
|
|||
|
|
返回:
|
|||
|
|
bool: 操作是否成功
|
|||
|
|
"""
|
|||
|
|
if df.empty:
|
|||
|
|
print("警告: 尝试保存空的DataFrame到数据库")
|
|||
|
|
return False
|
|||
|
|
|
|||
|
|
table_name = self.get_table_name(table_key)
|
|||
|
|
engine = self.get_engine()
|
|||
|
|
|
|||
|
|
try:
|
|||
|
|
df.to_sql(table_name, engine, if_exists=if_exists, index=False)
|
|||
|
|
return True
|
|||
|
|
except Exception as e:
|
|||
|
|
print(f"保存数据到数据库时出错: {e}")
|
|||
|
|
return False
|