Base code
This commit is contained in:
Kunthawat Greethong
206
backend/scripts/create_subscription_tables.py
Normal file
206
backend/scripts/create_subscription_tables.py
Normal file
@@ -0,0 +1,206 @@
|
||||
"""
|
||||
Database Migration Script for Subscription System
|
||||
Creates all tables needed for usage-based subscription and monitoring.
|
||||
"""
|
||||
|
||||
import sys
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
# Add the backend directory to Python path
|
||||
backend_dir = Path(__file__).parent.parent
|
||||
sys.path.insert(0, str(backend_dir))
|
||||
|
||||
from sqlalchemy import create_engine, text
|
||||
from sqlalchemy.orm import sessionmaker
|
||||
from loguru import logger
|
||||
import traceback
|
||||
|
||||
# Import models
|
||||
from models.subscription_models import Base as SubscriptionBase
|
||||
from services.database import DATABASE_URL
|
||||
from services.subscription.pricing_service import PricingService
|
||||
|
||||
def create_subscription_tables():
|
||||
"""Create all subscription-related tables."""
|
||||
|
||||
try:
|
||||
# Create engine
|
||||
engine = create_engine(DATABASE_URL, echo=True)
|
||||
|
||||
# Create all tables
|
||||
logger.info("Creating subscription system tables...")
|
||||
SubscriptionBase.metadata.create_all(bind=engine)
|
||||
logger.info("✅ Subscription tables created successfully")
|
||||
|
||||
# Create session for data initialization
|
||||
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
||||
db = SessionLocal()
|
||||
|
||||
try:
|
||||
# Initialize pricing and plans
|
||||
pricing_service = PricingService(db)
|
||||
|
||||
logger.info("Initializing default API pricing...")
|
||||
pricing_service.initialize_default_pricing()
|
||||
logger.info("✅ Default API pricing initialized")
|
||||
|
||||
logger.info("Initializing default subscription plans...")
|
||||
pricing_service.initialize_default_plans()
|
||||
logger.info("✅ Default subscription plans initialized")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error initializing default data: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
logger.info("🎉 Subscription system setup completed successfully!")
|
||||
|
||||
# Display summary
|
||||
display_setup_summary(engine)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Error creating subscription tables: {e}")
|
||||
logger.error(traceback.format_exc())
|
||||
raise
|
||||
|
||||
def display_setup_summary(engine):
|
||||
"""Display a summary of the created tables and data."""
|
||||
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("SUBSCRIPTION SYSTEM SETUP SUMMARY")
|
||||
logger.info("="*60)
|
||||
|
||||
# Check tables
|
||||
tables_query = text("""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND name LIKE '%subscription%' OR name LIKE '%usage%' OR name LIKE '%pricing%'
|
||||
ORDER BY name
|
||||
""")
|
||||
|
||||
result = conn.execute(tables_query)
|
||||
tables = result.fetchall()
|
||||
|
||||
logger.info(f"\n📊 Created Tables ({len(tables)}):")
|
||||
for table in tables:
|
||||
logger.info(f" • {table[0]}")
|
||||
|
||||
# Check subscription plans
|
||||
plans_query = text("SELECT COUNT(*) FROM subscription_plans")
|
||||
result = conn.execute(plans_query)
|
||||
plan_count = result.fetchone()[0]
|
||||
logger.info(f"\n💳 Subscription Plans: {plan_count}")
|
||||
|
||||
if plan_count > 0:
|
||||
plans_detail_query = text("""
|
||||
SELECT name, tier, price_monthly, price_yearly
|
||||
FROM subscription_plans
|
||||
ORDER BY price_monthly
|
||||
""")
|
||||
result = conn.execute(plans_detail_query)
|
||||
plans = result.fetchall()
|
||||
|
||||
for plan in plans:
|
||||
name, tier, monthly, yearly = plan
|
||||
logger.info(f" • {name} ({tier}): ${monthly}/month, ${yearly}/year")
|
||||
|
||||
# Check API pricing
|
||||
pricing_query = text("SELECT COUNT(*) FROM api_provider_pricing")
|
||||
result = conn.execute(pricing_query)
|
||||
pricing_count = result.fetchone()[0]
|
||||
logger.info(f"\n💰 API Pricing Entries: {pricing_count}")
|
||||
|
||||
if pricing_count > 0:
|
||||
pricing_detail_query = text("""
|
||||
SELECT provider, model_name, cost_per_input_token, cost_per_output_token
|
||||
FROM api_provider_pricing
|
||||
WHERE cost_per_input_token > 0 OR cost_per_output_token > 0
|
||||
ORDER BY provider, model_name
|
||||
""")
|
||||
result = conn.execute(pricing_detail_query)
|
||||
pricing_entries = result.fetchall()
|
||||
|
||||
logger.info("\n LLM Pricing (per token):")
|
||||
for entry in pricing_entries:
|
||||
provider, model, input_cost, output_cost = entry
|
||||
logger.info(f" • {provider}/{model}: ${input_cost:.8f} in, ${output_cost:.8f} out")
|
||||
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("NEXT STEPS:")
|
||||
logger.info("="*60)
|
||||
logger.info("1. Update your FastAPI app to include subscription routes:")
|
||||
logger.info(" from api.subscription_api import router as subscription_router")
|
||||
logger.info(" app.include_router(subscription_router)")
|
||||
logger.info("\n2. Update database service to include subscription models:")
|
||||
logger.info(" Add SubscriptionBase.metadata.create_all(bind=engine) to init_database()")
|
||||
logger.info("\n3. Test the API endpoints:")
|
||||
logger.info(" GET /api/subscription/plans")
|
||||
logger.info(" GET /api/subscription/usage/{user_id}")
|
||||
logger.info(" GET /api/subscription/dashboard/{user_id}")
|
||||
logger.info("\n4. Configure user identification in middleware")
|
||||
logger.info(" Ensure user_id is properly extracted from requests")
|
||||
logger.info("\n5. Set up monitoring dashboard frontend integration")
|
||||
logger.info("="*60)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error displaying summary: {e}")
|
||||
|
||||
def check_existing_tables(engine):
|
||||
"""Check if subscription tables already exist."""
|
||||
|
||||
try:
|
||||
with engine.connect() as conn:
|
||||
# Check for subscription tables
|
||||
check_query = text("""
|
||||
SELECT name FROM sqlite_master
|
||||
WHERE type='table' AND (
|
||||
name = 'subscription_plans' OR
|
||||
name = 'user_subscriptions' OR
|
||||
name = 'api_usage_logs' OR
|
||||
name = 'usage_summaries'
|
||||
)
|
||||
""")
|
||||
|
||||
result = conn.execute(check_query)
|
||||
existing_tables = result.fetchall()
|
||||
|
||||
if existing_tables:
|
||||
logger.warning(f"Found existing subscription tables: {[t[0] for t in existing_tables]}")
|
||||
response = input("Tables already exist. Do you want to continue and potentially overwrite data? (y/N): ")
|
||||
if response.lower() != 'y':
|
||||
logger.info("Migration cancelled by user")
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error checking existing tables: {e}")
|
||||
return True # Proceed anyway
|
||||
|
||||
if __name__ == "__main__":
|
||||
logger.info("🚀 Starting subscription system database migration...")
|
||||
|
||||
try:
|
||||
# Create engine to check existing tables
|
||||
engine = create_engine(DATABASE_URL, echo=False)
|
||||
|
||||
# Check existing tables
|
||||
if not check_existing_tables(engine):
|
||||
sys.exit(0)
|
||||
|
||||
# Create tables and initialize data
|
||||
create_subscription_tables()
|
||||
|
||||
logger.info("✅ Migration completed successfully!")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
logger.info("Migration cancelled by user")
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
logger.error(f"❌ Migration failed: {e}")
|
||||
sys.exit(1)
|
||||
Reference in New Issue
Block a user