132 lines
4.3 KiB
Python
132 lines
4.3 KiB
Python
#!/usr/bin/env python3
|
||
"""
|
||
Migration script to add missing wavespeed columns to usage_summaries table
|
||
"""
|
||
|
||
import sqlite3
|
||
import sys
|
||
import os
|
||
from pathlib import Path
|
||
|
||
def get_db_path():
|
||
"""Find the database file"""
|
||
# Look for common database locations
|
||
possible_paths = [
|
||
"backend/database.db",
|
||
"backend/data/database.db",
|
||
"database.db",
|
||
"data/database.db"
|
||
]
|
||
|
||
for path in possible_paths:
|
||
if os.path.exists(path):
|
||
return path
|
||
|
||
# If not found, check if there's a .db file in backend directory
|
||
backend_dir = Path("backend")
|
||
if backend_dir.exists():
|
||
for db_file in backend_dir.glob("*.db"):
|
||
return str(db_file)
|
||
|
||
return None
|
||
|
||
def migrate_usage_summaries():
|
||
"""Add missing wavespeed columns to usage_summaries table"""
|
||
|
||
db_path = get_db_path()
|
||
if not db_path:
|
||
print("❌ Database file not found!")
|
||
print("Looked in:")
|
||
for path in ["backend/database.db", "backend/data/database.db", "database.db", "data/database.db"]:
|
||
print(f" - {path}")
|
||
return False
|
||
|
||
print(f"📁 Using database: {db_path}")
|
||
|
||
try:
|
||
conn = sqlite3.connect(db_path)
|
||
cursor = conn.cursor()
|
||
|
||
# Check if table exists
|
||
cursor.execute("""
|
||
SELECT name FROM sqlite_master
|
||
WHERE type='table' AND name='usage_summaries'
|
||
""")
|
||
|
||
if not cursor.fetchone():
|
||
print("❌ Table 'usage_summaries' does not exist!")
|
||
return False
|
||
|
||
# Get current columns
|
||
cursor.execute("PRAGMA table_info(usage_summaries)")
|
||
columns = [row[1] for row in cursor.fetchall()]
|
||
print(f"📋 Current columns: {columns}")
|
||
|
||
# Columns to add
|
||
columns_to_add = [
|
||
("wavespeed_calls", "INTEGER DEFAULT 0"),
|
||
("tavily_calls", "INTEGER DEFAULT 0"),
|
||
("serper_calls", "INTEGER DEFAULT 0"),
|
||
("metaphor_calls", "INTEGER DEFAULT 0"),
|
||
("firecrawl_calls", "INTEGER DEFAULT 0"),
|
||
("stability_calls", "INTEGER DEFAULT 0"),
|
||
("exa_calls", "INTEGER DEFAULT 0"),
|
||
("video_calls", "INTEGER DEFAULT 0"),
|
||
("image_edit_calls", "INTEGER DEFAULT 0"),
|
||
("audio_calls", "INTEGER DEFAULT 0"),
|
||
("wavespeed_tokens", "INTEGER DEFAULT 0"),
|
||
("wavespeed_cost", "FLOAT DEFAULT 0.0"),
|
||
("tavily_cost", "FLOAT DEFAULT 0.0"),
|
||
("serper_cost", "FLOAT DEFAULT 0.0"),
|
||
("metaphor_cost", "FLOAT DEFAULT 0.0"),
|
||
("firecrawl_cost", "FLOAT DEFAULT 0.0"),
|
||
("stability_cost", "FLOAT DEFAULT 0.0"),
|
||
("exa_cost", "FLOAT DEFAULT 0.0"),
|
||
("video_cost", "FLOAT DEFAULT 0.0"),
|
||
("image_edit_cost", "FLOAT DEFAULT 0.0"),
|
||
("audio_cost", "FLOAT DEFAULT 0.0")
|
||
]
|
||
|
||
# Add missing columns
|
||
added_columns = []
|
||
for column_name, column_def in columns_to_add:
|
||
if column_name not in columns:
|
||
print(f"➕ Adding column: {column_name}")
|
||
cursor.execute(f"ALTER TABLE usage_summaries ADD COLUMN {column_name} {column_def}")
|
||
added_columns.append(column_name)
|
||
else:
|
||
print(f"✅ Column already exists: {column_name}")
|
||
|
||
if added_columns:
|
||
conn.commit()
|
||
print(f"🎉 Successfully added {len(added_columns)} columns:")
|
||
for col in added_columns:
|
||
print(f" - {col}")
|
||
else:
|
||
print("✅ All columns already exist!")
|
||
|
||
# Verify the changes
|
||
cursor.execute("PRAGMA table_info(usage_summaries)")
|
||
new_columns = [row[1] for row in cursor.fetchall()]
|
||
print(f"📋 Updated columns: {new_columns}")
|
||
|
||
conn.close()
|
||
return True
|
||
|
||
except sqlite3.Error as e:
|
||
print(f"❌ Database error: {e}")
|
||
return False
|
||
except Exception as e:
|
||
print(f"❌ Unexpected error: {e}")
|
||
return False
|
||
|
||
if __name__ == "__main__":
|
||
print("🚀 Starting usage_summaries migration...")
|
||
|
||
if migrate_usage_summaries():
|
||
print("✅ Migration completed successfully!")
|
||
sys.exit(0)
|
||
else:
|
||
print("❌ Migration failed!")
|
||
sys.exit(1)
|