背景:sqlacodegen提供了方法,可以通过链接数据库以后根据表自动创建model类。虽然schemas类根据不同场景和数据库表字段是有区别的,但是我在写脚本的时候其实用的schemas类字段大部分都是根据数据库表进行删改的,所以写了一个方法,可以根据model类映射出来一个schemas类,然后再根据需求手动删改字段,可以节省点时间
from sqlalchemy.ext.declarative import declarative_base
from datetime import datetime
from decimal import Decimal
from TestDB import DB
Base = declarative_base()
def get_column_python_type(column):
#sqlalchemy的数据库字段类型转python字段类型
python_type = column.type.python_type
return python_type
def add_import_if_needed(import_statement):
#增加导入包的文本(如datetime类型,需要增加from datetime import datetime)
import_texts = ''
for import_text in import_statement:
import_texts = import_texts + '\n' + import_text
return import_texts
def generate_schema_text(model_class):
schema_text = f"class {model_class.__name__}DTO(BaseModel):\n"
import_statements = set()
#遍历model的字段和类型
for column_name, column in model_class.__mapper__.c.items():
#剔除数据库中的id列映射
if column_name == 'id':
continue
column_type = get_column_python_type(column)
#根据model字段和类型生成schema类字段和类型
#数据库表中可空的字段映射schema类增加Optional关键字
if column.nullable==True:
schema_text += f" {column_name}: Optional[{column_type.__name__}]\n"
else:
schema_text += f" {column_name}: {column_type.__name__}\n"
if column_type in {datetime, Decimal}:
#model类型中的datetime和Decimal类型转schema类型的时候需要增加导入语句。
#这里根据需要可以再添加其他的类型
import_statements.add(f"from {column_type.__module__} import {column_type.__name__}")
schema_text = add_import_if_needed(import_statements)+'\n\n'+schema_text
return schema_text
def write_schema_to_file(model_class, file_path):
text = 'from pydantic import BaseModel\n'+'from typing import Optional'
schema_text = text + generate_schema_text(model_class)
model_name = model_class.__name__
file_path = file_path + model_name + 'DTO.py'
with open(file_path, 'w+') as f:
f.write(schema_text)
# 测试
write_schema_to_file(DB.Achieve,'/Users/username/PycharmProjects/my_script/Admin/schema/')
model文件内容是
class Achieve(Base):
__tablename__ = 'achievement'
user_id = Column(BIGINT(20),nullable==False)
pay_time = Column(DateTime)
pay_amount = Column(DECIMAL(10, 2))
运行结果是会产生一个python文件,内容如下。根据需要可以再这个类里面删改字段了
from pydantic import BaseModel
from typing import Optional
from decimal import Decimal
from datetime import datetime
class AchieveDTO(BaseModel):
user_id: int
pay_time: Optional[datetime]
pay_amount: Optional[Decimal]