python爬虫数据提取脚本

背景

使用scrapy_redis框架时，为了利用redis高效的io，通常回把数据放到redis上。在任务结束后则需要将redis里的结果数据迁移。

场景

将爬虫数据从redis迁移到mongodb，并且提取生成一份excel格式的文件

使用方法

参数2：指定爬虫结果集

参数3~：指定指定结果集中键在excel中的别名

python export.py [spider_name]:items redis_key1=excel_title1 redis_key2=excel_title2

脚本

有些地方些的不太好，后续优化~

import sys
import time
import redis
import pymongo
import json
import openpyxl


# 配置
REDIS_HOST = "redis host"
REDIS_PORT = 6379
REDIS_DB = 0
REDIS_PWD = "redis password"
REDIS_COL = ""  # 导出数据的目标集合
MONGO_HOST = "mongodb host"
MONGO_PORT = 27017
MONGO_DB = ""   # 数据存储的库名
MONGO_COL = ""  # 数据存储的集合名
MONGO_USER = "mongodb username"
MONGO_PWD = "mongodb password"


db_field = []   # 数据库字段
file_field = []     # 存入表的字段


def string_date():
    return time.strftime("%Y_%m_%d", time.localtime())


class MoveData:
    def __init__(self, redis_db, redis_col, db_field: list, file_field: list):
        """
        满足调用MoveData实例完成数据迁移
        :param redis_db: 迁移数据所在redis库
        :param redis_col: 迁移数据所在redis集合     DanKeRent:items
        :param db_field: 提取迁移数据的字段名
        :param file_field: 数据保存文件的字段名
        :param path: 保存数据的文件路径
        """
        self.redis_db = redis_db
        self.redis_col = redis_col
        self.db_field = db_field
        self.file_field = file_field
        self.mongo_db = self.redis_col.split(":")[0]       # mongo库名不改
        self.mongo_col = self.mongo_db + "_{}".format(string_date())   # mongo集合名添加日期后缀
        self.redis = redis.Redis(
            host=REDIS_HOST,
            port=REDIS_PORT,
            db=redis_db,
            password=REDIS_PWD
        )
        self.mongo = pymongo.MongoClient(
            host=MONGO_HOST,
            port=MONGO_PORT,
            username=MONGO_USER,
            password=MONGO_PWD
        )
        self.mongo_db_handle = self.mongo[self.mongo_db]
        self.mongo_col_handle = self.mongo_db_handle[self.mongo_col]
        self.workbook = None    # excel文件对象
        self.sheet = None
        self.write_line = 0

    def skip_redis(self):
        self.redis_empty()

    def move(self):
        # 将数据从redis迁移到mongo
        redis_data_len = self.redis.llen(self.redis_col)
        for index in range(redis_data_len):
            data = self.redis.lindex(self.redis_col, index)
            data = json.loads(data)
            print("导出redis缓存 {}: {}".format(index, data))
            self.mongo_col_handle.insert(data)
            print("导出redis缓存 {}: {}".format(index, data))
        if not redis_data_len:
            print("redis中无对应数据，直接从mongodb中导出数据")
            self.redis_empty()

    def redis_empty(self):
        # 当要迁移的redis列表不存在时(可能输入错误，也可能已经迁移过了)
        colls = self.mongo_db_handle.list_collection_names()
        select = max([time.mktime(time.strptime(coll[-10:], "%Y_%m_%d")) for coll in colls])
        select_date = time.strftime("%Y_%m_%d", time.localtime(select))
        self.mongo_col = self.mongo_db + "_{}".format(select_date)
        self.mongo_col_handle = self.mongo_db_handle[self.mongo_col]

    def create_file(self):
        self.workbook = openpyxl.Workbook()     # 新建工作簿
        self.sheet = self.workbook.create_sheet(self.redis_col.split(":")[1])    # 以redis集合名作为excel sheet名
        for cols in range(len(file_field)):
            self.sheet.cell(self.write_line + 1, cols + 1).value = file_field[cols]
        print("写入列名: {}".format(file_field))
        self.write_line += 1

    def save_row(self, data):
        for row in range(len(db_field)):
            self.sheet.cell(self.write_line + 1, row + 1).value = data.get(db_field[row], None)
        print("提取数据 {}: {}".format(self.write_line, data))
        self.write_line += 1

    def to_file(self):
        # 将迁移到mongodb的数据保存到文件(excel)
        # 1.创建文件excel
        self.create_file()
        # 2.遍历提取数据
        print("正在导表：{}".format(self.mongo_col))
        for data in self.mongo_col_handle.find():
            self.save_row(data)
        # 3.保存数据
        save_path = "{}.xlsx".format(self.mongo_col)
        self.workbook.save(save_path)

    def over(self):
        # 删除redis缓存
        self.redis.delete(self.redis_col)

    def run(self):
        self.move()
        # self.skip_redis()
        self.to_file()
        self.over()


if __name__ == '__main__':
    # 需要迁移的redis列表
    REDIS_COL = sys.argv[1]
    for i in sys.argv[2:]:
        before, after = i.split("=")
        db_field.append(before)
        file_field.append(after)
    m = MoveData(
        redis_db=REDIS_DB,
        redis_col=REDIS_COL,        # spider_name:items
        db_field=db_field,          # 字段列表
        file_field=file_field,      # 导表表头列表
    )
    m.run()
    time.sleep(10)