File size: 5,992 Bytes
0347fce | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 | """
Database Migration Script for Stash-Go
Handles schema updates for user data isolation
"""
from sqlalchemy import text
from sqlalchemy.orm import Session
from database import SessionLocal, engine
from models import Base
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
def check_column_exists(db: Session, table_name: str, column_name: str) -> bool:
"""Check if a column exists in a table"""
try:
query = f"PRAGMA table_info({table_name})"
result = db.execute(text(query)).fetchall()
return any(col[1] == column_name for col in result)
except Exception as e:
logger.error(f"Error checking column: {e}")
return False
def migrate_user_isolation(db: Session):
"""
Migrate database to add user_id to bills table and
ensure all sales/bills have valid user_id
"""
logger.info("Starting migration: User Data Isolation...")
try:
# Check if bills table has user_id column
if not check_column_exists(db, 'bills', 'user_id'):
logger.info("Adding user_id column to bills table...")
db.execute(text("ALTER TABLE bills ADD COLUMN user_id INTEGER"))
db.commit()
else:
logger.info("user_id already exists in bills table")
# Check if sales table has user_id column
if not check_column_exists(db, 'sales', 'user_id'):
logger.info("Adding user_id column to sales table...")
db.execute(text("ALTER TABLE sales ADD COLUMN user_id INTEGER"))
db.commit()
else:
logger.info("user_id already exists in sales table")
# Assign default user_id (1 = admin) to existing records
logger.info("Assigning default user_id to existing sales...")
db.execute(text("""
UPDATE sales
SET user_id = 1
WHERE user_id IS NULL
"""))
db.commit()
sales_updated = db.execute(text("SELECT COUNT(*) FROM sales WHERE user_id IS NOT NULL")).scalar()
logger.info(f"Sales records updated: {sales_updated}")
logger.info("Assigning default user_id to existing bills...")
db.execute(text("""
UPDATE bills
SET user_id = 1
WHERE user_id IS NULL
"""))
db.commit()
bills_updated = db.execute(text("SELECT COUNT(*) FROM bills WHERE user_id IS NOT NULL")).scalar()
logger.info(f"Bills records updated: {bills_updated}")
logger.info("✅ Migration completed successfully!")
return True
except Exception as e:
logger.error(f"❌ Migration failed: {e}")
db.rollback()
return False
finally:
db.close()
def create_tables():
"""Create all tables (idempotent)"""
logger.info("Creating tables...")
Base.metadata.create_all(bind=engine)
logger.info("✅ Tables created successfully!")
def verify_data_isolation():
"""Verify that data isolation is working"""
db = SessionLocal()
try:
logger.info("\n🔍 Verifying Data Isolation...")
# Check sales with NULL user_id
null_sales = db.execute(text("SELECT COUNT(*) FROM sales WHERE user_id IS NULL")).scalar()
if null_sales > 0:
logger.warning(f"⚠️ Found {null_sales} sales with NULL user_id")
else:
logger.info("✅ All sales have user_id assigned")
# Check bills with NULL user_id
null_bills = db.execute(text("SELECT COUNT(*) FROM bills WHERE user_id IS NULL")).scalar()
if null_bills > 0:
logger.warning(f"⚠️ Found {null_bills} bills with NULL user_id")
else:
logger.info("✅ All bills have user_id assigned")
# Show summary
total_sales = db.execute(text("SELECT COUNT(*) FROM sales")).scalar()
total_bills = db.execute(text("SELECT COUNT(*) FROM bills")).scalar()
logger.info(f"\n📊 Database Summary:")
logger.info(f" Total Sales: {total_sales}")
logger.info(f" Total Bills: {total_bills}")
return True
except Exception as e:
logger.error(f"Error during verification: {e}")
return False
finally:
db.close()
def rollback_migration():
"""Rollback migration (remove user_id columns)"""
db = SessionLocal()
try:
logger.warning("⚠️ Rolling back migration...")
db.execute(text("ALTER TABLE sales DROP COLUMN user_id"))
db.execute(text("ALTER TABLE bills DROP COLUMN user_id"))
db.commit()
logger.info("✅ Rollback completed!")
return True
except Exception as e:
logger.error(f"Error during rollback: {e}")
db.rollback()
return False
finally:
db.close()
def main():
"""Run migrations"""
import sys
logger.info("""
╔══════════════════════════════════════╗
║ Stash-Go Database Migration ║
║ User Data Isolation Setup ║
╚══════════════════════════════════════╝
""")
if len(sys.argv) > 1:
if sys.argv[1] == "create":
create_tables()
elif sys.argv[1] == "migrate":
migrate_user_isolation(SessionLocal())
elif sys.argv[1] == "verify":
verify_data_isolation()
elif sys.argv[1] == "rollback":
rollback_migration()
else:
print("Usage: python migrate.py [create|migrate|verify|rollback]")
else:
# Default: create tables then migrate
create_tables()
migrate_user_isolation(SessionLocal())
verify_data_isolation()
if __name__ == "__main__":
main()
|