python-fastapi-mm-zhcs-yj-api/common/websocketManager.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from fastapi import APIRouter, Request, Depends, Query, HTTPException, status,WebSocket,WebSocketDisconnect

from sqlalchemy.orm import Session
from models import *
from utils import *
from utils.ry_system_util import *

class ConnectionManager:
    def __init__(self):
        self.active_connections = {}  #: List[WebSocket]

    async def connect(self, websocket: WebSocket,pattern_id:str,db: Session):
        await websocket.accept()

        if pattern_id not in self.active_connections:

            self.active_connections[pattern_id] = []
        # else:
        self.active_connections[pattern_id].append(websocket)

        print('连接成功',len(self.active_connections[pattern_id]))
        # data = pattern_id_get_tp_pattern_ws_list(pattern_id, db)
        # message = []
        # for info in data:
        #     user = user_id_get_user_info(db, info.create_by)
        #     dept = dept_id_get_dept_info(db, user.dept_id)
        #     message.append({"id": info.id,
        #                     "name": info.name,
        #                     "pattern_id": info.pattern_id,
        #                     "content": info.content,
        #                     "visible": info.visible,
        #                     "user_id": info.create_by,
        #                     "nick_name": user.nick_name,
        #                     "dept_name": dept.dept_name})
        #     message = json.dumps(message)
        #     await websocket.send_text(message)
        #     await websocket.send_text('连接成功')
    def disconnect(self, websocket: WebSocket,pattern_id:str):
        # await websocket.send_text('已断开')
        print('已断开')
        self.active_connections[pattern_id].remove(websocket)
        if not self.active_connections[pattern_id]:
            del self.active_connections[pattern_id]
    async def broadcast(self,pattern_id:str,db: Session, websocket: WebSocket):#, message: str
        data = pattern_id_get_tp_pattern_ws_list(pattern_id,db)
        message = []
        for info in data:
            user= user_id_get_user_info(db,info.create_by)
            dept = dept_id_get_dept_info(db,user.dept_id)
            message.append({"id":info.id,
                            "name":info.name,
                            "pattern_id":info.pattern_id,
                            "content":info.content,
                            "visible":info.visible,
                            "user_id":info.create_by,
                            "nick_name":user.nick_name,
                            "dept_name":dept.dept_name})
        message = json.dumps(message)
        for connection in self.active_connections[pattern_id]:
            if connection!=websocket:
                await connection.send_text(message)
    async def ownmessage(self,websocket: WebSocket,pattern_id:str,db: Session):#, message: str
        data = pattern_id_get_tp_pattern_ws_list(pattern_id,db)
        message = []
        for info in data:
            user= user_id_get_user_info(db,info.create_by)
            dept = dept_id_get_dept_info(db,user.dept_id)
            message.append({"id":info.id,
                            "name":info.name,
                            "pattern_id":info.pattern_id,
                            "content":info.content,
                            "visible":info.visible,
                            "user_id":info.create_by,
                            "nick_name":user.nick_name,
                            "dept_name":dept.dept_name})
        message = json.dumps(message)
        await websocket.send_text(message)

manager = ConnectionManager()

def user_id_and_pattern_id_get_tp_pattern_ws_info(user_id:str,pattern_id:str,db: Session):
    query = db.query(TpPatternWSList)
    query = query.filter(TpPatternWSList.del_flag != '2')
    query = query.filter(TpPatternWSList.pattern_id == pattern_id)
    query = query.filter(TpPatternWSList.create_by == user_id)
    query = query.order_by(TpPatternWSList.create_time.desc())
    return query.first()

def pattern_id_get_tp_pattern_ws_users(pattern_id:str,db: Session):
    query = db.query(TpPatternWSList)
    query = query.filter(TpPatternWSList.del_flag != '2')
    query = query.filter(TpPatternWSList.pattern_id == pattern_id)
    query = query.group_by(TpPatternWSList.create_by)
    return query.all()

def pattern_id_get_tp_pattern_ws_list(pattern_id:str,db: Session):
    query = db.query(TpPatternWSList)
    query = query.filter(TpPatternWSList.del_flag != '2')
    query = query.filter(TpPatternWSList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSList.create_time)
    # query = query.group_by(TpPatternWSList.create_id)
    return query.all()

def pattern_id_get_tp_pattern_ws_count(pattern_id:str,db: Session):
    query = db.query(TpPatternWSList)
    query = query.filter(TpPatternWSList.del_flag != '2')
    query = query.filter(TpPatternWSList.pattern_id == pattern_id)
    return query.count()

def pattern_id_get_tp_pattern_ws_max_time(pattern_id:str,db: Session):
    query = db.query(TpPatternWSList.update_time)
    query = query.filter(TpPatternWSList.del_flag != '2')
    query = query.filter(TpPatternWSList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSUserList.update_time.desc())
    return query.first()

def pattern_id_get_tp_pattern_ws_user_list(pattern_id:str,db: Session):
    query = db.query(TpPatternWSUserList)
    query = query.filter(TpPatternWSUserList.del_flag != '2')
    query = query.filter(TpPatternWSUserList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSUserList.create_time.desc())
    return query.all()

def pattern_id_get_tp_pattern_ws_group_def_info(pattern_id:str,db: Session):
    query = db.query(TpPatternWSGroupList)
    query = query.filter(TpPatternWSGroupList.del_flag != '2')
    query = query.filter(TpPatternWSGroupList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSGroupList.create_time.desc())
    return query.first()

def user_id_get_tp_pattern_ws_user_list(user_id:str,db: Session):
    query = db.query(TpPatternWSUserList)
    query = query.filter(TpPatternWSUserList.del_flag != '2')
    query = query.filter(TpPatternWSUserList.user_id == user_id)
    query = query.order_by(TpPatternWSUserList.create_time.desc())
    return query.all()

def user_id_and_pattern_id_get_tp_pattern_ws_user_info(user_id:str,pattern_id:str,db: Session):
    query = db.query(TpPatternWSUserList)
    query = query.filter(TpPatternWSUserList.del_flag != '2')
    query = query.filter(TpPatternWSUserList.user_id == user_id)
    query = query.filter(TpPatternWSUserList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSUserList.create_time.desc())
    return query.first()
def pattern_id_get_tp_pattern_ws_group_list(pattern_id:str,db: Session):
    query = db.query(TpPatternWSGroupList)
    query = query.filter(TpPatternWSGroupList.del_flag != '2')
    query = query.filter(TpPatternWSGroupList.pattern_id == pattern_id)
    query = query.order_by(TpPatternWSGroupList.create_time.desc())
    return query.all()