158 lines
5.6 KiB
Python
158 lines
5.6 KiB
Python
import pymysql
|
||
import pandas as pd
|
||
from typing import Optional
|
||
from contextlib import contextmanager
|
||
import os
|
||
from dotenv import load_dotenv
|
||
from sqlalchemy import create_engine, text
|
||
|
||
load_dotenv()
|
||
|
||
|
||
class DatabaseManager:
|
||
"""管理 OHLCV 数据的 MySQL 数据库"""
|
||
|
||
def __init__(self, db_path: str = None):
|
||
"""初始化数据库管理器 (db_path 参数已废弃，仅为兼容性保留)"""
|
||
self._engine = None
|
||
self.initialize_db()
|
||
|
||
def _get_sqlalchemy_engine(self):
|
||
"""获取 SQLAlchemy 引擎（用于 pandas read_sql）"""
|
||
if self._engine is None:
|
||
host = os.getenv('MYSQL_HOST')
|
||
port = os.getenv('MYSQL_PORT', 3306)
|
||
user = os.getenv('MYSQL_USER')
|
||
password = os.getenv('MYSQL_PASSWORD')
|
||
database = os.getenv('MYSQL_DB')
|
||
self._engine = create_engine(
|
||
f"mysql+pymysql://{user}:{password}@{host}:{port}/{database}?charset=utf8mb4"
|
||
)
|
||
return self._engine
|
||
|
||
@contextmanager
|
||
def get_connection(self):
|
||
"""获取 MySQL 数据库连接"""
|
||
conn = pymysql.connect(
|
||
host=os.getenv('MYSQL_HOST'),
|
||
port=int(os.getenv('MYSQL_PORT', 3306)),
|
||
user=os.getenv('MYSQL_USER'),
|
||
password=os.getenv('MYSQL_PASSWORD'),
|
||
database=os.getenv('MYSQL_DB'),
|
||
charset='utf8mb4',
|
||
)
|
||
try:
|
||
yield conn
|
||
finally:
|
||
conn.close()
|
||
|
||
def initialize_db(self):
|
||
"""创建表（如果不存在）"""
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
CREATE TABLE IF NOT EXISTS ohlcv (
|
||
timestamp DATETIME,
|
||
symbol VARCHAR(20),
|
||
timeframe VARCHAR(10),
|
||
open DOUBLE,
|
||
high DOUBLE,
|
||
low DOUBLE,
|
||
close DOUBLE,
|
||
volume DOUBLE,
|
||
PRIMARY KEY (symbol, timeframe, timestamp),
|
||
INDEX idx_ohlcv_time (symbol, timeframe, timestamp)
|
||
) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
|
||
""")
|
||
conn.commit()
|
||
|
||
def save_ohlcv(self, df: pd.DataFrame, symbol: str, timeframe: str):
|
||
"""保存 OHLCV 数据到数据库 (Upsert)"""
|
||
if df.empty:
|
||
return
|
||
|
||
# 确保 timestamp 是 datetime 类型
|
||
df = df.copy()
|
||
if not pd.api.types.is_datetime64_any_dtype(df['timestamp']):
|
||
df['timestamp'] = pd.to_datetime(df['timestamp'])
|
||
|
||
# 转换为字符串格式
|
||
df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')
|
||
|
||
# 添加元数据列
|
||
if 'symbol' not in df.columns:
|
||
df['symbol'] = symbol
|
||
if 'timeframe' not in df.columns:
|
||
df['timeframe'] = timeframe
|
||
|
||
cols = ['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']
|
||
data = [tuple(row) for row in df[cols].values]
|
||
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
sql = """
|
||
REPLACE INTO ohlcv
|
||
(timestamp, symbol, timeframe, open, high, low, close, volume)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
|
||
"""
|
||
cursor.executemany(sql, data)
|
||
conn.commit()
|
||
|
||
def get_ohlcv(self, symbol: str, timeframe: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
|
||
"""从数据库获取 OHLCV 数据"""
|
||
|
||
query = """
|
||
SELECT timestamp, open, high, low, close, volume
|
||
FROM ohlcv
|
||
WHERE symbol = :p0 AND timeframe = :p1
|
||
"""
|
||
params = {'p0': symbol, 'p1': timeframe}
|
||
param_idx = 2
|
||
|
||
if start_date:
|
||
query += f" AND timestamp >= :p{param_idx}"
|
||
params[f'p{param_idx}'] = start_date
|
||
param_idx += 1
|
||
|
||
if end_date:
|
||
# 处理 'YYYY-MM-DD' 格式
|
||
if len(str(end_date)) == 10:
|
||
query += f" AND timestamp < DATE_ADD(:p{param_idx}, INTERVAL 1 DAY)"
|
||
else:
|
||
query += f" AND timestamp <= :p{param_idx}"
|
||
params[f'p{param_idx}'] = end_date
|
||
|
||
query += " ORDER BY timestamp ASC"
|
||
|
||
engine = self._get_sqlalchemy_engine()
|
||
with engine.connect() as conn:
|
||
df = pd.read_sql_query(text(query), conn, params=params, parse_dates=['timestamp'])
|
||
|
||
return df
|
||
|
||
def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
|
||
"""获取指定 symbol/timeframe 的最新时间戳"""
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
SELECT MAX(timestamp)
|
||
FROM ohlcv
|
||
WHERE symbol = %s AND timeframe = %s
|
||
""", (symbol, timeframe))
|
||
result = cursor.fetchone()
|
||
|
||
return pd.to_datetime(result[0]) if result and result[0] else None
|
||
|
||
def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
|
||
"""获取指定 symbol/timeframe 的最早时间戳"""
|
||
with self.get_connection() as conn:
|
||
cursor = conn.cursor()
|
||
cursor.execute("""
|
||
SELECT MIN(timestamp)
|
||
FROM ohlcv
|
||
WHERE symbol = %s AND timeframe = %s
|
||
""", (symbol, timeframe))
|
||
result = cursor.fetchone()
|
||
|
||
return pd.to_datetime(result[0]) if result and result[0] else None
|