SqlModel 教程之外的实用操作

 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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70

from contextlib import asynccontextmanager,contextmanager
from typing import AsyncIterator, Iterator
from sqlalchemy.ext.asyncio import create_async_engine, async_sessionmaker
from sqlalchemy import AsyncAdaptedQueuePool, create_engine
from sqlalchemy.orm import sessionmaker
from sqlmodel import Session
from sqlmodel.ext.asyncio.session import AsyncSession


async_engine = create_async_engine(
    async_url, # 异步dsn  postgresql+asyncpg://
    echo=app.settings.ECHO,
    future=True,
    pool_size=2,
    max_overflow=30,
    pool_recycle=3600,
    pool_pre_ping=True,
    poolclass=AsyncAdaptedQueuePool,
)

AsyncSessionLocal = async_sessionmaker(
    bind=async_engine,
    class_=AsyncSession,
    expire_on_commit=False,  # 避免commit后属性过期
)


@asynccontextmanager
async def get_async_session() -> AsyncIterator[AsyncSession]:
    """ """

    async with AsyncSessionLocal() as session:
        yield session
        await session.commit()

sync_engine = create_engine(
    sync_url,  # postgresql+psycopg2://
    echo=app.settings.ECHO,
    pool_size=2,
    max_overflow=30,
    pool_recycle=3600,
    pool_pre_ping=True,
)

SyncSessionLocal = sessionmaker(
    bind=sync_engine,
    class_=Session,
    expire_on_commit=False,  # 避免commit后属性过期
)

@contextmanager
def get_sync_session()->Iterator[Session]:
    with SyncSessionLocal() as session:
        yield session
        session.commit()


def session_scope(func):
    """会话管理装饰器（自动开启、提交、回滚、关闭会话）"""
    @wraps(func)
    def wrapper(*args, **kwargs):
        with SyncSessionLocal() as session:
            try:
                result = func(*args, session_=session, **kwargs)
                session.commit()
                return result
            except Exception as e:
                session.rollback()
                raise e
    return wrapper

1
2
3
4
5

    with get_session() as session:
        stmt = select(func.count()).where(col(TestData.status)==StatusEnums.ACTIVE)
        count = session.exec(stmt)
        if cnt:=count.first():
            print(f"get cnt {cnt}")  # get cnt xx

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

    with get_session() as session:
        # 更新 status 为inactive 的变成 banned 状态
        stmt_update =update(TestData).where(col(TestData.status)==StatusEnums.INACTIVE).values(status=StatusEnums.BANNED)
        session.exec(stmt_update)
        session.commit()
        # 删除所有 banned 状态的数据，慎重
        stmt_delete =delete(TestData).where(col(TestData.status)==StatusEnums.BANNED)
        session.exec(stmt_delete)
        session.commit()

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14

from sqlmodel import Session, select,insert
from sqlalchemy.dialects.postgresql import insert as pg_insert ,Insert
from .models.creation_template import CreationTemplate

def update_creation_template(session:Session,uuid:list[str],creation_type:Optional[str],template_id:Optional[str],commit:bool=False):
    if len(uuid)==0 or creation_type is None : # 兼容旧传参
        return 
    data  = [{"res_uuid": x,"creation_type":creation_type,"template_id":template_id } for x in uuid]
    stmt:Insert = pg_insert(CreationTemplate).values(data)
    # 这里一定要 重新赋值
    stmt = stmt.on_conflict_do_update(index_elements=[CreationTemplate.res_uuid],set_={CreationTemplate.creation_type:creation_type,CreationTemplate.template_id:template_id})
    session.execute(stmt)
    if commit:
        session.commit()

 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

class PGConflictBase:
    @session_scope # 装饰器注入session
    def conflict_or_do_nothing(self,session_: Session):
        """
            仅创建一次
        """
        stmt:Insert = pg_insert(self.__class__).values(self.model_dump(exclude_unset=True, exclude_none=False))
        pk = self.__class__.get_pk_fields()
        stmt = stmt.on_conflict_do_nothing(index_elements=pk)
        session_.exec(stmt)
        session_.commit()

    @session_scope # 装饰器注入session
    def conflict_or_do_update(self,session_: Session,exclude_fields: Optional[List[str]] = None):
        """
            仅创建一次
        """
        pk = self.__class__.get_pk_fields()
        exclude_fields = exclude_fields.extend(pk)
        instance_data = self.model_dump(exclude_unset=True，exclude_none=False)
        stmt:Insert = pg_insert(self.__class__).values(instance_data)
        update_values = {
            field: getattr(self, field)  # 直观获取当前实例的字段值，可读性极强
            for field in instance_data.keys()
            if field not in exclude_fields
        }
        stmt = stmt.on_conflict_do_update(index_elements=pk,set_=update_values)
        session_.exec(stmt)
        session_.commit()
    
    @classmethod
    def get_pk_fields(cls):
        # 检查模型并获取mapper对象
        mapper = inspect(cls)
        # 提取主键字段名
        return [col.key for col in mapper.primary_key]

1
2
3
4

class obj:
    rank_no:Optional[int]=Field(default=None,description="排序数，可以为空")
# 
order = col(obj.rank_no).asc().nulls_last() # 默认的情况下会将null 值排前面，使用nulls_last() 可以置后