Add database seeding functionality and enhance model relationships
- Implement initial data seeding for plans in the database. - Create a new `seed_all_data` function to manage seeding process. - Update `Sound` and `User` models to include relationships for `SoundPlayed` and `Stream`.
This commit is contained in:
@@ -6,6 +6,7 @@ from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.logging import get_logger
|
||||
from app.core.seeds import seed_all_data
|
||||
from app.models import ( # noqa: F401
|
||||
plan,
|
||||
playlist,
|
||||
@@ -45,6 +46,19 @@ async def init_db() -> None:
|
||||
async with engine.begin() as conn:
|
||||
await conn.run_sync(SQLModel.metadata.create_all)
|
||||
logger.info("Database tables created successfully")
|
||||
|
||||
# Seed initial data
|
||||
await seed_initial_data()
|
||||
|
||||
except Exception:
|
||||
logger.exception("Failed to initialize database")
|
||||
raise
|
||||
|
||||
|
||||
async def seed_initial_data() -> None:
|
||||
"""Seed initial data into the database."""
|
||||
logger = get_logger(__name__)
|
||||
logger.info("Starting initial data seeding")
|
||||
|
||||
async with AsyncSession(engine) as session:
|
||||
await seed_all_data(session)
|
||||
|
||||
65
app/core/seeds.py
Normal file
65
app/core/seeds.py
Normal file
@@ -0,0 +1,65 @@
|
||||
"""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
|
||||
|
||||
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_all_data(session: AsyncSession) -> None:
|
||||
"""Seed all initial data."""
|
||||
logger.info("Starting data seeding")
|
||||
|
||||
try:
|
||||
await seed_plans(session)
|
||||
logger.info("Data seeding completed successfully")
|
||||
except Exception:
|
||||
logger.exception("Failed to seed data")
|
||||
await session.rollback()
|
||||
raise
|
||||
@@ -6,6 +6,7 @@ from app.models.base import BaseModel
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from app.models.playlist_sound import PlaylistSound
|
||||
from app.models.sound_played import SoundPlayed
|
||||
from app.models.stream import Stream
|
||||
|
||||
|
||||
@@ -31,3 +32,4 @@ class Sound(BaseModel, table=True):
|
||||
# relationships
|
||||
playlist_sounds: list["PlaylistSound"] = Relationship(back_populates="sound")
|
||||
streams: list["Stream"] = Relationship(back_populates="sound")
|
||||
play_history: list["SoundPlayed"] = Relationship(back_populates="sound")
|
||||
|
||||
@@ -8,6 +8,8 @@ from app.models.base import BaseModel
|
||||
if TYPE_CHECKING:
|
||||
from app.models.plan import Plan
|
||||
from app.models.playlist import Playlist
|
||||
from app.models.sound_played import SoundPlayed
|
||||
from app.models.stream import Stream
|
||||
from app.models.user_oauth import UserOauth
|
||||
|
||||
|
||||
@@ -29,3 +31,5 @@ class User(BaseModel, table=True):
|
||||
oauths: list["UserOauth"] = Relationship(back_populates="user")
|
||||
plan: "Plan" = Relationship(back_populates="users")
|
||||
playlists: list["Playlist"] = Relationship(back_populates="user")
|
||||
sounds_played: list["SoundPlayed"] = Relationship(back_populates="user")
|
||||
streams: list["Stream"] = Relationship(back_populates="user")
|
||||
|
||||
Reference in New Issue
Block a user