drama-studio/audio-studio-api/app/database.py

"""데이터베이스 연결 설정

MongoDB (motor async) + GridFS (오디오 저장)
"""

import os
import logging
from typing import Optional

from motor.motor_asyncio import AsyncIOMotorClient, AsyncIOMotorDatabase, AsyncIOMotorGridFSBucket
from redis.asyncio import Redis

logger = logging.getLogger(__name__)


class Database:
    """데이터베이스 연결 관리"""

    def __init__(self):
        self.client: Optional[AsyncIOMotorClient] = None
        self.db: Optional[AsyncIOMotorDatabase] = None
        self.gridfs: Optional[AsyncIOMotorGridFSBucket] = None
        self.redis: Optional[Redis] = None

    async def connect(self):
        """데이터베이스 연결"""
        # MongoDB
        mongodb_url = os.getenv("MONGODB_URL", "mongodb://localhost:27017/")
        db_name = os.getenv("DB_NAME", "audio_studio")

        logger.info(f"MongoDB 연결 중: {db_name}")
        self.client = AsyncIOMotorClient(mongodb_url)
        self.db = self.client[db_name]

        # GridFS (오디오 파일 저장용)
        self.gridfs = AsyncIOMotorGridFSBucket(self.db, bucket_name="audio_files")

        # 연결 테스트
        await self.client.admin.command("ping")
        logger.info("MongoDB 연결 성공")

        # Redis
        redis_url = os.getenv("REDIS_URL", "redis://localhost:6379")
        logger.info("Redis 연결 중...")
        self.redis = Redis.from_url(redis_url, decode_responses=True)

        # 연결 테스트
        await self.redis.ping()
        logger.info("Redis 연결 성공")

        # 인덱스 생성
        await self._create_indexes()

    async def _create_indexes(self):
        """컬렉션 인덱스 생성"""
        # voices 컬렉션
        await self.db.voices.create_index("voice_id", unique=True)
        await self.db.voices.create_index("owner_id")
        await self.db.voices.create_index("type")
        await self.db.voices.create_index("language")
        await self.db.voices.create_index("is_public")

        # tts_generations 컬렉션
        await self.db.tts_generations.create_index("generation_id", unique=True)
        await self.db.tts_generations.create_index("user_id")
        await self.db.tts_generations.create_index("voice_id")
        await self.db.tts_generations.create_index("created_at")

        # sound_effects 컬렉션
        await self.db.sound_effects.create_index("source_id")
        await self.db.sound_effects.create_index("categories")
        await self.db.sound_effects.create_index("tags")

        # music_tracks 컬렉션
        await self.db.music_tracks.create_index("source")
        await self.db.music_tracks.create_index("genre")
        await self.db.music_tracks.create_index("mood")

        logger.info("인덱스 생성 완료")

    async def disconnect(self):
        """데이터베이스 연결 해제"""
        if self.client:
            self.client.close()
            logger.info("MongoDB 연결 해제")

        if self.redis:
            await self.redis.close()
            logger.info("Redis 연결 해제")

    # ========================================
    # 컬렉션 접근자
    # ========================================

    @property
    def voices(self):
        """voices 컬렉션"""
        return self.db.voices

    @property
    def tts_generations(self):
        """tts_generations 컬렉션"""
        return self.db.tts_generations

    @property
    def sound_effects(self):
        """sound_effects 컬렉션"""
        return self.db.sound_effects

    @property
    def music_tracks(self):
        """music_tracks 컬렉션"""
        return self.db.music_tracks

    @property
    def user_voice_library(self):
        """user_voice_library 컬렉션"""
        return self.db.user_voice_library

    # ========================================
    # GridFS 오디오 저장
    # ========================================

    async def save_audio(
        self,
        audio_bytes: bytes,
        filename: str,
        content_type: str = "audio/wav",
        metadata: dict = None,
    ) -> str:
        """오디오 파일을 GridFS에 저장

        Returns:
            file_id (str)
        """
        file_id = await self.gridfs.upload_from_stream(
            filename,
            audio_bytes,
            metadata={
                "content_type": content_type,
                **(metadata or {}),
            }
        )
        return str(file_id)

    async def get_audio(self, file_id: str) -> bytes:
        """GridFS에서 오디오 파일 읽기"""
        from bson import ObjectId
        from io import BytesIO

        buffer = BytesIO()
        await self.gridfs.download_to_stream(ObjectId(file_id), buffer)
        buffer.seek(0)
        return buffer.read()

    async def delete_audio(self, file_id: str):
        """GridFS에서 오디오 파일 삭제"""
        from bson import ObjectId
        await self.gridfs.delete(ObjectId(file_id))


# 싱글톤 인스턴스
db = Database()


# FastAPI 의존성
async def get_db() -> Database:
    """데이터베이스 인스턴스 반환 (의존성 주입용)"""
    return db