140 lines
5.0 KiB
Python
140 lines
5.0 KiB
Python
import sqlite3
|
|
import pandas as pd
|
|
from pathlib import Path
|
|
from typing import Optional, List, Tuple
|
|
from contextlib import contextmanager
|
|
import os
|
|
|
|
class DatabaseManager:
|
|
"""Manages SQLite database for storing OHLCV data."""
|
|
|
|
def __init__(self, db_path: str = "./data/market_data.db"):
|
|
self.db_path = Path(db_path)
|
|
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
|
self.initialize_db()
|
|
|
|
@contextmanager
|
|
def get_connection(self):
|
|
"""Context manager for database connection."""
|
|
conn = sqlite3.connect(self.db_path)
|
|
try:
|
|
yield conn
|
|
finally:
|
|
conn.close()
|
|
|
|
def initialize_db(self):
|
|
"""Create tables if they don't exist."""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# Create OHLCV table
|
|
# Composite primary key ensures uniqueness for a specific candle
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS ohlcv (
|
|
timestamp DATETIME,
|
|
symbol TEXT,
|
|
timeframe TEXT,
|
|
open REAL,
|
|
high REAL,
|
|
low REAL,
|
|
close REAL,
|
|
volume REAL,
|
|
PRIMARY KEY (symbol, timeframe, timestamp)
|
|
)
|
|
""")
|
|
|
|
# Create index for faster querying by time range
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_ohlcv_time
|
|
ON ohlcv (symbol, timeframe, timestamp)
|
|
""")
|
|
|
|
conn.commit()
|
|
|
|
def save_ohlcv(self, df: pd.DataFrame, symbol: str, timeframe: str):
|
|
"""Save OHLCV data to database (Upsert)."""
|
|
if df.empty:
|
|
return
|
|
|
|
# Ensure timestamp is datetime
|
|
df = df.copy()
|
|
if not pd.api.types.is_datetime64_any_dtype(df['timestamp']):
|
|
df['timestamp'] = pd.to_datetime(df['timestamp'])
|
|
|
|
# Convert to string for SQLite compatibility (ISO 8601)
|
|
df['timestamp'] = df['timestamp'].dt.strftime('%Y-%m-%d %H:%M:%S')
|
|
|
|
# Add metadata columns if they don't exist
|
|
if 'symbol' not in df.columns:
|
|
df['symbol'] = symbol
|
|
if 'timeframe' not in df.columns:
|
|
df['timeframe'] = timeframe
|
|
|
|
# Prepare data for insertion (ensure order)
|
|
records = df[['timestamp', 'symbol', 'timeframe', 'open', 'high', 'low', 'close', 'volume']].to_records(index=False)
|
|
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
# efficiently insert or replace
|
|
# sqlite3 adapter handles datetime objects correctly
|
|
cursor.executemany("""
|
|
INSERT OR REPLACE INTO ohlcv
|
|
(timestamp, symbol, timeframe, open, high, low, close, volume)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
|
""", records)
|
|
conn.commit()
|
|
|
|
def get_ohlcv(self, symbol: str, timeframe: str, start_date: Optional[str] = None, end_date: Optional[str] = None) -> pd.DataFrame:
|
|
"""Retrieve OHLCV data from database."""
|
|
query = """
|
|
SELECT timestamp, open, high, low, close, volume
|
|
FROM ohlcv
|
|
WHERE symbol = ? AND timeframe = ?
|
|
"""
|
|
params = [symbol, timeframe]
|
|
|
|
if start_date:
|
|
query += " AND timestamp >= ?"
|
|
params.append(start_date)
|
|
|
|
if end_date:
|
|
# Handle 'YYYY-MM-DD' vs full timestamp
|
|
if len(str(end_date)) == 10: # YYYY-MM-DD
|
|
query += " AND timestamp < date(?, '+1 day')"
|
|
else:
|
|
query += " AND timestamp <= ?"
|
|
params.append(end_date)
|
|
|
|
query += " ORDER BY timestamp ASC"
|
|
|
|
with self.get_connection() as conn:
|
|
df = pd.read_sql_query(query, conn, params=params, parse_dates=['timestamp'])
|
|
|
|
return df
|
|
|
|
def get_latest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
|
|
"""Get the latest timestamp for a symbol/timeframe."""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
SELECT MAX(timestamp)
|
|
FROM ohlcv
|
|
WHERE symbol = ? AND timeframe = ?
|
|
""", (symbol, timeframe))
|
|
result = cursor.fetchone()
|
|
|
|
return pd.to_datetime(result[0]) if result[0] else None
|
|
|
|
def get_earliest_timestamp(self, symbol: str, timeframe: str) -> Optional[pd.Timestamp]:
|
|
"""Get the earliest timestamp for a symbol/timeframe."""
|
|
with self.get_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("""
|
|
SELECT MIN(timestamp)
|
|
FROM ohlcv
|
|
WHERE symbol = ? AND timeframe = ?
|
|
""", (symbol, timeframe))
|
|
result = cursor.fetchone()
|
|
|
|
return pd.to_datetime(result[0]) if result[0] else None
|