from datetime import datetime
|
from typing import Type
|
|
from sqlalchemy import Column, Integer, DateTime, Text, String
|
from sqlalchemy.orm import Session
|
|
from app.config.const import RAGFLOW
|
# from Log import logger
|
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))
|
bisheng_token = Column(Text(10000))
|
ragflow_token = Column(Text(10000))
|
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() # 回滚事务
|
|
|
async def update_token(db: Session, user_id: int, access_token: str, token: dict):
|
# 参数验证
|
if not isinstance(user_id, int) or user_id <= 0:
|
return
|
db_token = None
|
# print(token)
|
try:
|
# 查询现有记录
|
db_token = db.query(TokenModel).filter_by(user_id=user_id).first()
|
if db_token:
|
# 记录存在，进行更新
|
db_token.token = access_token
|
db_token.ragflow_token = token.get(RAGFLOW, "")
|
# db_token.token = access_token
|
# for k, v in token.items():
|
# setattr(db_token, k.replace("app", "token"), v)
|
else:
|
# 记录不存在，进行插入
|
db_token = TokenModel(
|
user_id=user_id,
|
token=access_token,
|
ragflow_token = token.get(RAGFLOW, "")
|
)
|
# for k, v in token.items():
|
# setattr(db_token, k.replace("app", "token"), v)
|
db.add(db_token)
|
|
# 提交事务
|
db.commit()
|
db.refresh(db_token)
|
|
except Exception as e:
|
# logger.error(e)
|
# 异常处理
|
db.rollback() # 回滚事务
|
|
|
|
class ApiTokenModel(Base):
|
__tablename__ = "chat_api_tokens"
|
id = Column(String(36), primary_key=True)
|
app_id = Column(String(36), index=True)
|
type = Column(String(16))
|
token = Column(String(255))
|
created_at = Column(DateTime, default=datetime.now())
|
last_used_at = Column(DateTime, default=datetime.now())
|
