AI Video Generation Implementation

backend/scripts/update_image_edit_limits.py

@@ -0,0 +1,102 @@
+"""
+Script to update existing subscription plans with image_edit_calls_limit values.
+
+This script updates the SubscriptionPlan table to set image_edit_calls_limit
+for plans that were created before this column was added.
+
+Limits:
+- Free: 10 image editing calls/month
+- Basic: 30 image editing calls/month
+- Pro: 100 image editing calls/month
+- Enterprise: 0 (unlimited)
+"""
+
+import sys
+import os
+from pathlib import Path
+from datetime import datetime, timezone
+
+# 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
+from sqlalchemy.orm import sessionmaker
+from loguru import logger
+
+from models.subscription_models import SubscriptionPlan, SubscriptionTier
+from services.database import DATABASE_URL
+
+def update_image_edit_limits():
+    """Update existing subscription plans with image_edit_calls_limit values."""
+    
+    try:
+        # Create engine
+        engine = create_engine(DATABASE_URL, echo=False)
+        SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
+        db = SessionLocal()
+        
+        try:
+            # Ensure schema columns exist
+            from services.subscription.schema_utils import ensure_subscription_plan_columns
+            ensure_subscription_plan_columns(db)
+            
+            # Define limits for each tier
+            limits_by_tier = {
+                SubscriptionTier.FREE: 10,
+                SubscriptionTier.BASIC: 30,
+                SubscriptionTier.PRO: 100,
+                SubscriptionTier.ENTERPRISE: 0,  # Unlimited
+            }
+            
+            updated_count = 0
+            
+            # Update each plan
+            for tier, limit in limits_by_tier.items():
+                plans = db.query(SubscriptionPlan).filter(
+                    SubscriptionPlan.tier == tier,
+                    SubscriptionPlan.is_active == True
+                ).all()
+                
+                for plan in plans:
+                    current_limit = getattr(plan, 'image_edit_calls_limit', 0) or 0
+                    
+                    # Only update if limit is 0 (not set) or if it's different
+                    if current_limit != limit:
+                        setattr(plan, 'image_edit_calls_limit', limit)
+                        plan.updated_at = datetime.now(timezone.utc)
+                        updated_count += 1
+                        logger.info(f"Updated {plan.name} plan ({tier.value}): image_edit_calls_limit = {current_limit} -> {limit}")
+                    else:
+                        logger.debug(f"Plan {plan.name} ({tier.value}) already has image_edit_calls_limit = {limit}")
+            
+            # Commit changes
+            db.commit()
+            
+            if updated_count > 0:
+                logger.info(f"✅ Successfully updated {updated_count} subscription plan(s) with image_edit_calls_limit")
+            else:
+                logger.info("✅ All subscription plans already have correct image_edit_calls_limit values")
+            
+            return True
+            
+        except Exception as e:
+            logger.error(f"❌ Error updating image_edit_limits: {e}")
+            db.rollback()
+            raise
+        finally:
+            db.close()
+            
+    except Exception as e:
+        logger.error(f"❌ Error creating database connection: {e}")
+        raise
+
+if __name__ == "__main__":
+    logger.info("🔄 Updating subscription plans with image_edit_calls_limit...")
+    success = update_image_edit_limits()
+    if success:
+        logger.info("🎉 Image edit limits update completed successfully!")
+    else:
+        logger.error("❌ Image edit limits update failed")
+        sys.exit(1)
+