migrate_database.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172
  1. #!/usr/bin/env python3
  2. """
  3. Database migration script to add admin features.
  4. Adds:
  5. - is_admin column to users table
  6. - app_settings table
  7. """
  8. import sqlite3
  9. import sys
  10. def migrate_database():
  11. """Perform database migration."""
  12. db_path = "./absrecommend.db"
  13. try:
  14. conn = sqlite3.connect(db_path)
  15. cursor = conn.cursor()
  16. print("Starting database migration...")
  17. # Check if is_admin column exists
  18. cursor.execute("PRAGMA table_info(users)")
  19. columns = [col[1] for col in cursor.fetchall()]
  20. if 'is_admin' not in columns:
  21. print("Adding is_admin column to users table...")
  22. cursor.execute("ALTER TABLE users ADD COLUMN is_admin BOOLEAN DEFAULT 0")
  23. # Make the first user an admin
  24. cursor.execute("SELECT COUNT(*) FROM users")
  25. user_count = cursor.fetchone()[0]
  26. if user_count > 0:
  27. print("Making the first user an admin...")
  28. cursor.execute("UPDATE users SET is_admin = 1 ORDER BY id LIMIT 1")
  29. else:
  30. print("is_admin column already exists")
  31. # Check if app_settings table exists
  32. cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='app_settings'")
  33. if not cursor.fetchone():
  34. print("Creating app_settings table...")
  35. cursor.execute("""
  36. CREATE TABLE app_settings (
  37. id INTEGER PRIMARY KEY AUTOINCREMENT,
  38. key VARCHAR NOT NULL UNIQUE,
  39. value VARCHAR NOT NULL,
  40. updated_at DATETIME DEFAULT CURRENT_TIMESTAMP
  41. )
  42. """)
  43. # Add default settings
  44. cursor.execute("INSERT INTO app_settings (key, value) VALUES ('allow_registration', 'true')")
  45. else:
  46. print("app_settings table already exists")
  47. conn.commit()
  48. print("\n✓ Migration completed successfully!")
  49. print("\nYou can now restart the service:")
  50. print(" ./restart-service.sh")
  51. except sqlite3.Error as e:
  52. print(f"\n✗ Migration failed: {e}", file=sys.stderr)
  53. sys.exit(1)
  54. finally:
  55. if conn:
  56. conn.close()
  57. if __name__ == "__main__":
  58. migrate_database()