Modularized the tool.

rstat_tool/database.py

@@ -0,0 +1,158 @@
+# rstat_tool/database.py
+
+import sqlite3
+import time
+# --- IMPORT ADDED BACK IN ---
+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():
+    # ... (This function is unchanged)
+    conn = get_db_connection()
+    cursor = conn.cursor()
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS tickers (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            symbol TEXT NOT NULL UNIQUE,
+            market_cap INTEGER,
+            last_updated INTEGER
+        )
+    """)
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS subreddits (
+            id INTEGER PRIMARY KEY AUTOINCREMENT,
+            name TEXT NOT NULL UNIQUE
+        )
+    """)
+    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)
+        )
+    """)
+    conn.commit()
+    conn.close()
+    print("Database initialized successfully.")
+
+# --- CLEANUP FUNCTION ADDED BACK IN ---
+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()
+
+    # Find ticker IDs that match the blacklist
+    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})...")
+        
+        # 1. Delete all mentions associated with this ticker ID
+        cursor.execute("DELETE FROM mentions WHERE ticker_id = ?", (ticker_id,))
+        
+        # 2. Delete the ticker itself
+        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):
+    # ... (This function is unchanged)
+    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
+
+# ... (get_or_create_entity, update_ticker_market_cap, get_ticker_info are unchanged)
+def get_or_create_entity(conn, table_name, column_name, value):
+    # ...
+    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):
+    # ...
+    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):
+    # ...
+    cursor = conn.cursor()
+    cursor.execute("SELECT * FROM tickers WHERE id = ?", (ticker_id,))
+    return cursor.fetchone()
+
+def generate_summary_report(limit=20):
+    # ... (This function is unchanged)
+    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()