sqlalchemy根据数据库的model类生成schemas类_sqlalchemy model 设置schema

背景：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/')
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52

model文件内容是

class Achieve(Base):
    __tablename__ = 'achievement'

    user_id = Column(BIGINT(20),nullable==False)
    pay_time = Column(DateTime)
    pay_amount = Column(DECIMAL(10, 2))
1
2
3
4
5
6

运行结果是会产生一个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]
1
2
3
4
5
6
7
8
9