Update app.py

app.py

@@ -8,81 +8,118 @@ from PIL import Image
 from io import BytesIO
 import tempfile
 import ffmpeg
-import json
+import sqlite3
 from datetime import datetime, date
 from pathlib import Path
+from threading import Lock
 
 # --- Database Setup ---
 DATA_DIR = Path("/data")
 DATA_DIR.mkdir(exist_ok=True)
-USAGE_DB_PATH = DATA_DIR / "usage_limits.json"
+DB_PATH = DATA_DIR / "usage_limits.db"
 
 DAILY_LIMIT = 50
+db_lock = Lock()
 
-def load_usage_db() -> dict:
-    """Load the usage database from disk."""
-    if USAGE_DB_PATH.exists():
-        try:
-            with open(USAGE_DB_PATH, 'r') as f:
-                return json.load(f)
-        except Exception as e:
-            print(f"Error loading usage database: {e}")
-            return {}
-    return {}
-
-def save_usage_db(db: dict):
-    """Save the usage database to disk."""
+def init_db():
+    """Initialize the SQLite database."""
+    print(f"Initializing database at: {DB_PATH}")
     try:
-        with open(USAGE_DB_PATH, 'w') as f:
-            json.dump(db, f, indent=2)
+        with sqlite3.connect(DB_PATH) as conn:
+            conn.execute('''
+                CREATE TABLE IF NOT EXISTS usage (
+                    username TEXT PRIMARY KEY,
+                    date TEXT NOT NULL,
+                    count INTEGER NOT NULL
+                )
+            ''')
+            conn.commit()
+        print("Database initialized successfully")
     except Exception as e:
-        print(f"Error saving usage database: {e}")
+        print(f"Error initializing database: {e}")
+        import traceback
+        traceback.print_exc()
 
 def check_and_update_usage(username: str) -> bool:
     """
     Check if user has reached daily limit and update usage.
     Returns True if user can generate, False if limit reached.
     """
-    db = load_usage_db()
-    today = str(date.today())
-    
-    # Initialize user record if not exists
-    if username not in db:
-        db[username] = {"date": today, "count": 0}
-    
-    user_data = db[username]
-    
-    # Reset count if it's a new day
-    if user_data["date"] != today:
-        user_data["date"] = today
-        user_data["count"] = 0
-    
-    # Check if limit reached
-    if user_data["count"] >= DAILY_LIMIT:
-        return False
-    
-    # Increment count
-    user_data["count"] += 1
-    db[username] = user_data
-    save_usage_db(db)
-    
-    return True
+    with db_lock:
+        try:
+            with sqlite3.connect(DB_PATH) as conn:
+                today = str(date.today())
+                cursor = conn.cursor()
+                
+                # Get user record
+                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
+                result = cursor.fetchone()
+                
+                if result is None:
+                    # New user - create record
+                    cursor.execute("INSERT INTO usage (username, date, count) VALUES (?, ?, ?)", 
+                                   (username, today, 1))
+                    conn.commit()
+                    print(f"New user {username}: 1/{DAILY_LIMIT}")
+                    return True
+                
+                user_date, user_count = result
+                
+                # Reset if new day
+                if user_date != today:
+                    cursor.execute("UPDATE usage SET date = ?, count = ? WHERE username = ?", 
+                                   (today, 1, username))
+                    conn.commit()
+                    print(f"User {username} reset for new day: 1/{DAILY_LIMIT}")
+                    return True
+                
+                # Check limit
+                if user_count >= DAILY_LIMIT:
+                    print(f"User {username} reached limit: {user_count}/{DAILY_LIMIT}")
+                    return False
+                
+                # Increment count
+                new_count = user_count + 1
+                cursor.execute("UPDATE usage SET count = ? WHERE username = ?", 
+                               (new_count, username))
+                conn.commit()
+                print(f"User {username} usage: {new_count}/{DAILY_LIMIT}")
+                return True
+                
+        except Exception as e:
+            print(f"Error checking usage for {username}: {e}")
+            import traceback
+            traceback.print_exc()
+            # On error, allow the request (fail open)
+            return True
 
 def get_remaining_generations(username: str) -> int:
     """Get the number of remaining generations for today."""
-    db = load_usage_db()
-    today = str(date.today())
-    
-    if username not in db:
-        return DAILY_LIMIT
-    
-    user_data = db[username]
-    
-    # Reset if it's a new day
-    if user_data["date"] != today:
-        return DAILY_LIMIT
-    
-    return max(0, DAILY_LIMIT - user_data["count"])
+    with db_lock:
+        try:
+            with sqlite3.connect(DB_PATH) as conn:
+                today = str(date.today())
+                cursor = conn.cursor()
+                
+                cursor.execute("SELECT date, count FROM usage WHERE username = ?", (username,))
+                result = cursor.fetchone()
+                
+                if result is None:
+                    return DAILY_LIMIT
+                
+                user_date, user_count = result
+                
+                # Reset if new day
+                if user_date != today:
+                    return DAILY_LIMIT
+                
+                return max(0, DAILY_LIMIT - user_count)
+        except Exception as e:
+            print(f"Error getting remaining generations for {username}: {e}")
+            return DAILY_LIMIT
+
+# Initialize database on module load
+init_db()
 
 # --- Google Gemini API Configuration ---
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY", "")