198 lines
7.6 KiB
Python
198 lines
7.6 KiB
Python
# 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)
|
|
)
|
|
""")
|
|
|
|
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] |