# rstat_tool/database.py import sqlite3 import time from .ticker_extractor import COMMON_WORDS_BLACKLIST DB_FILE = "reddit_stocks.db" def get_db_connection(): """Establishes a connection to the SQLite database.""" conn = sqlite3.connect(DB_FILE) conn.row_factory = sqlite3.Row return conn def initialize_db(): """ Initializes the database and creates the necessary tables if they don't exist. """ conn = get_db_connection() cursor = conn.cursor() # --- Create tickers table --- cursor.execute(""" CREATE TABLE IF NOT EXISTS tickers ( id INTEGER PRIMARY KEY AUTOINCREMENT, symbol TEXT NOT NULL UNIQUE, market_cap INTEGER, last_updated INTEGER ) """) # --- Create subreddits table --- cursor.execute(""" CREATE TABLE IF NOT EXISTS subreddits ( id INTEGER PRIMARY KEY AUTOINCREMENT, name TEXT NOT NULL UNIQUE ) """) # --- Create mentions table --- cursor.execute(""" CREATE TABLE IF NOT EXISTS mentions ( id INTEGER PRIMARY KEY AUTOINCREMENT, ticker_id INTEGER, subreddit_id INTEGER, post_id TEXT NOT NULL, mention_timestamp INTEGER NOT NULL, sentiment_score REAL, FOREIGN KEY (ticker_id) REFERENCES tickers (id), FOREIGN KEY (subreddit_id) REFERENCES subreddits (id), UNIQUE(ticker_id, post_id, sentiment_score) ) """) cursor.execute(""" CREATE TABLE IF NOT EXISTS posts ( id INTEGER PRIMARY KEY AUTOINCREMENT, post_id TEXT NOT NULL UNIQUE, title TEXT NOT NULL, post_url TEXT, subreddit_id INTEGER, post_timestamp INTEGER, comment_count INTEGER, avg_comment_sentiment REAL, FOREIGN KEY (subreddit_id) REFERENCES subreddits (id) ) """) conn.commit() conn.close() print("Database initialized successfully.") def clean_stale_tickers(): """ Removes tickers and their associated mentions from the database if the ticker symbol exists in the COMMON_WORDS_BLACKLIST. """ print("\n--- Cleaning Stale Tickers from Database ---") conn = get_db_connection() cursor = conn.cursor() placeholders = ','.join('?' for _ in COMMON_WORDS_BLACKLIST) query = f"SELECT id, symbol FROM tickers WHERE symbol IN ({placeholders})" cursor.execute(query, tuple(COMMON_WORDS_BLACKLIST)) stale_tickers = cursor.fetchall() if not stale_tickers: print("No stale tickers to clean.") conn.close() return for ticker in stale_tickers: ticker_id = ticker['id'] ticker_symbol = ticker['symbol'] print(f"Removing stale ticker '{ticker_symbol}' (ID: {ticker_id})...") cursor.execute("DELETE FROM mentions WHERE ticker_id = ?", (ticker_id,)) cursor.execute("DELETE FROM tickers WHERE id = ?", (ticker_id,)) deleted_count = conn.total_changes conn.commit() conn.close() print(f"Cleanup complete. Removed {deleted_count} records.") def add_mention(conn, ticker_id, subreddit_id, post_id, timestamp, sentiment): """Adds a new mention with its sentiment score to the database.""" cursor = conn.cursor() try: cursor.execute( "INSERT INTO mentions (ticker_id, subreddit_id, post_id, mention_timestamp, sentiment_score) VALUES (?, ?, ?, ?, ?)", (ticker_id, subreddit_id, post_id, timestamp, sentiment) ) conn.commit() except sqlite3.IntegrityError: pass def get_or_create_entity(conn, table_name, column_name, value): """Generic function to get or create an entity and return its ID.""" cursor = conn.cursor() cursor.execute(f"SELECT id FROM {table_name} WHERE {column_name} = ?", (value,)) result = cursor.fetchone() if result: return result['id'] else: cursor.execute(f"INSERT INTO {table_name} ({column_name}) VALUES (?)", (value,)) conn.commit() return cursor.lastrowid def update_ticker_market_cap(conn, ticker_id, market_cap): """Updates the market cap and timestamp for a specific ticker.""" cursor = conn.cursor() current_timestamp = int(time.time()) cursor.execute("UPDATE tickers SET market_cap = ?, last_updated = ? WHERE id = ?", (market_cap, current_timestamp, ticker_id)) conn.commit() def get_ticker_info(conn, ticker_id): """Retrieves all info for a specific ticker by its ID.""" cursor = conn.cursor() cursor.execute("SELECT * FROM tickers WHERE id = ?", (ticker_id,)) return cursor.fetchone() def generate_summary_report(limit=20): """Queries the DB to generate a summary for the command-line tool.""" print(f"\n--- Top {limit} Tickers by Mention Count ---") conn = get_db_connection() cursor = conn.cursor() query = """ SELECT t.symbol, t.market_cap, COUNT(m.id) as mention_count, SUM(CASE WHEN m.sentiment_score > 0.1 THEN 1 ELSE 0 END) as bullish_mentions, SUM(CASE WHEN m.sentiment_score < -0.1 THEN 1 ELSE 0 END) as bearish_mentions, SUM(CASE WHEN m.sentiment_score BETWEEN -0.1 AND 0.1 THEN 1 ELSE 0 END) as neutral_mentions FROM mentions m JOIN tickers t ON m.ticker_id = t.id GROUP BY t.symbol, t.market_cap ORDER BY mention_count DESC LIMIT ?; """ results = cursor.execute(query, (limit,)).fetchall() header = f"{'Ticker':<8} | {'Mentions':<8} | {'Bullish':<8} | {'Bearish':<8} | {'Neutral':<8} | {'Market Cap':<15}" print(header) print("-" * len(header)) for row in results: market_cap_str = "N/A" if row['market_cap'] and row['market_cap'] > 0: mc = row['market_cap'] if mc >= 1e12: market_cap_str = f"${mc/1e12:.2f}T" elif mc >= 1e9: market_cap_str = f"${mc/1e9:.2f}B" else: market_cap_str = f"${mc/1e6:.2f}M" print(f"{row['symbol']:<8} | {row['mention_count']:<8} | {row['bullish_mentions']:<8} | {row['bearish_mentions']:<8} | {row['neutral_mentions']:<8} | {market_cap_str:<15}") conn.close() def get_overall_summary(limit=50): """Gets the top tickers across all subreddits for the dashboard.""" conn = get_db_connection() query = """ SELECT t.symbol, t.market_cap, COUNT(m.id) as mention_count, SUM(CASE WHEN m.sentiment_score > 0.1 THEN 1 ELSE 0 END) as bullish_mentions, SUM(CASE WHEN m.sentiment_score < -0.1 THEN 1 ELSE 0 END) as bearish_mentions, SUM(CASE WHEN m.sentiment_score BETWEEN -0.1 AND 0.1 THEN 1 ELSE 0 END) as neutral_mentions FROM mentions m JOIN tickers t ON m.ticker_id = t.id GROUP BY t.symbol, t.market_cap ORDER BY mention_count DESC LIMIT ?; """ results = conn.execute(query, (limit,)).fetchall() conn.close() return results def get_subreddit_summary(subreddit_name, limit=50): """Gets the top tickers for a specific subreddit for the dashboard.""" conn = get_db_connection() query = """ SELECT t.symbol, t.market_cap, COUNT(m.id) as mention_count, SUM(CASE WHEN m.sentiment_score > 0.1 THEN 1 ELSE 0 END) as bullish_mentions, SUM(CASE WHEN m.sentiment_score < -0.1 THEN 1 ELSE 0 END) as bearish_mentions, SUM(CASE WHEN m.sentiment_score BETWEEN -0.1 AND 0.1 THEN 1 ELSE 0 END) as neutral_mentions FROM mentions m JOIN tickers t ON m.ticker_id = t.id JOIN subreddits s ON m.subreddit_id = s.id WHERE s.name = ? GROUP BY t.symbol, t.market_cap ORDER BY mention_count DESC LIMIT ?; """ results = conn.execute(query, (subreddit_name, limit)).fetchall() conn.close() return results def get_all_scanned_subreddits(): """Gets a unique list of all subreddits we have data for.""" # --- THIS IS THE CORRECTED LINE --- conn = get_db_connection() results = conn.execute("SELECT DISTINCT name FROM subreddits ORDER BY name ASC;").fetchall() conn.close() return [row['name'] for row in results] def add_or_update_post_analysis(conn, post_data): """ Inserts a new post analysis record or updates an existing one. This prevents duplicate entries for the same post. """ cursor = conn.cursor() # Use the UNIQUE post_id to replace old data with new on conflict cursor.execute( """ INSERT INTO posts (post_id, title, post_url, subreddit_id, post_timestamp, comment_count, avg_comment_sentiment) VALUES (:post_id, :title, :post_url, :subreddit_id, :post_timestamp, :comment_count, :avg_comment_sentiment) ON CONFLICT(post_id) DO UPDATE SET comment_count = excluded.comment_count, avg_comment_sentiment = excluded.avg_comment_sentiment; """, post_data ) conn.commit() def get_deep_dive_details(ticker_symbol): """ Gets all analyzed posts that mention a specific ticker. """ conn = get_db_connection() query = """ SELECT DISTINCT p.*, s.name as subreddit_name FROM posts p JOIN mentions m ON p.post_id = m.post_id JOIN tickers t ON m.ticker_id = t.id JOIN subreddits s ON p.subreddit_id = s.id WHERE t.symbol = ? ORDER BY p.post_timestamp DESC; """ results = conn.execute(query, (ticker_symbol,)).fetchall() conn.close() return results