两种方式实现websocket获取数据库查询进度

原创

已于 2022-12-31 17:18:18 修改 · 2.5k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#websocket #python #pandas

于 2022-11-20 17:31:36 首次发布

本文介绍如何利用WebSocket实现实时查询数据库查询进度。详细展示了两种方式：一是通过额外接口获取进度，二是边查询边返回进度。同时，还介绍了如何使aiomysql支持pandas的read_sql_query方法。

两种方式实现websocket获取数据库查询进度

本文实现了两种方式用websocket连接实现前端通过API获取数据库查询进度，作为websocket初步学习的记录

内容目录概要：

使用额外接口获取指定数据库查询进度，查询进度的接口与获取数据的接口分开实现
查询数据的同时可以同时返回进度
使aiomysql支持pandas的read_sql_query方法

食用前提：

sanic web服务
pandas分块读取mysql数据库，示例中伪造了dataframe代替数据库中的数据，pandas异步查询mysql的实现方式放在最后额外的part

第一种方式，查询进度的接口和获取数据的接口分开实现

# -*- coding: utf-8 -*-
import asyncio
import json

import pandas as pd
from sanic import Sanic
from sanic import Websocket, Request

app = Sanic("sanic")

# 这里为了简单，将进度存储设置成了全局变量容器，合理的做法可以放到一个可共用的类、第三方存储位置等，并设置过期自动清理
query_progress = {
   
   }


class MysqlDatabaseQuery:
    async def read_query(self):
        data = pd.DataFrame({
   
   "id": [3, 5, 8], "name": ["a", "b", "c"]}, index=[1, 2, 3])
        rowcount = len(data)
        for idx, row in data.iterrows():
            await asyncio.sleep(1)
            rownumber = idx
            yield dict(row), f"{
     
     (rownumber / rowcount): .2%}"


# 方式一
# 尝试单独用一个接口来查询数据库query的进度
async def read_query_progress(requests: Request, ws: Websocket):
    query_id = requests.args.get("query_id")
    while True:
        progress = query_progress.get(str(query_id)) or "0%"
        if ws.recv_cancel:
            break
        _ = await ws.recv()
        await ws.send(progress)


# 异步处理的任务：储存数据库查询的进度
async def dispatch_query_progress(**context):
    global query_progress
    query_progress = context


# 获取的数据接口
async def read_query(requests: Request, ws: Websocket):
    data_reader = MysqlDatabaseQuery().read_query()
    records = []
    query_id = "1"  # 需要提前定义好查询id
    async for record, progress in data_reader:
        records.append(record)
        # sanic提供的发布信号并在后台执行任务的方法，定义事件名称，传递参数，任务会异步执行：https://sanic.dev/zh/guide/advanced/signals.html#%E6%B7%BB%E5%8A%A0%E4%BF%A1%E5%8F%B7-adding-a-signal
        await requests.app.dispatch("query.read.progress", context={
   
   query_id: progress})
        if ws.recv_cancel:
            break
    else:
        await ws.send(json.dumps(records))


app.add_websocket_route(read_query, "/query1")
app.add_websocket_route(read_query_progress, "/read_progress1")
app.add_signal(dispatch_query_progress, "query.read.progress")  # 添加接收信号后的处理方法handler到app中

if __name__ == '__main__':
    app.run(auto_reload=True, dev=True)

接口调用效果：

# 查询接口ws://127.0.0.1:8000/query1的返回数据：
[{"id": 3, "name": "a"}, {"id": 5, "name": "b"}, {"id": 8, "name": "c"}]

在查询接口查询的过程中，调用进度接口
在这里插入图片描述

第二种方式，数据边查询边返回进度

# -*- coding: utf-8 -*-
import asyncio
import json

import pandas as pd
from sanic import Sanic
from sanic import Websocket, Request

app = Sanic("sanic")

# 这里为了简单，将进度存储设置成了全局变量容器，合理的做法可以放到一个可共用的类、第三方存储位置等，并设置过期自动清理
query_progress = {
   
   }


class MysqlDatabaseQuery:
    async def read_query(self):
        data = pd.DataFrame({
   
   "id": [3, 5, 8], "name": ["a", "b", "c"]}, index=[1, 2, 3])
        rowcount = len(data)
        for idx