123 lines
3.7 KiB
Python
123 lines
3.7 KiB
Python
"""Database initialization and seeding functions."""
|
|
|
|
from app.database import db
|
|
from app.models.plan import Plan
|
|
|
|
|
|
def init_database():
|
|
"""Initialize database tables and seed with default data."""
|
|
# Create all tables
|
|
db.create_all()
|
|
|
|
# Seed plans if they don't exist
|
|
seed_plans()
|
|
|
|
# Migrate existing users to have plans
|
|
migrate_users_to_plans()
|
|
|
|
|
|
def seed_plans():
|
|
"""Seed the plans table with default plans if empty."""
|
|
# Check if plans already exist
|
|
if Plan.query.count() > 0:
|
|
return
|
|
|
|
# Create default plans
|
|
plans_data = [
|
|
{
|
|
"code": "free",
|
|
"name": "Free Plan",
|
|
"description": "Basic features with limited usage",
|
|
"credits": 25,
|
|
"max_credits": 75,
|
|
},
|
|
{
|
|
"code": "premium",
|
|
"name": "Premium Plan",
|
|
"description": "Enhanced features with increased usage limits",
|
|
"credits": 50,
|
|
"max_credits": 150,
|
|
},
|
|
{
|
|
"code": "pro",
|
|
"name": "Pro Plan",
|
|
"description": "Full access with unlimited usage",
|
|
"credits": 100,
|
|
"max_credits": 300,
|
|
},
|
|
]
|
|
|
|
for plan_data in plans_data:
|
|
plan = Plan(**plan_data)
|
|
db.session.add(plan)
|
|
|
|
db.session.commit()
|
|
print(f"Seeded {len(plans_data)} plans into database")
|
|
|
|
|
|
def migrate_users_to_plans():
|
|
"""Assign plans to existing users who don't have one."""
|
|
from app.models.user import User
|
|
|
|
try:
|
|
# Find users without plans
|
|
users_without_plans = User.query.filter(User.plan_id.is_(None)).all()
|
|
|
|
# Find users with plans but NULL credits (only if credits column exists)
|
|
# Note: We only migrate users with NULL credits, not 0 credits
|
|
# 0 credits means they spent them, NULL means they never got assigned
|
|
try:
|
|
users_without_credits = User.query.filter(
|
|
User.plan_id.isnot(None), User.credits.is_(None),
|
|
).all()
|
|
except Exception:
|
|
# Credits column doesn't exist yet, will be handled by create_all
|
|
users_without_credits = []
|
|
|
|
if not users_without_plans and not users_without_credits:
|
|
return
|
|
|
|
# Get default and pro plans
|
|
default_plan = Plan.get_default_plan()
|
|
pro_plan = Plan.get_pro_plan()
|
|
|
|
# Get the first user (admin) from all users ordered by ID
|
|
first_user = User.query.order_by(User.id).first()
|
|
|
|
updated_count = 0
|
|
|
|
# Assign plans to users without plans
|
|
for user in users_without_plans:
|
|
# First user gets pro plan, others get free plan
|
|
if user.id == first_user.id:
|
|
user.plan_id = pro_plan.id
|
|
# Only set credits if the column exists
|
|
try:
|
|
user.credits = pro_plan.credits
|
|
except Exception:
|
|
pass
|
|
else:
|
|
user.plan_id = default_plan.id
|
|
# Only set credits if the column exists
|
|
try:
|
|
user.credits = default_plan.credits
|
|
except Exception:
|
|
pass
|
|
updated_count += 1
|
|
|
|
# Assign credits to users with plans but no credits
|
|
for user in users_without_credits:
|
|
user.credits = user.plan.credits
|
|
updated_count += 1
|
|
|
|
if updated_count > 0:
|
|
db.session.commit()
|
|
print(
|
|
f"Updated {updated_count} existing users with plans and credits",
|
|
)
|
|
|
|
except Exception:
|
|
# If there's any error (like missing columns), just skip migration
|
|
# The database will be properly created by create_all()
|
|
pass
|