chongming-plugins/example/router/count.py

from sqlalchemy import select, update

from chongming_worker import rpc_method

from ..models.pydantic import CountGetResponse
from ..models.counter import Counter


@rpc_method(
    subject="example.count",
    method="GET",
    summary="Increment and get counter value",
    response_model=CountGetResponse,  # 需要定义 CountResponse
)
async def count_get() -> CountGetResponse:
    from ..global_var import db
    
    if db is None:
        raise RuntimeError("Database not initialized")
    async with db.get_session() as session:
        statement = select(Counter).where(Counter.name == "visits") # type: ignore
        result = await session.execute(statement)
        counter = result.scalar_one_or_none()
        
        if counter is None:
            counter = Counter(name="visits", value=1)
            session.add(counter)
            await session.commit()
            return CountGetResponse(name="visits", value=1)
        
        else:
            new_value = counter.value + 1

            update_stmt = (
                update(Counter)
                .where(Counter.name == "visits") # type: ignore
                .values(value=new_value)
            )
            await session.execute(update_stmt)
            await session.commit()

            return CountGetResponse(name="visits", value=new_value)