import pymysql

# Database connection
DB_CONFIG = {
    "host": "rm-f8ze60yirdj8786u2wo.mysql.rds.aliyuncs.com",
    "user": "root",
    "password": "csqz@20255",
    "db": "csqz-client",
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor
}

def update_pdf_schema():
    """Update genealogy_pdfs table to add required fields"""
    conn = pymysql.connect(**DB_CONFIG)
    try:
        with conn.cursor() as cursor:
            # Add version_name field
            try:
                cursor.execute("ALTER TABLE genealogy_pdfs ADD COLUMN version_name VARCHAR(255) DEFAULT ''")
                print("Added version_name column")
            except Exception as e:
                print(f"version_name column may already exist: {e}")
            
            # Add version_source field
            try:
                cursor.execute("ALTER TABLE genealogy_pdfs ADD COLUMN version_source VARCHAR(255) DEFAULT ''")
                print("Added version_source column")
            except Exception as e:
                print(f"version_source column may already exist: {e}")
            
            # Add file_provider field
            try:
                cursor.execute("ALTER TABLE genealogy_pdfs ADD COLUMN file_provider VARCHAR(255) DEFAULT ''")
                print("Added file_provider column")
            except Exception as e:
                print(f"file_provider column may already exist: {e}")
            
            # Update the existing record with sample data
            cursor.execute("UPDATE genealogy_pdfs SET version_name = '留总正式版', version_source = '家族收藏', file_provider = '留总' WHERE id = 1")
            print("Updated existing record")
            
        conn.commit()
        print("Database schema updated successfully")
        
    except Exception as e:
        print(f"Error: {e}")
        conn.rollback()
    finally:
        conn.close()

if __name__ == "__main__":
    update_pdf_schema()