使用Python从数据库读取表生成JAVA实体类

import pymysql


#连接数据库
def generate_java_entity(table_name, class_name):
    connection = pymysql.connect(
        host='****',# 数据库host
        user='*****',# 数据库用户名
        password='****',# 数据库密码
        database='*****' # 数据库名称
    )

    try:
        with connection.cursor() as cursor:
            cursor.execute(f"DESCRIBE {table_name}")
            columns = cursor.fetchall()

        java_code = f"import lombok.Data;\n"
        java_code += f"import javax.persistence.*;\n"
        if any('datetime' in column[1] for column in columns):
            java_code += f"import java.util.Date;\n"
        if any('decimal' in column[1] for column in columns):
            java_code += f"import java.math.BigDecimal;\n"
        java_code += f"\n@Data\n@Entity\n@Table(name = \"{table_name}\")\n"
        java_code += f"public class {class_name} {{\n"

        for column in columns:
            column_name = column[0]
            column_type = column[1]
            java_type = get_java_type(column_type)

            if column[3] == 'PRI':
                java_code += f"    @Id\n    @GeneratedValue(strategy = GenerationType.IDENTITY)\n"

            java_code += f"    @Column(name = \"{column_name}\")\n"
            java_code += f"    private {java_type} {to_camel_case(column_name)};\n\n"

        java_code += "}"

        with open(f"{class_name}.java", "w") as file:
            file.write(java_code)

        print(f"{class_name}.java has been generated.")

    finally:
        connection.close()


def get_java_type(sql_type):
    if 'int' in sql_type:
        return 'Integer'
    elif 'varchar' in sql_type or 'text' in sql_type:
        return 'String'
    elif 'decimal' in sql_type:
        return 'BigDecimal'
    elif 'date' in sql_type or 'timestamp' in sql_type:
        return 'Date'
    elif 'datetime' in sql_type:
        return 'Date'
    elif 'bigint' in sql_type:
        return 'Long'
    elif 'text' in sql_type:
        return 'String'
    return 'String'


def to_camel_case(snake_str):
    components = snake_str.split('_')
    return components[0] + ''.join(x.title() for x in components[1:])

#这里输入要读取的表名称和要转化成的java的实体类名称
generate_java_entity('test_table', 'TestTable')