92 lines
2.5 KiB
Python
92 lines
2.5 KiB
Python
"""Database seeding functionality."""
|
|
|
|
from sqlmodel import select
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
|
|
from app.core.logging import get_logger
|
|
from app.models.plan import Plan
|
|
from app.models.playlist import Playlist
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
async def seed_plans(session: AsyncSession) -> None:
|
|
"""Seed initial plan data."""
|
|
logger.info("Seeding plan data")
|
|
|
|
# Check if plans already exist
|
|
existing_plans = await session.exec(select(Plan))
|
|
if existing_plans.first():
|
|
logger.info("Plans already exist, skipping seeding")
|
|
return
|
|
|
|
# Define initial plans
|
|
initial_plans = [
|
|
Plan(
|
|
code="free",
|
|
name="Free Plan",
|
|
description="Basic free plan with limited features",
|
|
credits=25,
|
|
max_credits=75,
|
|
),
|
|
Plan(
|
|
code="premium",
|
|
name="Premium Plan",
|
|
description="Premium plan with more features",
|
|
credits=50,
|
|
max_credits=150,
|
|
),
|
|
Plan(
|
|
code="pro",
|
|
name="Pro Plan",
|
|
description="Pro plan with unlimited features",
|
|
credits=100,
|
|
max_credits=300,
|
|
),
|
|
]
|
|
|
|
# Add plans to session
|
|
for plan in initial_plans:
|
|
session.add(plan)
|
|
|
|
await session.commit()
|
|
logger.info("Successfully seeded %d plans", len(initial_plans))
|
|
|
|
|
|
async def seed_main_playlist(session: AsyncSession) -> None:
|
|
"""Seed the main playlist with initial data."""
|
|
logger.info("Seeding main playlist data")
|
|
|
|
# Check if the main playlist already exists
|
|
existing_playlist = await session.exec(select(Playlist).where(Playlist.is_main))
|
|
if existing_playlist.first():
|
|
logger.info("Main playlist already exists, skipping seeding")
|
|
return
|
|
|
|
# Create the main playlist
|
|
main_playlist = Playlist(
|
|
name="All",
|
|
description="The default main playlist with all the tracks",
|
|
is_main=True,
|
|
is_deletable=False,
|
|
is_current=True,
|
|
)
|
|
session.add(main_playlist)
|
|
|
|
await session.commit()
|
|
logger.info("Successfully seeded main playlist")
|
|
|
|
|
|
async def seed_all_data(session: AsyncSession) -> None:
|
|
"""Seed all initial data."""
|
|
logger.info("Starting data seeding")
|
|
|
|
try:
|
|
await seed_plans(session)
|
|
await seed_main_playlist(session)
|
|
logger.info("Data seeding completed successfully")
|
|
except Exception:
|
|
logger.exception("Failed to seed data")
|
|
await session.rollback()
|
|
raise
|