codegen/db.py

import pymysql
from config import DB

def get_conn():
    return pymysql.connect(
        **DB,
        cursorclass=pymysql.cursors.DictCursor
    )

def get_table(table_name):
    sql = """
    SELECT table_name, table_comment
    FROM information_schema.tables
    WHERE table_schema=%s AND table_name=%s
    """
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (DB["database"], table_name))
            return cur.fetchone()

def get_columns(table_name):
    sql = """
    SELECT column_name, data_type, column_comment
    FROM information_schema.columns
    WHERE table_schema=%s AND table_name=%s
    ORDER BY ordinal_position
    """
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, (DB["database"], table_name))
            return cur.fetchall()

def mysql_to_java(mysql_type):
    mapping = {
        "bigint": "Long",
        "int": "Integer",
        "varchar": "String",
        "datetime": "LocalDateTime",
        "decimal": "BigDecimal"
    }
    return mapping.get(mysql_type, "String")