日志优化 (dataelement#261)

src/backend/bisheng/api/v1/chat.py

@@ -118,13 +118,20 @@ async def chat(
         *,
         flow_id: str,
         websocket: WebSocket,
+        t: Optional[str] = None,
         chat_id: Optional[str] = None,
         Authorize: AuthJWT = Depends(),
 ):
     """Websocket endpoint for chat."""
     try:
-        Authorize.jwt_required(auth_from='websocket', websocket=websocket)
-        payload = json.loads(Authorize.get_jwt_subject())
+        if t:
+            Authorize.jwt_required(auth_from='websocket', token=t)
+            Authorize._token = t
+        else:
+            Authorize.jwt_required(auth_from='websocket', websocket=websocket)
+
+        payload = Authorize.get_jwt_subject()
+        payload = json.loads(payload)
         user_id = payload.get('user_id')
         if chat_id:
             with next(get_session()) as session:

src/backend/bisheng/api/v2/rpc.py

@@ -5,7 +5,7 @@
 from bisheng.database.models.user import User
 from bisheng.database.models.user_role import UserRole
 from bisheng.settings import settings
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, HTTPException
 from fastapi.responses import RedirectResponse
 from fastapi_jwt_auth import AuthJWT
 from loguru import logger
@@ -25,45 +25,38 @@ def set_cookie(*,
                session: Session = Depends(get_session),
                Authorize: AuthJWT = Depends()):
     """设置默认"""
-    if deptId:
-        # this interface should update user model, and now the main ref don't mathes
-        db_user = session.exec(select(User).where(User.dept_id == deptId)).first()
-        if not db_user:
-            db_user = User(user_name=deptName, password='none', dept_id=deptId)
-            session.add(db_user)
-            session.flush()
-            db_user_role = UserRole(user_id=db_user.user_id, role_id=2)
-            session.add(db_user_role)
-            session.commit()
-            session.refresh(db_user)
-    if not deptId:
-        dept_user_id = settings.get_from_db('default_operator').get('user')
-    else:
-        dept_user_id = db_user.user_id
 
-    payload = {'user_name': deptName, 'user_id': dept_user_id, 'role': [2]}
     # admin
-    role_admin = session.query(UserRole).where(UserRole.user_id == user_id,
-                                               UserRole.role_id == 1).first()
     try:
-        if user_id and role_id == 1:
-            if not role_admin:
-                # keep
-                admin_user = session.get(User, user_id)
-                if not admin_user:
-                    db_user = User(id=user_id, user_name=deptName, password='none', dept_id=deptId)
-                    session.add(db_user)
-                db_user_role = UserRole(user_id=user_id, role_id=1)
+        if deptId:
+            # this interface should update user model, and now the main ref don't mathes
+            db_user = session.exec(select(User).where(User.dept_id == deptId)).first()
+            if not db_user:
+                db_user = User(user_name=deptName, password='none', dept_id=deptId)
+                session.add(db_user)
+                session.flush()
+                db_user_role = UserRole(user_id=db_user.user_id, role_id=2)
+                session.add(db_user_role)
+                session.commit()
+                session.refresh(db_user)
+        else:
+            raise ValueError('deptId 必须传递')
+        payload = {'user_name': deptName, 'user_id': db_user.user_id, 'role': [2]}
+        if role_id == 1:
+            admin_user = session.query(User).where(User.user_name == 'root').first()
+            if not admin_user:
+                admin_user = User(user_name='root', password='none')
+                session.add(admin_user)
+                session.flush()
+                session.refresh(admin_user)
+                db_user_role = UserRole(user_id=admin_user.user_id, role_id=1)
                 session.add(db_user_role)
-            payload = {'user_name': deptName, 'user_id': user_id, 'role': 'admin'}
-        elif user_id:
-            if role_admin:
-                # delete
-                session.delete(role_admin)
-        session.commit()
+            payload = {'user_name': 'root', 'user_id': admin_user.user_id, 'role': 'admin'}
+            session.commit()
     except Exception as e:
         logger.error(str(e))
         session.rollback()
+        return HTTPException(status_code=500, detail=str(e))
 
     # Create the tokens and passing to set_access_cookies or set_refresh_cookies
     access_token = Authorize.create_access_token(subject=json.dumps(payload), expires_time=864000)

src/backend/bisheng/database/base.py

@@ -68,7 +68,7 @@ def session_getter():
         session = Session(db_service.engine)
         yield session
     except Exception as e:
-        print('Session rollback because of exception:', e)
+        logger.info('Session rollback because of exception:{}', e)
         session.rollback()
         raise
     finally:

src/backend/bisheng/utils/threadpool.py

@@ -92,7 +92,8 @@ async def as_completed(self) -> List[Tuple[str, concurrent.futures.Future]]:
                             pending_count += 1
                 if len(lf) == 0:
                     self.async_task.pop(k)
-            logger.info('async_wait_count={}', pending_count)
+            if pending_count > 0:
+                logger.info('async_wait_count={}', pending_count)
             return completed_futures
 
     # async def async_done_callback(self, future):