142 lines
4.8 KiB
Python
142 lines
4.8 KiB
Python
# database.py
|
|
|
|
import sqlite3
|
|
import time
|
|
|
|
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 (This is the corrected section) ---
|
|
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 (This is the corrected section) ---
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS subreddits (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
name TEXT NOT NULL UNIQUE
|
|
)
|
|
""")
|
|
|
|
# --- Create mentions table with sentiment_score column ---
|
|
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)
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
print("Database initialized successfully.")
|
|
|
|
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 # Ignore duplicate mentions
|
|
|
|
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():
|
|
"""Queries the DB to generate a summary with market caps and avg. sentiment."""
|
|
print("\n--- Summary Report ---")
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
query = """
|
|
SELECT
|
|
t.symbol,
|
|
t.market_cap,
|
|
COUNT(m.id) as mention_count,
|
|
AVG(m.sentiment_score) as avg_sentiment
|
|
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 20;
|
|
"""
|
|
results = cursor.execute(query).fetchall()
|
|
|
|
print(f"{'Ticker':<10} | {'Mentions':<10} | {'Sentiment':<18} | {'Market Cap':<20}")
|
|
print("-" * 65)
|
|
|
|
for row in results:
|
|
# Format Market Cap
|
|
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"
|
|
elif mc >= 1e6: market_cap_str = f"${mc/1e6:.2f}M"
|
|
else: market_cap_str = f"${mc:,}"
|
|
|
|
# Determine Sentiment Label
|
|
sentiment_score = row['avg_sentiment']
|
|
if sentiment_score is not None:
|
|
if sentiment_score > 0.1: sentiment_label = f"Bullish ({sentiment_score:+.2f})"
|
|
elif sentiment_score < -0.1: sentiment_label = f"Bearish ({sentiment_score:+.2f})"
|
|
else: sentiment_label = f"Neutral ({sentiment_score:+.2f})"
|
|
else:
|
|
sentiment_label = "N/A"
|
|
|
|
print(f"{row['symbol']:<10} | {row['mention_count']:<10} | {sentiment_label:<18} | {market_cap_str:<20}")
|
|
|
|
conn.close() |