web-page-backend/music-lab/app/db.py

import os
import sqlite3
import json
from typing import Any, Dict, List, Optional

DB_PATH = "/app/data/music.db"


def _conn() -> sqlite3.Connection:
    os.makedirs(os.path.dirname(DB_PATH), exist_ok=True)
    conn = sqlite3.connect(DB_PATH)
    conn.row_factory = sqlite3.Row
    return conn


def init_db() -> None:
    with _conn() as conn:
        conn.execute("""
            CREATE TABLE IF NOT EXISTS music_tasks (
                id         TEXT PRIMARY KEY,
                status     TEXT NOT NULL DEFAULT 'queued',
                progress   INTEGER NOT NULL DEFAULT 0,
                message    TEXT NOT NULL DEFAULT '',
                audio_url  TEXT,
                error      TEXT,
                params     TEXT NOT NULL DEFAULT '{}',
                provider   TEXT NOT NULL DEFAULT 'local',
                created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
            )
        """)
        conn.execute("CREATE INDEX IF NOT EXISTS idx_tasks_created ON music_tasks(created_at DESC)")

        conn.execute("""
            CREATE TABLE IF NOT EXISTS music_library (
                id           INTEGER PRIMARY KEY AUTOINCREMENT,
                title        TEXT NOT NULL DEFAULT '',
                genre        TEXT NOT NULL DEFAULT '',
                moods        TEXT NOT NULL DEFAULT '[]',
                instruments  TEXT NOT NULL DEFAULT '[]',
                duration_sec INTEGER,
                bpm          INTEGER,
                key          TEXT NOT NULL DEFAULT '',
                scale        TEXT NOT NULL DEFAULT '',
                prompt       TEXT NOT NULL DEFAULT '',
                audio_url    TEXT NOT NULL DEFAULT '',
                file_path    TEXT NOT NULL DEFAULT '',
                task_id      TEXT,
                tags         TEXT NOT NULL DEFAULT '[]',
                provider     TEXT NOT NULL DEFAULT 'local',
                lyrics       TEXT NOT NULL DEFAULT '',
                image_url    TEXT NOT NULL DEFAULT '',
                suno_id      TEXT NOT NULL DEFAULT '',
                created_at   TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
            )
        """)
        conn.execute("CREATE INDEX IF NOT EXISTS idx_library_created ON music_library(created_at DESC)")

        conn.execute("""
            CREATE TABLE IF NOT EXISTS saved_lyrics (
                id         INTEGER PRIMARY KEY AUTOINCREMENT,
                title      TEXT NOT NULL DEFAULT '',
                text       TEXT NOT NULL DEFAULT '',
                prompt     TEXT NOT NULL DEFAULT '',
                created_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now')),
                updated_at TEXT NOT NULL DEFAULT (strftime('%Y-%m-%dT%H:%M:%fZ','now'))
            )
        """)
        conn.execute("CREATE INDEX IF NOT EXISTS idx_lyrics_created ON saved_lyrics(created_at DESC)")

        # 기존 테이블 마이그레이션 (컬럼 없으면 추가)
        for col, default in [
            ("provider", "'local'"), ("lyrics", "''"),
            ("image_url", "''"), ("suno_id", "''"),
        ]:
            try:
                conn.execute(f"ALTER TABLE music_library ADD COLUMN {col} TEXT NOT NULL DEFAULT {default}")
            except sqlite3.OperationalError:
                pass  # 이미 존재
        try:
            conn.execute("ALTER TABLE music_tasks ADD COLUMN provider TEXT NOT NULL DEFAULT 'local'")
        except sqlite3.OperationalError:
            pass


# ── music_tasks CRUD ──────────────────────────────────────────────────────────

def _task_row_to_dict(r) -> Dict[str, Any]:
    return {
        "task_id":   r["id"],
        "status":    r["status"],
        "progress":  r["progress"],
        "message":   r["message"],
        "audio_url": r["audio_url"],
        "error":     r["error"],
        "params":    json.loads(r["params"]),
        "provider":  r["provider"] if "provider" in r.keys() else "local",
        "created_at": r["created_at"],
        "updated_at": r["updated_at"],
    }


def create_task(task_id: str, params: Dict[str, Any], provider: str = "local") -> Dict[str, Any]:
    with _conn() as conn:
        conn.execute(
            "INSERT INTO music_tasks (id, params, provider) VALUES (?, ?, ?)",
            (task_id, json.dumps(params), provider),
        )
        row = conn.execute("SELECT * FROM music_tasks WHERE id = ?", (task_id,)).fetchone()
    return _task_row_to_dict(row)


def update_task(
    task_id: str,
    status: str,
    progress: int,
    message: str,
    audio_url: str = None,
    error: str = None,
) -> None:
    with _conn() as conn:
        conn.execute(
            """
            UPDATE music_tasks
            SET status = ?, progress = ?, message = ?, audio_url = ?, error = ?,
                updated_at = strftime('%Y-%m-%dT%H:%M:%fZ','now')
            WHERE id = ?
            """,
            (status, progress, message, audio_url, error, task_id),
        )


def get_task(task_id: str) -> Optional[Dict[str, Any]]:
    with _conn() as conn:
        row = conn.execute("SELECT * FROM music_tasks WHERE id = ?", (task_id,)).fetchone()
    return _task_row_to_dict(row) if row else None


# ── music_library CRUD ────────────────────────────────────────────────────────

def _track_row_to_dict(r) -> Dict[str, Any]:
    keys = r.keys()
    return {
        "id":           r["id"],
        "title":        r["title"],
        "genre":        r["genre"],
        "moods":        json.loads(r["moods"]) if r["moods"] else [],
        "instruments":  json.loads(r["instruments"]) if r["instruments"] else [],
        "duration_sec": r["duration_sec"],
        "bpm":          r["bpm"],
        "key":          r["key"],
        "scale":        r["scale"],
        "prompt":       r["prompt"],
        "audio_url":    r["audio_url"],
        "file_path":    r["file_path"],
        "task_id":      r["task_id"],
        "tags":         json.loads(r["tags"]) if r["tags"] else [],
        "provider":     r["provider"] if "provider" in keys else "local",
        "lyrics":       r["lyrics"] if "lyrics" in keys else "",
        "image_url":    r["image_url"] if "image_url" in keys else "",
        "suno_id":      r["suno_id"] if "suno_id" in keys else "",
        "created_at":   r["created_at"],
    }


def get_all_tracks() -> List[Dict[str, Any]]:
    with _conn() as conn:
        rows = conn.execute("SELECT * FROM music_library ORDER BY created_at DESC").fetchall()
    return [_track_row_to_dict(r) for r in rows]


def add_track(data: Dict[str, Any]) -> Dict[str, Any]:
    with _conn() as conn:
        conn.execute(
            """
            INSERT INTO music_library
                (title, genre, moods, instruments, duration_sec, bpm, key, scale,
                 prompt, audio_url, file_path, task_id, tags,
                 provider, lyrics, image_url, suno_id)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                data.get("title", ""),
                data.get("genre", ""),
                json.dumps(data.get("moods", [])),
                json.dumps(data.get("instruments", [])),
                data.get("duration_sec"),
                data.get("bpm"),
                data.get("key", ""),
                data.get("scale", ""),
                data.get("prompt", ""),
                data.get("audio_url", ""),
                data.get("file_path", ""),
                data.get("task_id"),
                json.dumps(data.get("tags", [])),
                data.get("provider", "local"),
                data.get("lyrics", ""),
                data.get("image_url", ""),
                data.get("suno_id", ""),
            ),
        )
        row = conn.execute("SELECT * FROM music_library WHERE rowid = last_insert_rowid()").fetchone()
    return _track_row_to_dict(row)


def delete_track(track_id: int) -> bool:
    with _conn() as conn:
        # 파일 경로 먼저 조회 (삭제 후 파일도 지울 수 있도록)
        row = conn.execute("SELECT file_path FROM music_library WHERE id = ?", (track_id,)).fetchone()
        if not row:
            return False
        conn.execute("DELETE FROM music_library WHERE id = ?", (track_id,))
    return True


def get_track_by_task_id(task_id: str) -> Optional[Dict[str, Any]]:
    with _conn() as conn:
        row = conn.execute("SELECT * FROM music_library WHERE task_id = ?", (task_id,)).fetchone()
    return _track_row_to_dict(row) if row else None


def update_track_duration(track_id: int, duration_sec: int) -> None:
    with _conn() as conn:
        conn.execute(
            "UPDATE music_library SET duration_sec = ? WHERE id = ? AND duration_sec IS NULL",
            (duration_sec, track_id),
        )


def get_track_file_path(track_id: int) -> Optional[str]:
    with _conn() as conn:
        row = conn.execute("SELECT file_path FROM music_library WHERE id = ?", (track_id,)).fetchone()
    return row["file_path"] if row else None


# ── saved_lyrics CRUD ────────────────────────────────────────────────────────

def _lyrics_row_to_dict(r) -> Dict[str, Any]:
    return {
        "id":         r["id"],
        "title":      r["title"],
        "text":       r["text"],
        "prompt":     r["prompt"],
        "created_at": r["created_at"],
        "updated_at": r["updated_at"],
    }


def get_all_lyrics() -> List[Dict[str, Any]]:
    with _conn() as conn:
        rows = conn.execute("SELECT * FROM saved_lyrics ORDER BY created_at DESC").fetchall()
    return [_lyrics_row_to_dict(r) for r in rows]


def add_lyrics(data: Dict[str, Any]) -> Dict[str, Any]:
    with _conn() as conn:
        conn.execute(
            "INSERT INTO saved_lyrics (title, text, prompt) VALUES (?, ?, ?)",
            (data.get("title", ""), data.get("text", ""), data.get("prompt", "")),
        )
        row = conn.execute("SELECT * FROM saved_lyrics WHERE rowid = last_insert_rowid()").fetchone()
    return _lyrics_row_to_dict(row)


def update_lyrics(lyrics_id: int, data: Dict[str, Any]) -> Optional[Dict[str, Any]]:
    with _conn() as conn:
        fields = []
        values = []
        for k in ("title", "text", "prompt"):
            if k in data:
                fields.append(f"{k} = ?")
                values.append(data[k])
        if not fields:
            return None
        fields.append("updated_at = strftime('%Y-%m-%dT%H:%M:%fZ','now')")
        values.append(lyrics_id)
        conn.execute(f"UPDATE saved_lyrics SET {', '.join(fields)} WHERE id = ?", values)
        row = conn.execute("SELECT * FROM saved_lyrics WHERE id = ?", (lyrics_id,)).fetchone()
    return _lyrics_row_to_dict(row) if row else None


def delete_lyrics(lyrics_id: int) -> bool:
    with _conn() as conn:
        row = conn.execute("SELECT id FROM saved_lyrics WHERE id = ?", (lyrics_id,)).fetchone()
        if not row:
            return False
        conn.execute("DELETE FROM saved_lyrics WHERE id = ?", (lyrics_id,))
    return True