141 lines
4.9 KiB
Python
141 lines
4.9 KiB
Python
"""
|
|
Database migration script to create comprehensive user data cache table.
|
|
Run this script to add the cache table to your database.
|
|
"""
|
|
|
|
import sys
|
|
import os
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from sqlalchemy import create_engine, text
|
|
from sqlalchemy.orm import sessionmaker
|
|
from loguru import logger
|
|
import os
|
|
|
|
def create_cache_table():
|
|
"""Create the comprehensive user data cache table."""
|
|
try:
|
|
# Get database URL from environment or use default
|
|
database_url = os.getenv('DATABASE_URL', 'sqlite:///alwrity.db')
|
|
|
|
# Create engine
|
|
engine = create_engine(database_url)
|
|
|
|
# Create session
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
db = SessionLocal()
|
|
|
|
# SQL to create the cache table
|
|
create_table_sql = """
|
|
CREATE TABLE IF NOT EXISTS comprehensive_user_data_cache (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
user_id INTEGER NOT NULL,
|
|
strategy_id INTEGER,
|
|
data_hash VARCHAR(64) NOT NULL,
|
|
comprehensive_data JSON NOT NULL,
|
|
created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
expires_at DATETIME NOT NULL,
|
|
last_accessed DATETIME DEFAULT CURRENT_TIMESTAMP,
|
|
access_count INTEGER DEFAULT 0
|
|
);
|
|
"""
|
|
|
|
# Create indexes
|
|
create_indexes_sql = [
|
|
"CREATE INDEX IF NOT EXISTS idx_user_strategy ON comprehensive_user_data_cache(user_id, strategy_id);",
|
|
"CREATE INDEX IF NOT EXISTS idx_expires_at ON comprehensive_user_data_cache(expires_at);",
|
|
"CREATE INDEX IF NOT EXISTS idx_data_hash ON comprehensive_user_data_cache(data_hash);"
|
|
]
|
|
|
|
# Execute table creation
|
|
logger.info("Creating comprehensive_user_data_cache table...")
|
|
db.execute(text(create_table_sql))
|
|
|
|
# Execute index creation
|
|
logger.info("Creating indexes...")
|
|
for index_sql in create_indexes_sql:
|
|
db.execute(text(index_sql))
|
|
|
|
# Commit changes
|
|
db.commit()
|
|
|
|
# Verify table creation
|
|
result = db.execute(text("SELECT name FROM sqlite_master WHERE type='table' AND name='comprehensive_user_data_cache';"))
|
|
table_exists = result.fetchone()
|
|
|
|
if table_exists:
|
|
logger.info("✅ Comprehensive user data cache table created successfully!")
|
|
|
|
# Show table structure
|
|
result = db.execute(text("PRAGMA table_info(comprehensive_user_data_cache);"))
|
|
columns = result.fetchall()
|
|
|
|
logger.info("Table structure:")
|
|
for column in columns:
|
|
logger.info(f" - {column[1]} ({column[2]})")
|
|
|
|
else:
|
|
logger.error("❌ Failed to create comprehensive_user_data_cache table")
|
|
return False
|
|
|
|
db.close()
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Error creating cache table: {str(e)}")
|
|
if 'db' in locals():
|
|
db.close()
|
|
return False
|
|
|
|
def drop_cache_table():
|
|
"""Drop the comprehensive user data cache table (for testing)."""
|
|
try:
|
|
# Get database URL from environment or use default
|
|
database_url = os.getenv('DATABASE_URL', 'sqlite:///alwrity.db')
|
|
|
|
# Create engine
|
|
engine = create_engine(database_url)
|
|
|
|
# Create session
|
|
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
|
|
db = SessionLocal()
|
|
|
|
# Drop table
|
|
logger.info("Dropping comprehensive_user_data_cache table...")
|
|
db.execute(text("DROP TABLE IF EXISTS comprehensive_user_data_cache;"))
|
|
db.commit()
|
|
|
|
logger.info("✅ Comprehensive user data cache table dropped successfully!")
|
|
db.close()
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"❌ Error dropping cache table: {str(e)}")
|
|
if 'db' in locals():
|
|
db.close()
|
|
return False
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
|
|
parser = argparse.ArgumentParser(description="Manage comprehensive user data cache table")
|
|
parser.add_argument("--action", choices=["create", "drop"], default="create",
|
|
help="Action to perform (create or drop table)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
if args.action == "create":
|
|
success = create_cache_table()
|
|
if success:
|
|
logger.info("🎉 Cache table setup completed successfully!")
|
|
else:
|
|
logger.error("💥 Cache table setup failed!")
|
|
sys.exit(1)
|
|
elif args.action == "drop":
|
|
success = drop_cache_table()
|
|
if success:
|
|
logger.info("🗑️ Cache table dropped successfully!")
|
|
else:
|
|
logger.error("💥 Failed to drop cache table!")
|
|
sys.exit(1)
|