136 lines
4.2 KiB
Python
136 lines
4.2 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 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)
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
conn.close()
|
|
print("Database initialized successfully.")
|
|
|
|
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 add_mention(conn, ticker_id, subreddit_id, post_id, timestamp):
|
|
"""Adds a new mention to the database, ignoring duplicates."""
|
|
cursor = conn.cursor()
|
|
try:
|
|
cursor.execute(
|
|
"INSERT INTO mentions (ticker_id, subreddit_id, post_id, mention_timestamp) VALUES (?, ?, ?, ?)",
|
|
(ticker_id, subreddit_id, post_id, timestamp)
|
|
)
|
|
conn.commit()
|
|
except sqlite3.IntegrityError:
|
|
pass
|
|
|
|
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 and print a summary with market caps."""
|
|
print("\n--- Summary Report ---")
|
|
conn = get_db_connection()
|
|
cursor = conn.cursor()
|
|
|
|
query = """
|
|
SELECT
|
|
t.symbol,
|
|
t.market_cap,
|
|
COUNT(m.id) as mention_count
|
|
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} | {'Market Cap':<20}")
|
|
print("-" * 45)
|
|
|
|
for row in results:
|
|
market_cap_str = "N/A"
|
|
if row['market_cap']:
|
|
# Format market cap into a readable string (e.g., $1.23T, $45.6B, $123.4M)
|
|
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:,}"
|
|
|
|
print(f"{row['symbol']:<10} | {row['mention_count']:<10} | {market_cap_str:<20}")
|
|
|
|
conn.close() |