python 线程池ThreadPoolExecutor

背景

最近项目中需要抽取数据库中某些数据组成一个缓存表，SQL倒是不复杂就是单纯的慢，慢到navicat跑崩了也没跑出来，然后就想着用python写脚本，根据时间字段做异步查询

代码

数据库连接写成伪码，过程抽象一下就是这样

import time
import pandas as pd
from concurrent.futures import ThreadPoolExecutor

# 需要保存的字段，sql中已查询
FIELDS = ["FIELDS_1",
          "FIELDS_2",
          "FIELDS_3",
          "FIELDS_4"
          "FIELDS_5"]
executor = ThreadPoolExecutor(max_workers=8)
df = pd.DataFrame(columns=FIELDS)


def get_data(i):
    # 创建oracle连接
    # 根据参数i，用时间参数format执行的sql
    # 执行sql并获取结果
    time.sleep(6)
    return [{"data": "data"}], i


def save_data(future):
    # 用pandas写入.xlsx格式文件中
    print("保存第{}组数据".format(future[1]))
    global df
    df = df.append(pd.DataFrame(data=future[0], columns=FIELDS))
    time.sleep(1)


def run():
    for i in range(100):
        future = executor.submit(get_data, i)
        # future.result()
        future.add_done_callback(save_data)


if __name__ == '__main__':
    run()
    executor.shutdown(wait=True)
    df.to_excel(SAVE_PATH)

线程池关闭

通过shutdown方法，可以等待线程池内任务结束，并且优雅的关闭线程池，避免了子线程还没跑完，主线程执行完了要退出然后就都一起退出了的尴尬

子线程结果获取

子线程的结果其实可以通过submit返回的future.result()方法来直接获取，但是那样会造成主线程阻塞，所以还是用future.add_done_callback()方法通过回调函数来处理结果。

回调函数传参

回调函数只支持一个参数future，所以执行过程中的其他参数需要通过submit提交时传到子线程的执行函数，然后通过返回参数，让回调函数入参