""" Advanced Local Database Manager for Stock Alchemist Stores data as JSON files organized by date, type, and ticker Uses MySQL for indexing and JSON files for data storage """ import json import mysql.connector from mysql.connector import Error from datetime import date, datetime, timedelta from pathlib import Path from typing import Any, Dict, List, Optional, Union from dataclasses import dataclass, asdict, field from enum import Enum import hashlib import gzip import numpy as np class DataType(Enum): """Supported data types - simplified naming""" # Calendar events (no prefix needed) EARNINGS = "earnings" IPO = "ipo" STOCK_SPLIT = "stock_split" DIVIDENDS = "dividends" ECONOMIC_EVENTS = "economic_events" # Other data types FUNDAMENTAL = "fundamental_analysis" NEWS = "news" TECHNICAL_ANALYSIS = "technical_analysis" @dataclass class DatabaseEntry: """Base class for database entries""" date: str # ISO format YYYY-MM-DD data_type: str # DataType enum value ticker: str data: Dict[str, Any] created_at: str = field(default_factory=lambda: datetime.now().isoformat()) updated_at: str = field(default_factory=lambda: datetime.now().isoformat()) expiry_date: Optional[str] = None metadata: Dict[str, Any] = field(default_factory=dict) def to_dict(self): """Convert to dictionary""" return asdict(self) @classmethod def from_dict(cls, data: Dict): """Create from dictionary""" return cls(**data) def generate_key(self): """ Generate unique key for this entry For calendar events, includes execution_date/ex_date to prevent duplicates """ # For calendar events, include the actual event date to ensure uniqueness if self.data_type in ['earnings', 'ipo', 'stock_split', 'dividends']: event_date = (self.data.get('execution_date') or self.data.get('ex_date') or self.data.get('date') or self.date) key_string = f"{self.data_type}_{self.ticker}_{event_date}" else: key_string = f"{self.date}_{self.data_type}_{self.ticker}" return hashlib.md5(key_string.encode()).hexdigest() class LocalDatabase: """ Advanced local database manager with MySQL index and JSON file storage Features: - MySQL metadata index for fast queries - JSON files for actual data storage - Optional gzip compression for large data - Automatic expiry and cleanup - Date/Type/Ticker indexing - Batch operations support """ def __init__(self, db_dir: str = "database", compress: bool = False): """ Initialize database manager Args: db_dir: Root directory for database storage compress: Whether to compress JSON files with gzip """ self.db_dir = Path(db_dir) try: self.db_dir.mkdir(exist_ok=True) self.data_dir = self.db_dir / "data" self.data_dir.mkdir(exist_ok=True) except (PermissionError, OSError) as e: print(f"⚠️ Error creating database dir at {self.db_dir}: {e}") import tempfile self.db_dir = Path(tempfile.gettempdir()) / "gotti_database" print(f"⚠️ Falling back to temporary directory: {self.db_dir}") self.db_dir.mkdir(exist_ok=True) self.data_dir = self.db_dir / "data" self.data_dir.mkdir(exist_ok=True) # Load environment variables from dotenv import load_dotenv import os load_dotenv() # MySQL connection parameters from environment variables self.mysql_config = { 'host': os.getenv('DB_HOST', 'localhost').strip(), 'user': os.getenv('DB_USERNAME', 'root').strip(), 'password': os.getenv('DB_PASSWORD', '').strip(), 'database': os.getenv('DB_DATABASE', 'gotti').strip(), 'port': int(os.getenv('DB_PORT', 3306)) } # SSL Configuration for TiDB ssl_ca = os.getenv('DB_SSL_CA') if ssl_ca: # Handle if the content itself is passed (begins with ---) if "-----BEGIN CERTIFICATE-----" in ssl_ca: try: import tempfile # Create a temporary file for the certificate # We use a fixed location in /tmp relative to workdir if possible or tempfile # But tempfile is safer. We need it to persist for the process life. # Be careful about file cleanup, but for a container it's fine. # Create a persistent temp file (won't be deleted automatically on close) # We'll save it to a known location for debugging: /tmp/tidb_ca.pem tmp_ca_path = Path("/tmp/tidb_ca.pem") # If on Windows (local dev), assume execution in app dir if os.name == 'nt': tmp_ca_path = Path("tidb_ca.pem") with open(tmp_ca_path, "w", encoding='utf-8') as f: f.write(ssl_ca) self.mysql_config['ssl_ca'] = str(tmp_ca_path) self.mysql_config['ssl_verify_cert'] = True self.mysql_config['ssl_verify_identity'] = True print(f"🔒 Using provided SSL Certificate content (saved to {tmp_ca_path})") except Exception as e: print(f"⚠️ Failed to write SSL CA content: {e}") else: # Resolve relative path if needed (existing logic) if not os.path.isabs(ssl_ca): project_root = Path(__file__).parent.parent.parent ssl_ca_path = project_root / ssl_ca if ssl_ca_path.exists(): self.mysql_config['ssl_ca'] = str(ssl_ca_path) self.mysql_config['ssl_verify_cert'] = True self.mysql_config['ssl_verify_identity'] = True print(f"🔒 Using SSL Certificate from file: {ssl_ca_path}") else: print(f"⚠️ SSL CA file not found at {ssl_ca_path}") elif os.path.exists(ssl_ca): self.mysql_config['ssl_ca'] = ssl_ca self.mysql_config['ssl_verify_cert'] = True self.mysql_config['ssl_verify_identity'] = True self.compress = compress self._init_database() def _create_connection(self): """Create and return a MySQL database connection""" try: connection = mysql.connector.connect(**self.mysql_config) return connection except Error as e: print(f"❌ Error connecting to MySQL: {e}") return None def _get_table_name(self, data_type: str) -> str: """Determine which table to use based on data_type""" # Calendar events if data_type in ['earnings', 'ipo', 'stock_split', 'dividends', 'economic_events']: return 'calendar' # News elif data_type == 'news': return 'news' # Fundamental analysis elif data_type == 'fundamental_analysis': return 'fundamental_analysis' else: raise ValueError(f"Unknown data type: {data_type}") def _init_database(self): """Initialize MySQL tables - three separate tables by data category""" conn = self._create_connection() if not conn: raise Exception("Failed to connect to MySQL database") cursor = conn.cursor() try: # Create calendar table cursor.execute(''' CREATE TABLE IF NOT EXISTS calendar ( entry_key VARCHAR(32) PRIMARY KEY, date DATE NOT NULL, event_type VARCHAR(50) NOT NULL, ticker VARCHAR(20) NOT NULL, data JSON NOT NULL, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, expiry_date DATE, metadata JSON, execution_date DATE, INDEX idx_date (date), INDEX idx_event_type (event_type), INDEX idx_ticker (ticker), INDEX idx_date_event (date, event_type), INDEX idx_ticker_event (ticker, event_type) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ''') # Create news table cursor.execute(''' CREATE TABLE IF NOT EXISTS news ( entry_key VARCHAR(32) PRIMARY KEY, date DATE NOT NULL, ticker VARCHAR(20) NOT NULL, data JSON NOT NULL, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, expiry_date DATE, metadata JSON, INDEX idx_date (date), INDEX idx_ticker (ticker), INDEX idx_date_ticker (date, ticker) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ''') # Create fundamental_analysis table cursor.execute(''' CREATE TABLE IF NOT EXISTS fundamental_analysis ( entry_key VARCHAR(32) PRIMARY KEY, date DATE NOT NULL, ticker VARCHAR(20) NOT NULL, data JSON NOT NULL, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, expiry_date DATE, metadata JSON, INDEX idx_date (date), INDEX idx_ticker (ticker), INDEX idx_date_ticker (date, ticker) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ''') # Create available_tickers table - whitelist of allowed tickers cursor.execute(''' CREATE TABLE IF NOT EXISTS available_tickers ( ticker VARCHAR(20) PRIMARY KEY, name VARCHAR(255), exchange VARCHAR(50), sector VARCHAR(100), is_active BOOLEAN DEFAULT TRUE, added_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, metadata JSON, INDEX idx_is_active (is_active), INDEX idx_exchange (exchange), INDEX idx_sector (sector) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ''') # Create signals table - tracks actionable ticker signals cursor.execute(''' CREATE TABLE IF NOT EXISTS signals ( signal_id VARCHAR(32) PRIMARY KEY, ticker VARCHAR(20) NOT NULL, signal_date DATE NOT NULL, signal_position VARCHAR(10) NOT NULL, calendar_event_keys JSON, news_keys JSON, fundamental_analysis_key VARCHAR(32), sentiment JSON, created_at DATETIME NOT NULL, updated_at DATETIME NOT NULL, metadata JSON, INDEX idx_ticker (ticker), INDEX idx_signal_date (signal_date), INDEX idx_ticker_date (ticker, signal_date), INDEX idx_created_at (created_at) ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 ''') # Add execution_date column to calendar table if it doesn't exist try: cursor.execute("ALTER TABLE calendar ADD COLUMN execution_date DATE AFTER date") print("✅ Added execution_date column to calendar table") except Error as e: if e.errno != 1060: # Error 1060 = Duplicate column name pass # Column already exists, ignore # Add sentiment column to signals table if it doesn't exist try: cursor.execute("ALTER TABLE signals ADD COLUMN sentiment JSON AFTER fundamental_analysis_key") print("✅ Added sentiment column to signals table") except Error as e: if e.errno != 1060: # Error 1060 = Duplicate column name pass # Column already exists, ignore # Add signal_position column to signals table if it doesn't exist try: cursor.execute("ALTER TABLE signals ADD COLUMN signal_position VARCHAR(10) AFTER signal_date") print("✅ Added signal_position column to signals table") except Error as e: if e.errno != 1060: # Error 1060 = Duplicate column name pass # Column already exists, ignore conn.commit() print("✅ MySQL database tables initialized successfully (calendar, news, fundamental_analysis, available_tickers, signals)") except Error as e: print(f"❌ Error initializing database: {e}") # Try without IF NOT EXISTS for MySQL versions that don't support it try: cursor.execute("SHOW COLUMNS FROM calendar LIKE 'execution_date'") if cursor.fetchone() is None: cursor.execute("ALTER TABLE calendar ADD COLUMN execution_date DATE AFTER date") conn.commit() print("✅ Added execution_date column to calendar table") except Exception as alter_error: print(f"⚠️ Could not add execution_date column: {alter_error}") # Try adding sentiment column for older MySQL versions try: cursor.execute("SHOW COLUMNS FROM signals LIKE 'sentiment'") if cursor.fetchone() is None: cursor.execute("ALTER TABLE signals ADD COLUMN sentiment JSON AFTER fundamental_analysis_key") conn.commit() print("✅ Added sentiment column to signals table") except Exception as alter_error: print(f"⚠️ Could not add sentiment column: {alter_error}") # Try adding signal_position column for older MySQL versions try: cursor.execute("SHOW COLUMNS FROM signals LIKE 'signal_position'") if cursor.fetchone() is None: cursor.execute("ALTER TABLE signals ADD COLUMN signal_position VARCHAR(10) AFTER signal_date") conn.commit() print("✅ Added signal_position column to signals table") except Exception as alter_error: print(f"⚠️ Could not add signal_position column: {alter_error}") finally: cursor.close() conn.close() def _generate_file_path(self, entry_key: str, data_type: str, date_str: str) -> Path: """Generate organized file path for data storage""" # Organize by type/year/month/ year_month = datetime.fromisoformat(date_str).strftime("%Y/%m") type_dir = self.data_dir / data_type / year_month type_dir.mkdir(parents=True, exist_ok=True) extension = ".json.gz" if self.compress else ".json" return type_dir / f"{entry_key}{extension}" def _write_json(self, file_path: Path, data: Dict): """Write JSON data with optional compression""" json_str = json.dumps(data, indent=2, default=str) if self.compress: with gzip.open(file_path, 'wt', encoding='utf-8') as f: f.write(json_str) else: with open(file_path, 'w', encoding='utf-8') as f: f.write(json_str) def _read_json(self, file_path: Path, compressed: bool) -> Dict: """Read JSON data with optional decompression""" if compressed: with gzip.open(file_path, 'rt', encoding='utf-8') as f: return json.load(f) else: with open(file_path, 'r', encoding='utf-8') as f: return json.load(f) def _clean_data_for_json(self, data: Any) -> Any: """ Recursively clean data to ensure it's JSON serializable. - Converts NaN/Infinity to None - Converts numpy types to native python types """ if isinstance(data, dict): return {k: self._clean_data_for_json(v) for k, v in data.items()} elif isinstance(data, list): return [self._clean_data_for_json(v) for v in data] elif isinstance(data, float): if np.isnan(data) or np.isinf(data): return None return float(data) elif isinstance(data, np.integer): return int(data) elif isinstance(data, np.floating): if np.isnan(data) or np.isinf(data): return None return float(data) elif isinstance(data, np.ndarray): return self._clean_data_for_json(data.tolist()) return data def add_ticker(self, ticker: str, name: str = None, exchange: str = None, sector: str = None, metadata: Dict = None) -> bool: """ Add a ticker to the available_tickers whitelist. Args: ticker: Ticker symbol name: Company name exchange: Exchange name (e.g., 'NASDAQ', 'NYSE') sector: Company sector metadata: Additional metadata as JSON Returns: True if successful, False otherwise """ conn = self._create_connection() if not conn: return False cursor = conn.cursor() try: now = datetime.now() cursor.execute(''' INSERT INTO available_tickers (ticker, name, exchange, sector, is_active, added_at, updated_at, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE name = VALUES(name), exchange = VALUES(exchange), sector = VALUES(sector), is_active = VALUES(is_active), updated_at = VALUES(updated_at), metadata = VALUES(metadata) ''', ( ticker.upper(), name, exchange, sector, True, now, now, json.dumps(metadata) if metadata else None )) conn.commit() return True except Error as e: print(f"❌ Error adding ticker {ticker}: {e}") return False finally: cursor.close() conn.close() def get_macroeconomic_indicators(self) -> Dict[str, Any]: """ Retrieve macroeconomic indicators from the database. Returns: Dictionary of macroeconomic indicators """ conn = self._create_connection() if not conn: return {} cursor = conn.cursor() try: cursor.execute(''' SELECT data FROM macroeconomic_indicators ORDER BY date DESC LIMIT 1 ''') row = cursor.fetchone() if row: return json.loads(row[0]) return {} except Error as e: print(f"❌ Error fetching macroeconomic indicators: {e}") return {} finally: cursor.close() conn.close() def remove_ticker(self, ticker: str) -> bool: """ Deactivate a ticker (soft delete - sets is_active to False). Args: ticker: Ticker symbol to deactivate Returns: True if successful, False otherwise """ conn = self._create_connection() if not conn: return False cursor = conn.cursor() try: cursor.execute( "UPDATE available_tickers SET is_active = FALSE, updated_at = %s WHERE ticker = %s", (datetime.now(), ticker.upper()) ) conn.commit() return cursor.rowcount > 0 except Error as e: print(f"❌ Error removing ticker {ticker}: {e}") return False finally: cursor.close() conn.close() def get_all_available_tickers(self) -> List[str]: """ Get all active tickers from the whitelist. Returns: List of ticker symbols """ conn = self._create_connection() if not conn: return [] cursor = conn.cursor() try: cursor.execute("SELECT ticker FROM available_tickers WHERE is_active = TRUE ORDER BY ticker") return [row[0] for row in cursor.fetchall()] except Error as e: print(f"❌ Error fetching available tickers: {e}") return [] finally: cursor.close() conn.close() def is_ticker_available(self, ticker: str) -> bool: """ Check if ticker is in the available_tickers whitelist. Args: ticker: Ticker symbol to check Returns: True if ticker is available and active, False otherwise """ conn = self._create_connection() if not conn: return False cursor = conn.cursor() try: cursor.execute( "SELECT is_active FROM available_tickers WHERE ticker = %s", (ticker.upper(),) ) result = cursor.fetchone() if result and result[0]: # Ticker exists and is_active = True return True return False except Error as e: print(f"❌ Error checking ticker availability for {ticker}: {e}") return False finally: cursor.close() conn.close() def save(self, entry: DatabaseEntry, expiry_days: Optional[int] = None) -> bool: """ Save entry to database. Updates existing entry if duplicate is found. IMPORTANT: Checks if ticker is in available_tickers whitelist before saving. Args: entry: DatabaseEntry to save expiry_days: Optional expiry in days Returns: True if successful, False if ticker not available or save fails """ try: # CRITICAL: Check if ticker is in the available_tickers whitelist # Skip check for economic events as they use country names as tickers if entry.data_type != DataType.ECONOMIC_EVENTS.value and not self.is_ticker_available(entry.ticker): print(f"⚠️ Skipping {entry.data_type} for {entry.ticker} - ticker not in available_tickers whitelist") return False entry_key = entry.generate_key() # Get the appropriate table name table_name = self._get_table_name(entry.data_type) # Check if entry already exists conn = self._create_connection() if not conn: return False cursor = conn.cursor() cursor.execute(f'SELECT created_at FROM {table_name} WHERE entry_key = %s', (entry_key,)) existing = cursor.fetchone() # Preserve original created_at if updating if existing: entry.created_at = str(existing[0]) # Update the updated_at timestamp entry.updated_at = datetime.now().isoformat() # Calculate expiry date if specified if expiry_days: expiry_date = (datetime.now() + timedelta(days=expiry_days)).date().isoformat() entry.expiry_date = expiry_date # Store data directly as JSON in database # Different INSERT statement based on table structure if table_name == 'calendar': # Calendar table has event_type column cursor.execute(''' INSERT INTO calendar (entry_key, date, event_type, ticker, data, created_at, updated_at, expiry_date, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE data = VALUES(data), updated_at = VALUES(updated_at), expiry_date = VALUES(expiry_date), metadata = VALUES(metadata) ''', ( entry_key, entry.date, entry.data_type, # event_type (earnings, ipo, etc.) entry.ticker, json.dumps(self._clean_data_for_json(entry.data), default=str), entry.created_at, entry.updated_at, entry.expiry_date, json.dumps(entry.metadata) )) else: # News and fundamental_analysis tables don't have event_type cursor.execute(f''' INSERT INTO {table_name} (entry_key, date, ticker, data, created_at, updated_at, expiry_date, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE data = VALUES(data), updated_at = VALUES(updated_at), expiry_date = VALUES(expiry_date), metadata = VALUES(metadata) ''', ( entry_key, entry.date, entry.ticker, json.dumps(self._clean_data_for_json(entry.data), default=str), entry.created_at, entry.updated_at, entry.expiry_date, json.dumps(entry.metadata) )) conn.commit() return True except Exception as e: print(f"Error saving entry {entry.ticker}: {e}") return False finally: if 'cursor' in locals(): cursor.close() if 'conn' in locals() and conn: conn.close() def save_batch(self, entries: List[DatabaseEntry], expiry_days: Optional[int] = None) -> int: """ Save multiple entries in batch. Updates existing entries if duplicates are found. Args: entries: List of DatabaseEntry objects expiry_days: Optional expiry in days Returns: Number of successfully saved entries """ success_count = 0 conn = self._create_connection() if not conn: return 0 cursor = conn.cursor() try: for entry in entries: try: # CRITICAL: Check if ticker is in the available_tickers whitelist # Skip check for economic events as they use country names as tickers if entry.data_type != DataType.ECONOMIC_EVENTS.value and not self.is_ticker_available(entry.ticker): print(f"⚠️ Skipping {entry.data_type} for {entry.ticker} - ticker not in available_tickers whitelist") continue entry_key = entry.generate_key() table_name = self._get_table_name(entry.data_type) # Check if entry already exists cursor.execute(f'SELECT created_at FROM {table_name} WHERE entry_key = %s', (entry_key,)) existing = cursor.fetchone() # Preserve original created_at if updating if existing: entry.created_at = str(existing[0]) # Update the updated_at timestamp entry.updated_at = datetime.now().isoformat() if expiry_days: expiry_date = (datetime.now() + timedelta(days=expiry_days)).date().isoformat() entry.expiry_date = expiry_date # Store data directly in database - different format for calendar if table_name == 'calendar': cursor.execute(''' INSERT INTO calendar (entry_key, date, event_type, ticker, data, created_at, updated_at, expiry_date, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE data = VALUES(data), updated_at = VALUES(updated_at), expiry_date = VALUES(expiry_date), metadata = VALUES(metadata) ''', ( entry_key, entry.date, entry.data_type, # event_type entry.ticker, json.dumps(self._clean_data_for_json(entry.data), default=str), entry.created_at, entry.updated_at, entry.expiry_date, json.dumps(entry.metadata) )) else: cursor.execute(f''' INSERT INTO {table_name} (entry_key, date, ticker, data, created_at, updated_at, expiry_date, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE data = VALUES(data), updated_at = VALUES(updated_at), expiry_date = VALUES(expiry_date), metadata = VALUES(metadata) ''', ( entry_key, entry.date, entry.ticker, json.dumps(self._clean_data_for_json(entry.data), default=str), entry.created_at, entry.updated_at, entry.expiry_date, json.dumps(entry.metadata) )) success_count += 1 except Exception as e: print(f"Error saving entry {entry.ticker}: {e}") continue conn.commit() finally: cursor.close() conn.close() return success_count def save_signal(self, ticker: str, calendar_event_keys: List[str], news_keys: List[str], fundamental_key: str, signal_position: str, sentiment: Dict = None) -> bool: """Save a generated signal to the database""" try: conn = self._create_connection() if not conn: return False cursor = conn.cursor() signal_date = datetime.now().date().isoformat() signal_id = hashlib.md5(f"{ticker}_{signal_date}".encode()).hexdigest() now = datetime.now() # Merge with existing sentiment if provided final_sentiment = sentiment if sentiment: cursor.execute("SELECT sentiment FROM signals WHERE signal_id = %s", (signal_id,)) existing = cursor.fetchone() if existing and existing[0]: existing_sentiment = json.loads(existing[0]) if isinstance(existing[0], str) else existing[0] if isinstance(existing_sentiment, dict): existing_sentiment.update(sentiment) final_sentiment = existing_sentiment cursor.execute(''' INSERT INTO signals (signal_id, ticker, signal_date, signal_position, calendar_event_keys, news_keys, fundamental_analysis_key, sentiment, created_at, updated_at, metadata) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s) ON DUPLICATE KEY UPDATE signal_position = VALUES(signal_position), calendar_event_keys = VALUES(calendar_event_keys), news_keys = VALUES(news_keys), fundamental_analysis_key = VALUES(fundamental_analysis_key), sentiment = VALUES(sentiment), updated_at = VALUES(updated_at) ''', ( signal_id, ticker, signal_date, signal_position, json.dumps(calendar_event_keys), json.dumps(news_keys), fundamental_key, json.dumps(final_sentiment) if final_sentiment else None, now, now, json.dumps({}) )) conn.commit() return True except Error as e: print(f"❌ Error saving signal: {e}") return False finally: if 'cursor' in locals(): cursor.close() if 'conn' in locals(): conn.close() def get_signal(self, ticker: str, date_str: str = None) -> Optional[Dict]: """Get signal for a ticker on a specific date (defaults to today)""" try: if not date_str: date_str = datetime.now().date().isoformat() conn = self._create_connection() if not conn: return None cursor = conn.cursor(dictionary=True) cursor.execute(''' SELECT * FROM signals WHERE ticker = %s AND signal_date = %s ''', (ticker, date_str)) result = cursor.fetchone() if result: # Parse JSON fields for field in ['calendar_event_keys', 'news_keys', 'metadata', 'sentiment']: if result.get(field): if isinstance(result[field], str): result[field] = json.loads(result[field]) return result except Error as e: print(f"❌ Error getting signal: {e}") return None finally: if 'cursor' in locals(): cursor.close() if 'conn' in locals(): conn.close() def get_recent_signals(self, limit: int = 50) -> List[Dict]: """Get most recent signals""" try: conn = self._create_connection() if not conn: return [] cursor = conn.cursor(dictionary=True) cursor.execute(''' SELECT * FROM signals ORDER BY created_at DESC LIMIT %s ''', (limit,)) results = cursor.fetchall() parsed_results = [] for result in results: # Parse JSON fields for field in ['calendar_event_keys', 'news_keys', 'metadata', 'sentiment']: if result.get(field) and isinstance(result[field], str): try: result[field] = json.loads(result[field]) except: pass parsed_results.append(result) return parsed_results except Error as e: print(f"❌ Error getting recent signals: {e}") return [] finally: if 'cursor' in locals(): cursor.close() if 'conn' in locals(): conn.close() def get(self, date_str: str, data_type: str, ticker: str) -> Optional[DatabaseEntry]: """ Retrieve entry by date, type, and ticker Args: date_str: Date in YYYY-MM-DD format data_type: Data type (earnings, ipo, news, fundamental_analysis, etc.) ticker: Stock ticker Returns: DatabaseEntry if found, None otherwise """ try: # Generate key key_string = f"{date_str}_{data_type}_{ticker}" entry_key = hashlib.md5(key_string.encode()).hexdigest() # Get table name for this data type table_name = self._get_table_name(data_type) conn = self._create_connection() if not conn: return None cursor = conn.cursor() # Different SELECT based on table structure if table_name == 'calendar': cursor.execute(''' SELECT date, event_type, ticker, data, created_at, updated_at, expiry_date, metadata FROM calendar WHERE entry_key = %s ''', (entry_key,)) else: cursor.execute(f''' SELECT date, ticker, data, created_at, updated_at, expiry_date, metadata FROM {table_name} WHERE entry_key = %s ''', (entry_key,)) result = cursor.fetchone() cursor.close() conn.close() if not result: return None # Parse result based on table structure if table_name == 'calendar': date_val, event_type_val, ticker_val, data_json, created_at, updated_at, expiry_date, metadata_json = result data_type_val = event_type_val # event_type is the data_type else: date_val, ticker_val, data_json, created_at, updated_at, expiry_date, metadata_json = result data_type_val = data_type # Use the data_type parameter # Check if expired if expiry_date: if str(expiry_date) < datetime.now().date().isoformat(): return None # Parse JSON data from database data_dict = json.loads(data_json) if isinstance(data_json, str) else data_json metadata_dict = json.loads(metadata_json) if isinstance(metadata_json, str) else (metadata_json or {}) # Create DatabaseEntry entry = DatabaseEntry( date=str(date_val), data_type=str(data_type_val), ticker=str(ticker_val), data=data_dict, created_at=str(created_at), updated_at=str(updated_at), expiry_date=str(expiry_date) if expiry_date else None, metadata=metadata_dict ) return entry except Exception as e: print(f"Error retrieving entry: {e}") return None def query(self, date_from: Optional[str] = None, date_to: Optional[str] = None, data_type: Optional[str] = None, ticker: Optional[str] = None, limit: Optional[int] = None, include_expired: bool = False) -> List[DatabaseEntry]: """ Query database with flexible filters across all tables Args: date_from: Start date (inclusive) date_to: End date (inclusive) data_type: Filter by data type (e.g., 'earnings', 'news', 'fundamental_analysis') ticker: Filter by ticker limit: Max results include_expired: Whether to include expired entries Returns: List of DatabaseEntry objects """ try: conn = self._create_connection() if not conn: return [] cursor = conn.cursor() entries = [] # Determine which tables to query tables_to_query = [] if data_type: # Query specific table based on data_type table_name = self._get_table_name(data_type) tables_to_query.append((table_name, data_type)) else: # Query all tables tables_to_query = [ ('calendar', None), # Will get all calendar events ('news', 'news'), ('fundamental_analysis', 'fundamental_analysis') ] # Query each table for table_name, specific_type in tables_to_query: params = [] if table_name == 'calendar': # Calendar has event_type column query = "SELECT date, event_type, ticker, data, created_at, updated_at, expiry_date, metadata FROM calendar WHERE 1=1" if specific_type: # Specific calendar event type query += " AND event_type = %s" params.append(specific_type) else: # News and fundamental_analysis don't have event_type query = f"SELECT date, ticker, data, created_at, updated_at, expiry_date, metadata FROM {table_name} WHERE 1=1" if date_from: query += " AND date >= %s" params.append(date_from) if date_to: query += " AND date <= %s" params.append(date_to) if ticker: query += " AND ticker = %s" params.append(ticker) if not include_expired: query += " AND (expiry_date IS NULL OR expiry_date >= %s)" params.append(datetime.now().date().isoformat()) query += " ORDER BY date DESC, created_at DESC" if limit and len(tables_to_query) == 1: # Only apply limit if querying a single table query += f" LIMIT {limit}" cursor.execute(query, tuple(params)) results = cursor.fetchall() # Parse results based on table structure for row in results: try: if table_name == 'calendar': date_val, event_type_val, ticker_val, data_json, created_at, updated_at, expiry_date, metadata_json = row data_type_val = event_type_val else: date_val, ticker_val, data_json, created_at, updated_at, expiry_date, metadata_json = row data_type_val = specific_type # Parse JSON data data_dict = json.loads(data_json) if isinstance(data_json, str) else data_json metadata_dict = json.loads(metadata_json) if isinstance(metadata_json, str) else (metadata_json or {}) entry = DatabaseEntry( date=str(date_val), data_type=str(data_type_val), ticker=str(ticker_val), data=data_dict, created_at=str(created_at), updated_at=str(updated_at), expiry_date=str(expiry_date) if expiry_date else None, metadata=metadata_dict ) entries.append(entry) except Exception as e: print(f"Error loading entry from {table_name}: {e}") continue cursor.close() conn.close() # Sort all entries by date and apply limit if needed entries.sort(key=lambda x: (x.date, x.created_at), reverse=True) if limit: entries = entries[:limit] return entries except Exception as e: print(f"Error querying database: {e}") return [] def delete(self, date_str: str, data_type: str, ticker: str) -> bool: """Delete entry by date, type, and ticker""" try: # Generate key key_string = f"{date_str}_{data_type}_{ticker}" entry_key = hashlib.md5(key_string.encode()).hexdigest() table_name = self._get_table_name(data_type) conn = self._create_connection() if not conn: return False cursor = conn.cursor() cursor.execute(f'DELETE FROM {table_name} WHERE entry_key = %s', (entry_key,)) conn.commit() conn.close() return True except Exception as e: print(f"Error deleting entry: {e}") return False def clean_expired(self) -> int: """Remove expired entries""" try: conn = self._create_connection() if not conn: return 0 cursor = conn.cursor() total_cleaned = 0 for table_name in ['calendar', 'news', 'fundamental_analysis']: cursor.execute(f''' DELETE FROM {table_name} WHERE expiry_date IS NOT NULL AND expiry_date < %s ''', (datetime.now().date().isoformat(),)) total_cleaned += cursor.rowcount conn.commit() conn.close() print(f"✓ Cleaned {total_cleaned} expired entries") return total_cleaned except Exception as e: print(f"Error cleaning expired entries: {e}") return 0 def get_stats(self) -> Dict[str, Any]: """Get database statistics across all tables""" try: conn = self._create_connection() if not conn: return {} cursor = conn.cursor() # Initialize counters total_entries = 0 by_type = {} all_tickers = {} total_size = 0 expired_count = 0 min_date = None max_date = None # Query each table for table_name in ['calendar', 'news', 'fundamental_analysis']: # Count entries cursor.execute(f'SELECT COUNT(*) FROM {table_name}') table_count = cursor.fetchone()[0] total_entries += table_count if table_name == 'calendar': # Get counts by event_type cursor.execute('SELECT event_type, COUNT(*) FROM calendar GROUP BY event_type') for event_type, count in cursor.fetchall(): by_type[event_type] = count else: # For news and fundamental_analysis, use table name as type by_type[table_name] = table_count # Get ticker counts cursor.execute(f'SELECT ticker, COUNT(*) FROM {table_name} GROUP BY ticker') for ticker, count in cursor.fetchall(): all_tickers[ticker] = all_tickers.get(ticker, 0) + count # Get data size cursor.execute(f'SELECT SUM(LENGTH(data)) FROM {table_name}') table_size = cursor.fetchone()[0] or 0 total_size += table_size # Count expired entries cursor.execute(f''' SELECT COUNT(*) FROM {table_name} WHERE expiry_date IS NOT NULL AND expiry_date < %s ''', (datetime.now().date().isoformat(),)) expired_count += cursor.fetchone()[0] # Get date range cursor.execute(f'SELECT MIN(date), MAX(date) FROM {table_name}') table_date_range = cursor.fetchone() if table_date_range[0]: if min_date is None or table_date_range[0] < min_date: min_date = table_date_range[0] if max_date is None or table_date_range[1] > max_date: max_date = table_date_range[1] # Get top 10 tickers top_tickers = dict(sorted(all_tickers.items(), key=lambda x: x[1], reverse=True)[:10]) conn.close() stats = { 'total_entries': total_entries, 'by_type': by_type, 'top_tickers': top_tickers, 'total_size_bytes': total_size, 'total_size_mb': round(total_size / (1024 * 1024), 2), 'expired_entries': expired_count, 'date_range': {'from': str(min_date), 'to': str(max_date)} if min_date else None, 'compression': 'enabled' if self.compress else 'disabled' } return stats except Exception as e: print(f"Error getting stats: {e}") return {} def clear_all(self) -> bool: """Clear all data (use with caution!)""" try: conn = self._create_connection() if not conn: return False cursor = conn.cursor() # Truncate all tables for table_name in ['calendar', 'news', 'fundamental_analysis', 'signals']: cursor.execute(f'TRUNCATE TABLE {table_name}') conn.commit() conn.close() print("✓ All data cleared") return True except Exception as e: print(f"Error clearing data: {e}") return False # Global database instance db_instance = None def get_database(db_dir: str = "database", compress: bool = False) -> LocalDatabase: """Get or create database instance""" global db_instance if db_instance is None or db_instance.db_dir != Path(db_dir): db_instance = LocalDatabase(db_dir=db_dir, compress=compress) return db_instance