rag-gateway/app/models/token_model.py

from datetime import datetime
from typing import Type

from sqlalchemy import Column, Integer, DateTime, Text
from sqlalchemy.orm import Session

from app.models.base_model import Base


class TokenModel(Base):
    __tablename__ = "token"
    id = Column(Integer, primary_key=True, index=True)
    user_id = Column(Integer, index=True)
    token = Column(Text(10000), unique=True, index=True)
    bisheng_token = Column(Text(10000), unique=True, index=True)
    ragflow_token = Column(Text(10000), unique=True, index=True)
    created_at = Column(DateTime, default=datetime.utcnow)


def upsert_token(db: Session, user_id: int, access_token: str, bisheng_token: str, ragflow_token: str):
    # 参数验证
    if not isinstance(user_id, int) or user_id <= 0:
        return
    if not access_token or not bisheng_token or not ragflow_token:
        return
    db_token = None
    try:
        # 查询现有记录
        existing_token = db.query(TokenModel).filter_by(user_id=user_id).first()

        if existing_token:
            # 记录存在，进行更新
            existing_token.token = access_token
            existing_token.bisheng_token = bisheng_token
            existing_token.ragflow_token = ragflow_token
        else:
            # 记录不存在，进行插入
            db_token = TokenModel(
                user_id=user_id,
                token=access_token,
                bisheng_token=bisheng_token,
                ragflow_token=ragflow_token
            )
            db.add(db_token)

        # 提交事务
        db.commit()
        db.refresh(db_token)

    except Exception as e:
        # 异常处理
        db.rollback()  # 回滚事务


def get_token(db: Session, user_id: int) -> Type[TokenModel] | None:
    return db.query(TokenModel).filter_by(user_id=user_id).first()