Spaces:
Running
Running
File size: 6,933 Bytes
5a30a94 9321b3c 5a30a94 9321b3c 5a30a94 9321b3c 5a30a94 9321b3c 5a30a94 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """SQLite database helpers for the music dashboard."""
import sqlite3
import os
from datetime import datetime, timezone, timedelta
from pathlib import Path
from dotenv import load_dotenv
load_dotenv()
DB_PATH = Path(__file__).parent.parent / "db" / "music_dashboard.db"
SCHEMA_PATH = Path(__file__).parent.parent / "db" / "schema.sql"
def get_connection():
"""Get a SQLite connection with WAL mode for better concurrency."""
conn = sqlite3.connect(str(DB_PATH))
conn.row_factory = sqlite3.Row
conn.execute("PRAGMA journal_mode=WAL")
conn.execute("PRAGMA foreign_keys=ON")
return conn
def init_db():
"""Initialize the database from schema.sql and run migrations."""
os.makedirs(DB_PATH.parent, exist_ok=True)
conn = get_connection()
with open(SCHEMA_PATH, "r") as f:
conn.executescript(f.read())
# Migration: add cycle_id if missing
cols = [row["name"] for row in conn.execute("PRAGMA table_info(play_snapshots)").fetchall()]
if "cycle_id" not in cols:
conn.execute("ALTER TABLE play_snapshots ADD COLUMN cycle_id TEXT")
conn.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_cycle ON play_snapshots(cycle_id)")
conn.commit()
print("[migration] Added cycle_id to play_snapshots")
# Migration: add ytmusic_play_count if missing
cols = [row["name"] for row in conn.execute("PRAGMA table_info(play_snapshots)").fetchall()]
if "ytmusic_play_count" not in cols:
conn.execute("ALTER TABLE play_snapshots ADD COLUMN ytmusic_play_count INTEGER")
conn.commit()
print("[migration] Added ytmusic_play_count to play_snapshots")
# Composite indexes for the dashboard's hot queries (latest snapshot per
# song by platform, songs by artist+platform, alerts by platform)
conn.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_platform_song ON play_snapshots(platform, song_id, id)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_snapshots_song_platform_time ON play_snapshots(song_id, platform, collected_at)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_songs_artist_platform ON songs(artist_id, platform)")
conn.execute("CREATE INDEX IF NOT EXISTS idx_alerts_platform ON viral_alerts(platform)")
conn.commit()
# Auto-seed if database is empty
cursor = conn.execute("SELECT COUNT(*) FROM artists")
count = cursor.fetchone()[0]
if count == 0:
print("[init_db] Database is empty. Seeding default artists...")
try:
from seed_artists import ARTISTS, upsert_artist
for artist in ARTISTS:
upsert_artist(conn, artist)
conn.commit()
print(f"[init_db] Successfully seeded {len(ARTISTS)} artists.")
except Exception as e:
print(f"[init_db] Failed to auto-seed: {e}")
conn.close()
print(f"Database initialized at {DB_PATH}")
def upsert_artist(conn, artist: dict):
"""Insert or update an artist."""
conn.execute(
"""INSERT INTO artists (id, name, spotify_id, youtube_channel_id, genre, image_url, is_watched)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(id) DO UPDATE SET
name=excluded.name,
spotify_id=excluded.spotify_id,
youtube_channel_id=excluded.youtube_channel_id,
genre=excluded.genre,
image_url=excluded.image_url""",
(
artist["id"],
artist["name"],
artist.get("spotify_id"),
artist.get("youtube_channel_id"),
artist.get("genre"),
artist.get("image_url"),
artist.get("is_watched", 0),
),
)
def upsert_song(conn, song: dict):
"""Insert or update a song."""
conn.execute(
"""INSERT INTO songs (id, artist_id, title, platform, platform_id, album_name, release_date, thumbnail_url)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(platform, platform_id) DO UPDATE SET
title=excluded.title,
album_name=excluded.album_name,
release_date=excluded.release_date,
thumbnail_url=excluded.thumbnail_url""",
(
song["id"],
song["artist_id"],
song["title"],
song["platform"],
song["platform_id"],
song.get("album_name"),
song.get("release_date"),
song.get("thumbnail_url"),
),
)
def generate_cycle_id(platform: str = "youtube") -> str:
"""Generate a cycle_id for the current collection run.
Format: {platform}_{YYYYMMDD_HHMM_IST}
All snapshots in the same run share this ID.
"""
ist = datetime.now(timezone.utc) + timedelta(hours=5, minutes=30)
return f"{platform}_{ist.strftime('%Y%m%d_%H%M')}"
def insert_snapshot(conn, song_id: str, play_count: int, like_count: int, comment_count: int, platform: str, cycle_id: str = None, ytmusic_play_count: int = None):
"""Record a play count snapshot with optional cycle_id and ytmusic_play_count."""
conn.execute(
"""INSERT INTO play_snapshots (song_id, play_count, like_count, comment_count, platform, cycle_id, ytmusic_play_count)
VALUES (?, ?, ?, ?, ?, ?, ?)""",
(song_id, play_count, like_count, comment_count, platform, cycle_id, ytmusic_play_count),
)
def get_previous_snapshot(conn, song_id: str, platform: str):
"""Get the most recent snapshot for a song (before the current one)."""
row = conn.execute(
"""SELECT play_count, collected_at FROM play_snapshots
WHERE song_id = ? AND platform = ?
ORDER BY collected_at DESC LIMIT 1""",
(song_id, platform),
).fetchone()
return dict(row) if row else None
def insert_viral_alert(conn, alert: dict):
"""Insert a viral alert."""
conn.execute(
"""INSERT INTO viral_alerts (song_id, previous_count, current_count, growth_factor, platform)
VALUES (?, ?, ?, ?, ?)""",
(
alert["song_id"],
alert["previous_count"],
alert["current_count"],
alert["growth_factor"],
alert["platform"],
),
)
def get_all_artists(conn, platform_filter=None):
"""Get all artists, optionally filtered by platform availability."""
if platform_filter == "youtube":
return conn.execute(
"SELECT * FROM artists WHERE youtube_channel_id IS NOT NULL ORDER BY name"
).fetchall()
elif platform_filter == "spotify":
return conn.execute(
"SELECT * FROM artists WHERE spotify_id IS NOT NULL ORDER BY name"
).fetchall()
return conn.execute("SELECT * FROM artists ORDER BY name").fetchall()
def get_watched_artists(conn):
"""Get all watched artists."""
return conn.execute(
"SELECT * FROM artists WHERE is_watched = 1 ORDER BY name"
).fetchall()
if __name__ == "__main__":
init_db()
|