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 '{}',
                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 '[]',
                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)")


# ── 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"]),
        "created_at": r["created_at"],
        "updated_at": r["updated_at"],
    }


def create_task(task_id: str, params: Dict[str, Any]) -> Dict[str, Any]:
    with _conn() as conn:
        conn.execute(
            "INSERT INTO music_tasks (id, params) VALUES (?, ?)",
            (task_id, json.dumps(params)),
        )
        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]:
    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 [],
        "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)
            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", [])),
            ),
        )
        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 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