使用 datamodel-code-generator 从 MySQL 生成 Python 模型

简介

datamodel-code-generator 是一个强大的工具，可以从多种数据源（包括 MySQL）自动生成 Python 数据模型。本文将详细介绍如何使用它从 MySQL 数据库生成 Pydantic 模型。

安装必要组件

首先需要安装相关的 Python 包：

# 安装主要工具
pip install datamodel-code-generator

# 安装 MySQL 支持
pip install sqlalchemy
pip install pymysql

使用方法

1. 从数据库直接生成

最基本的用法是直接从数据库连接生成：

datamodel-codegen \
    --db-url="mysql+pymysql://username:password@localhost:3306/database_name" \
    --output models.py

2. 从 SQL 文件生成

如果您有数据库的 SQL 文件，也可以直接从 SQL 文件生成：

datamodel-codegen \
    --input your_schema.sql \
    --input-file-type sql \
    --output models.py

常用参数说明

基本参数

• --output: 输出文件路径
• --input: 输入文件（使用 SQL 文件时）
• --input-file-type: 输入文件类型
• --db-url: 数据库连接字符串

高级参数

• --target-python-version: 指定目标 Python 版本

  datamodel-codegen --db-url="..." --output models.py --target-python-version 3.9
  

• --use-schema: 使用数据库 schema

  datamodel-codegen --db-url="..." --output models.py --use-schema
  

• --snake-case-field: 将字段名转换为蛇形命名

  datamodel-codegen --db-url="..." --output models.py --snake-case-field
  

生成代码示例

假设有以下 MySQL 表结构：

CREATE TABLE users (
    id INT PRIMARY KEY AUTO_INCREMENT,
    username VARCHAR(50) NOT NULL,
    email VARCHAR(100),
    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
);

生成的 Python 代码将类似于：

from datetime import datetime
from typing import Optional
from pydantic import BaseModel

class User(BaseModel):
    id: int
    username: str
    email: Optional[str] = None
    created_at: datetime

    class Config:
        orm_mode = True

类型映射

MySQL 到 Python 的主要类型映射关系：

MySQL 类型	Python 类型
INT	int
VARCHAR	str
TEXT	str
DATETIME	datetime
BOOLEAN	bool
DECIMAL	Decimal
JSON	dict

最佳实践

1. 错误处理

   • 始终检查生成的代码是否符合预期
   • 注意处理可能的类型转换问题

2. 自定义模板
   如果需要自定义生成的代码格式：

   datamodel-codegen --db-url="..." --output models.py --custom-template path/to/template.j2
   

3. 字段验证
   可以添加额外的验证规则：

   from pydantic import BaseModel, EmailStr
   
   class User(BaseModel):
       email: Optional[EmailStr] = None
   

常见问题解决

1. 连接问题

   • 确保数据库用户有足够权限
   • 检查网络连接和防火墙设置

2. 类型转换问题

   • 某些特殊类型可能需要手动调整
   • 考虑使用自定义类型映射

是的，默认情况下 datamodel-codegen 会将所有表的模型都写入同一个文件。如果您想要将不同的表分别生成到不同的文件，有以下几种解决方案：

1. 使用 --target-python-version 3.7 --class-name "*.py" 参数

datamodel-codegen \
    --db-url="mysql+pymysql://username:password@localhost:3306/database_name" \
    --output ./models \
    --target-python-version 3.7 \
    --class-name "*.py"

这样会在 models 目录下为每个表生成单独的文件。

2. 使用 --custom-file-mapping 参数

创建一个 JSON 文件来定义映射关系（例如 mapping.json）：

{
    "user.py": {
        "tables": ["users", "user_profiles"]
    },
    "product.py": {
        "tables": ["products", "categories"]
    }
}

然后使用这个映射文件：

datamodel-codegen \
    --db-url="mysql+pymysql://username:password@localhost:3306/database_name" \
    --output ./models \
    --custom-file-mapping mapping.json

3. 手动分割生成的文件

如果表不多，也可以先生成到一个文件，然后手动分割到不同的文件中。

推荐的项目结构

models/
├── __init__.py
├── user.py
├── product.py
├── order.py
└── common.py

在 __init__.py 中导入所有模型：

from .user import User, UserProfile
from .product import Product, Category
from .order import Order

这样的结构更清晰，也更容易维护。每个文件包含相关的模型：

from datetime import datetime
from typing import Optional
from pydantic import BaseModel

class User(BaseModel):
    id: int
    username: str
    # ...

class UserProfile(BaseModel):
    user_id: int
    # ...

from typing import Optional
from pydantic import BaseModel

class Product(BaseModel):
    id: int
    name: str
    # ...

class Category(BaseModel):
    id: int
    name: str
    # ...

注意事项

1. 使用分割文件时要注意处理模型之间的依赖关系
2. 确保 __init__.py 中正确导入了所有需要的模型
3. 如果使用自定义映射，要确保所有表都被包含在内
4. 考虑按业务领域或功能模块来组织文件结构

这样的组织方式会让代码更容易维护和理解。

总结

datamodel-code-generator 是一个非常实用的工具，可以大大提高开发效率。通过合理使用其提供的参数和选项，可以生成更符合项目需求的数据模型代码。

参考资料

• datamodel-code-generator 官方文档
• Pydantic 文档
• SQLAlchemy 文档