python异步切片下载文件（内置redis获取任务 mongo更新任务状态等）

异步切片下载二进制文件并上传桶删除本地文件

import json
import os
import asyncio
from urllib import parseimport aiohttp
import aioredis
from motor.motor_asyncio import AsyncIOMotorClient
from retrying import retry
from minio import Minio
from minio.error import S3Error
import loguruclass Async_download:def __init__(self, sema_number=20, redis_address='redis://ip:6379/db', redis_password='passwd'):self.__sema_number = sema_numberself.__sema = Noneself.redis_address = redis_addressself.redis_password = redis_passwordself.minio_config = {"endpoint": "ip:port","access_key": "******","secret_key": "******","secure": False}self.mongo_config = {"uri": "mongodb://{0}:{1}@ip:27017".format(parse.quote_plus("user"),parse.quote_plus("passwd")),"db": "******","collection": "******",}@classmethoddef _mkdir(cls, file_path, full_path):if os.path.exists(file_path):return Truetry:os.makedirs(full_path)except Exception as e:passreturn Falsedef __init_check(self, file_path: str):full_path, file_name = file_path.rsplit('/', 1)file_size = os.path.getsize(file_path) if self._mkdir(file_path, full_path) else 0return file_name, file_size@classmethoddef __sync_save_local(cls, r_headers, results, file_path):done, padding = resultsfor d in done:for index, value in d.result().items():r_headers[index] = valuewith open(file_path, 'ab') as f:for _, value in r_headers.items():f.write(value)return True@classmethoddef __generate_headers(cls, headers, file_size, first_byte):r_headers = {}index = 0if first_byte > 51200000:byte = 2048000  # 2M 为一片else:byte = 1024000  # 1M 为一片while True:file_size_two = file_size + byteif file_size_two >= first_byte:r_headers[index] = {"Range": f"bytes={file_size}-{first_byte}"}breakr_headers[index] = {"Range": f"bytes={file_size}-{file_size_two - 1}"}index += 1file_size = file_size_twofor key in r_headers:r_headers[key].update(headers)return r_headers@retry(stop_max_attempt_number=3)async def __download_one(self, session, method, url, r_headers, **kwargs):index, headers = r_headersasync with self.__sema:async with session.request(method, url, headers=headers, **kwargs) as response:binary = await response.content.read()return {index: binary}async def __async_section_download(self, session, method, url, r_headers, **kwargs):tasks = [asyncio.create_task(self.__download_one(session, method, url, (key, r_headers[key]), **kwargs)) for key inr_headers]return await asyncio.wait(tasks)@classmethodasync def __get_content_length(cls, session, method, url, headers, **kwargs):async with session.request(method, url, headers=headers, **kwargs) as response:return response.headers.get('Content-Length') or response.headers.get('content-length') or 0@classmethodasync def __sync_download(cls, session, method, url, headers, file_path, **kwargs):async with session.request(method, url, headers=headers, **kwargs) as response:with open(file_path, 'wb') as f:binary = await response.content.read()f.write(binary)async def __async_download_main(self, method, url, headers, file_path, **kwargs):file_name, file_size = self.__init_check(file_path)self.__sema = asyncio.Semaphore(self.__sema_number)async with aiohttp.ClientSession() as session:content_length = await self.__get_content_length(session, method, url, headers, **kwargs)if content_length and content_length.isdigit():content_length = int(content_length)if file_size >= content_length:await self.__upload_to_minio(file_path, file_name)  # Upload to MinIOawait self.__update_mongo_status(file_name, True)  # Update MongoDB statusos.remove(file_path)  # Delete local filereturn True, file_pathr_headers = self.__generate_headers(headers, file_size, content_length)results = await self.__async_section_download(session, method, url, r_headers, **kwargs)self.__sync_save_local(r_headers, results, file_path)else:await self.__sync_download(session, method, url, headers, file_path, **kwargs)if os.path.getsize(file_path) >= int(content_length):await self.__upload_to_minio(file_path, file_name)  # Upload to MinIOawait self.__update_mongo_status(file_name, True)  # Update MongoDB statusos.remove(file_path)  # Delete local filereturn True, file_pathreturn False, file_pathasync def __get_task_from_redis(self):async with aioredis.from_url(self.redis_address, password=self.redis_password) as redis:task = await redis.lpop('file_file_all')return taskasync def __process_redis_tasks(self):while True:task_info = await self.__get_task_from_redis()if task_info is None:breaktask = json.loads(task_info)try:method = 'get'url = task["file_link"]file_path = './{}'.format(task["file_name"])headers = {'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.93 Safari/537.36'}if url.startswith('http') or url.startswith('https'):try:ddd = url.split('http')if len(ddd) != 0:url = 'http' + ddd[-1]except:continue# Perform download for the current taskawait self.__async_download_main(method, url, headers, file_path)except Exception as e:loguru.logger.error("Error processing Redis task:", e)async def __upload_to_minio(self, file_path, object_name):"""上传minio"""try:minioClient = Minio(**self.minio_config)check_bucket = minioClient.bucket_exists("******")if not check_bucket:minioClient.make_bucket("******")loguru.logger.info("start upload file to MinIO")minioClient.fput_object(bucket_name="******", object_name=object_name, file_path=file_path)loguru.logger.info("file {0} is successfully uploaded to MinIO".format(object_name))except FileNotFoundError as err:loguru.logger.info('*' * 10)loguru.logger.error('MinIO upload failed: ' + str(err))except S3Error as err:loguru.logger.error("MinIO upload failed:", err)async def __update_mongo_status(self, file_name, status):"""更新mongo采集状态"""try:mongo_uri = self.mongo_config["uri"]db_name = self.mongo_config["db"]collection_name = self.mongo_config["collection"]client = AsyncIOMotorClient(mongo_uri)db = client.get_database(db_name)collection = db.get_collection(collection_name)await collection.update_one({"file_name": file_name}, {"$set": {"status": status}})except Exception as e:loguru.logger.error("MongoDB update failed:", e)async def start(self):await self.__process_redis_tasks()loguru.logger.add("download_file_output.log", rotation="500 MB", level="DEBUG")
if __name__ == '__main__':as_dw = Async_download(20)asyncio.run(as_dw.start())

部分代码来源于y小白的笔记