tortoise-orm 动态建表

1,496 阅读2分钟

本文已参与「新人创作礼」活动,一起开启掘金创作之路。

FastAPI+Tortoise-orm+MySQL动态建表

需求:提供一个接口,根据前端传来的数据,查看数据库里是否有对应的表,没有的话则创建,然后把其他数据插入到该表中。

Django的话,很容易就能实现这个需求,网上也有现成的栗子。而这个tortoise-orm,Google了老半天,都没找到一个能打的例子。最后还是看到github上库的作者说,目前未支持该功能,才去阅读源代码,然后勉强把需求实现了。

动态创建Model类

from typing import Type
from .base import AbstModel

def make_model(table_name: str) -> Type[AbstModel]:
    klass_name = table_name.title()
    
    class Meta:
        table = table_name
        
    return type(klass_name, (AbstModel,), {'__module__': 'models', 'Meta': Meta})

使新建的类能够query和insert

from tortoise import Tortoise

CREATE_TABLE_SQL = """
create table `{}`
...
""".strip()

async def create_table(name: str) -> None:
    sql = CREATE_TABLE_SQL.format(name)
    Tortoise.get_connection('default').execute_query(sql)
    
async def init_model(model: Type[AbstModel]) -> None:
    model._meta.default_connection = 'default'
    # 主要就是参考了tortoise源码里的__init__.py文件
    Tortoise._init_relations()
    Tortoise._build_initial_querysets()

create table的SQL语句,可以这样获取:

>>> from tortoise.utils import get_schema_sql
>>> from tortoise import Tortoise
>>> from settings import TORTOISE_ORM
>>> Tortoise.init(TORTOISE_ORM)
>>> conn = Tortoise.get_connection('default')
>>> sql = get_schema_sql(conn, True)
>>> print(sql)

封装成一个对外函数

from typing import Dict
cached_models: Dict[str, Type[AbstModel]] = {}

async def get_or_create_model(table_name: str) -> Type[AbstModel]:
    if m := cached_models.get(table_name):
        return m
    model = make_model(table_name)
    await create_table(table_name)
    await init_model(model)
    cached_models[table_name] = model
    return model

接口示例

from fastapi import APIRouter
from pydantic import BaseModel
from models import AbstModel_Pydantic

router = APIRouter()

class PostData(BaseModel):
    table_name: str
    field1: str
    field2: int
       
@router.post('/', response_model=AbstModel_Pydantic)
async def create_and_insert(data: PostData):
    post = data.dict()
    model = await get_or_create_model(post.pop('table_name'))
    obj = await model.create(**post)
    return obj

其他

  • 文件结构:
% tree -a                                                       22-04-15 - 23:11:44
.
├── .env
├── app.py
├── models
│   ├── __init__.py
│   ├── account.py
│   ├── base.py
│   └── utils.py
├── poetry.lock
├── pyproject.toml
├── routers
│   ├── __init__.py
│   └── account.py
└── settings.py
  • 包版本:
% pip list                                                 :( 1 22-04-15 - 23:24:26
Package           Version
----------------- -------
aiomysql          0.1.0
aiosqlite         0.17.0
anyio             3.5.0
fastapi           0.75.1
idna              3.3
iso8601           1.0.2
pip               22.0.4
pydantic          1.9.0
PyMySQL           1.0.2
pypika-tortoise   0.1.4
python-dotenv     0.20.0
pytz              2022.1
setuptools        59.6.0
sniffio           1.2.0
starlette         0.17.1
tortoise-orm      0.19.0
typing_extensions 4.1.1
wheel             0.37.0
  • settings.py
from dotenv import load_dotenv

load_dotenv()
DEBUG = os.getenv("DEBUG", "").lower() in ("true", "y", "yes", "1")
DB_NAME = "database_name"

DB_URL = f"mysql://root:123456@127.0.0.1:3306/{DB_NAME}"
TORTOISE_ORM = {
    "connections": {"default": DB_URL},
    "apps": {"models": {"models": ["models"]}},
}
  • models/base.py
from tortoise import models, fields
from tortoise.contrib.pydantic import pydantic_model_creator

class AbstModel(models.Model):
    id = fields.IntField(pk=True)
    field1 = fields.CharField(100)
    field2 = fields.IntField()
    
    class Meta:
        abstract = True
        
# https://tortoise.github.io/examples/fastapi.html
AbstModel_Pydantic = pydantic_model_creator(AbstModel)