import pymysql

db_host = "rm-f8ze60yirdj8786u2wo.mysql.rds.aliyuncs.com"
db_user = "root"
db_pass = "csqz@20255"
db_name = "csqz-client"

def add_super_admin_column():
    try:
        conn = pymysql.connect(
            host=db_host,
            user=db_user,
            password=db_pass,
            db=db_name,
            port=3306,
            charset='utf8mb4'
        )
        cur = conn.cursor()
        
        # Check if column exists
        cur.execute("""
            SELECT COLUMN_NAME 
            FROM INFORMATION_SCHEMA.COLUMNS 
            WHERE TABLE_NAME = 'users' AND COLUMN_NAME = 'is_super_admin'
        """)
        if cur.fetchone():
            print("Column 'is_super_admin' already exists")
        else:
            # Add is_super_admin column with default 0
            cur.execute("""
                ALTER TABLE users 
                ADD COLUMN is_super_admin TINYINT(1) DEFAULT 0 NOT NULL
            """)
            print("Added column 'is_super_admin'")
        
        # Set lihai and liuyue as super admins
        cur.execute("UPDATE users SET is_super_admin = 1 WHERE username IN ('linhai', 'liuyue')")
        affected = cur.rowcount
        print(f"Updated {affected} users as super admins")
        
        conn.commit()
        print("Operation complete.")
        return True
        
    except Exception as e:
        print(f"Error: {e}")
        return False
    finally:
        if 'conn' in locals() and conn.open:
            conn.close()

if __name__ == "__main__":
    if add_super_admin_column():
        print("SUCCESS")
    else:
        import sys
        sys.exit(1)