368 lines
14 KiB
Python
368 lines
14 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""
|
||
@Time : 2025/6/16 14:12
|
||
@Auth : 九月的海
|
||
@File : suCaiDownload.py
|
||
@IDE : PyCharm
|
||
@Motto : September_sea
|
||
"""
|
||
import os
|
||
import time
|
||
import requests
|
||
import traceback
|
||
import random
|
||
from urllib.parse import urljoin
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from db_base import Material1_db # 导入数据库模块
|
||
|
||
|
||
class MaterialDownloader:
|
||
def __init__(self, base_url="https://tos.mobgi.com"):
|
||
"""
|
||
素材下载器初始化
|
||
|
||
:param base_url: 素材基础URL前缀
|
||
"""
|
||
self.base_url = base_url
|
||
self.download_dir = "D:\\创量\\A素材\\数据库素材批量下载" # 使用双反斜杠避免转义问题
|
||
self.base_save_directory = None
|
||
self.subfolder_count = 0
|
||
self.files_per_folder = 0
|
||
self.current_file_counts = {}
|
||
|
||
# 确保下载目录存在
|
||
self.ensure_directory_exists(self.download_dir)
|
||
|
||
def set_download_dir(self, directory):
|
||
"""设置下载目录并确保它存在"""
|
||
self.download_dir = directory
|
||
return self.ensure_directory_exists(directory)
|
||
|
||
def setup_subfolder_download(self, base_directory, subfolder_count=12, files_per_folder=40):
|
||
"""
|
||
设置分层文件夹下载结构
|
||
:param base_directory: 基础目录路径
|
||
:param subfolder_count: 子文件夹数量
|
||
:param files_per_folder: 每个子文件夹最大文件数量
|
||
:return: True表示设置成功，False表示失败
|
||
"""
|
||
self.base_save_directory = base_directory
|
||
self.subfolder_count = subfolder_count
|
||
self.files_per_folder = files_per_folder
|
||
|
||
# 创建基础目录
|
||
if not self.ensure_directory_exists(self.base_save_directory):
|
||
return False
|
||
|
||
# 创建子文件夹（1到subfolder_count）
|
||
for i in range(1, self.subfolder_count + 1):
|
||
folder_path = os.path.join(self.base_save_directory, str(i))
|
||
if not self.ensure_directory_exists(folder_path):
|
||
return False
|
||
|
||
# 初始化文件夹文件计数器
|
||
self.current_file_counts = {str(i): 0 for i in range(1, self.subfolder_count + 1)}
|
||
print(f"已创建子文件夹: 1-{subfolder_count}，每个文件夹最多保存 {files_per_folder} 个文件")
|
||
return True
|
||
|
||
def get_next_subfolder(self):
|
||
"""获取下一个有空间的子文件夹"""
|
||
if not self.current_file_counts:
|
||
return "1"
|
||
|
||
# 找出文件数量最少的文件夹
|
||
min_count = min(self.current_file_counts.values())
|
||
min_folders = [folder for folder, count in self.current_file_counts.items() if count == min_count]
|
||
|
||
# 随机选择一个最小文件的文件夹（确保分布均匀）
|
||
folder = random.choice(min_folders)
|
||
|
||
# 增加该文件夹的计数
|
||
self.current_file_counts[folder] += 1
|
||
return folder
|
||
|
||
def ensure_directory_exists(self, directory):
|
||
"""确保目录存在，如果不存在则创建"""
|
||
if not os.path.exists(directory):
|
||
try:
|
||
os.makedirs(directory)
|
||
print(f"创建下载目录: {directory}")
|
||
return True
|
||
except Exception as e:
|
||
print(f"无法创建目录 {directory}: {str(e)}")
|
||
return False
|
||
return True
|
||
|
||
def sanitize_filename(self, filename):
|
||
"""清理文件名中的无效字符"""
|
||
invalid_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*']
|
||
for char in invalid_chars:
|
||
filename = filename.replace(char, '_')
|
||
return filename.strip()
|
||
|
||
def download_file(self, full_url, save_path):
|
||
"""下载单个文件到本地"""
|
||
try:
|
||
start_time = time.time()
|
||
|
||
# 确保文件所在目录存在
|
||
save_dir = os.path.dirname(save_path)
|
||
if not self.ensure_directory_exists(save_dir):
|
||
return False, f"无法创建文件目录: {save_dir}"
|
||
response = requests.get(full_url, stream=True, timeout=30) # 增加超时时间到30秒
|
||
response.raise_for_status()
|
||
file_size = int(response.headers.get('Content-Length', 0))
|
||
downloaded = 0
|
||
last_percent = 0
|
||
# 使用临时文件名下载，完成后重命名为目标文件名
|
||
temp_path = f"{save_path}.download"
|
||
|
||
with open(temp_path, 'wb') as f:
|
||
for chunk in response.iter_content(chunk_size=8192):
|
||
if chunk:
|
||
f.write(chunk)
|
||
downloaded += len(chunk)
|
||
# 显示下载进度（当文件大小已知时）
|
||
if file_size > 0:
|
||
percent = int(downloaded * 100 / file_size)
|
||
# 每10%的进度更新一次显示
|
||
if percent - last_percent >= 10 or percent == 100:
|
||
print(
|
||
f"已下载: {percent}% ({downloaded / 1024 / 1024:.2f}MB/{file_size / 1024 / 1024:.2f}MB)")
|
||
last_percent = percent
|
||
|
||
# 下载完成后重命名
|
||
os.rename(temp_path, save_path)
|
||
|
||
elapsed = time.time() - start_time
|
||
size_mb = os.path.getsize(save_path) / (1024 * 1024)
|
||
speed = size_mb / elapsed if elapsed > 0 else 0
|
||
|
||
return True, f"下载成功! 大小: {size_mb:.2f}MB, 耗时: {elapsed:.1f}秒, 速度: {speed:.1f}MB/s"
|
||
|
||
except Exception as e:
|
||
# 清理临时文件
|
||
if 'temp_path' in locals() and os.path.exists(temp_path):
|
||
os.remove(temp_path)
|
||
return False, f"下载失败: {str(e)}"
|
||
|
||
def process_material(self, material):
|
||
"""处理单个素材记录"""
|
||
# 确保我们有正确的三个字段
|
||
if len(material) < 3:
|
||
return {
|
||
'material_id': 'unknown',
|
||
'material_name': 'invalid_data',
|
||
'url': 'invalid_data',
|
||
'save_path': 'invalid_data',
|
||
'success': False,
|
||
'message': f"数据库记录字段不足，预期3个，实际得到{len(material)}个: {material}"
|
||
}
|
||
|
||
material_id = material[0]
|
||
material_name = material[1] if len(material) > 1 else 'unknown'
|
||
file_uri = material[2] if len(material) > 2 else 'unknown'
|
||
# 生成完整URL
|
||
full_url = urljoin(self.base_url, file_uri) if file_uri != 'unknown' else 'unknown'
|
||
# 清理文件名
|
||
safe_name = self.sanitize_filename(material_name)
|
||
# 提取文件扩展名
|
||
extension = os.path.splitext(file_uri)[1] if file_uri != 'unknown' else '.bin'
|
||
# 生成保存路径 (文件名格式: 素材ID_清理后的素材名.扩展名)
|
||
# 确保文件名中不包含路径分隔符
|
||
filename = f"{material_id}_{safe_name}{extension}"
|
||
|
||
# 确定保存位置 - 如果有设置子文件夹结构，则使用；否则使用统一目录
|
||
if self.base_save_directory and self.subfolder_count > 0:
|
||
# 获取下一个有空间的子文件夹
|
||
folder_name = self.get_next_subfolder()
|
||
# 生成路径：基础路径/文件夹名/文件名
|
||
save_path = os.path.join(self.base_save_directory, folder_name, filename)
|
||
else:
|
||
save_path = os.path.join(self.download_dir, filename)
|
||
|
||
# 检查文件是否已存在
|
||
if os.path.exists(save_path):
|
||
return {
|
||
'material_id': material_id,
|
||
'material_name': material_name,
|
||
'url': full_url,
|
||
'save_path': save_path,
|
||
'success': True,
|
||
'message': f"文件已存在 (大小: {os.path.getsize(save_path) / (1024 * 1024):.2f}MB)"
|
||
}
|
||
|
||
# 如果URL无效，直接返回失败
|
||
if full_url == 'unknown':
|
||
return {
|
||
'material_id': material_id,
|
||
'material_name': material_name,
|
||
'url': full_url,
|
||
'save_path': save_path,
|
||
'success': False,
|
||
'message': "无效的文件URI"
|
||
}
|
||
|
||
# 下载文件
|
||
success, message = self.download_file(full_url, save_path)
|
||
return {
|
||
'material_id': material_id,
|
||
'material_name': material_name,
|
||
'url': full_url,
|
||
'save_path': save_path,
|
||
'success': success,
|
||
'message': message
|
||
}
|
||
|
||
def batch_download(self, sql_query=None, max_workers=5):
|
||
"""
|
||
批量下载素材
|
||
|
||
:param sql_query: 可选的自定义SQL查询
|
||
:param max_workers: 并发下载线程数
|
||
:return: 下载结果列表
|
||
"""
|
||
# 默认查询语句
|
||
if not sql_query:
|
||
sql_query = """
|
||
SELECT material_id, material_name, file_uri
|
||
FROM `material`
|
||
USE INDEX (idx_cTime_mainUserId)
|
||
WHERE main_user_id = '12400078500'
|
||
AND create_time BETWEEN '2025-04-29 17:53:37' AND '2025-04-30 17:53:37'
|
||
AND create_time >= '2025-04-27'
|
||
AND material_type = 'video'
|
||
AND file_uri LIKE 'tos_beijing/%'
|
||
LIMIT 50
|
||
"""
|
||
|
||
# 使用正确的数据库查询方式
|
||
try:
|
||
print(f"正在执行SQL查询: {sql_query}")
|
||
# 通过模块方法查询数据库
|
||
materials = Material1_db.query_database(sql_query)
|
||
if not materials:
|
||
print("未查询到符合条件的素材")
|
||
return []
|
||
|
||
# 打印查询结果中的第一条记录
|
||
if materials:
|
||
print(f"查询到的第一条记录: {materials[0]}")
|
||
|
||
print(f"\n查询到 {len(materials)} 个素材, 开始下载...\n")
|
||
# 保存开始时间
|
||
start_time = time.time()
|
||
# 使用线程池并发下载
|
||
results = []
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
futures = [executor.submit(self.process_material, mat) for mat in materials]
|
||
|
||
# 使用tqdm显示进度（可选）
|
||
try:
|
||
from tqdm import tqdm
|
||
for future in tqdm(futures, desc="下载进度", unit="文件"):
|
||
result = future.result()
|
||
results.append(result)
|
||
# 立即显示结果
|
||
self.print_result(result)
|
||
except ImportError:
|
||
# 如果没有安装tqdm，使用普通循环
|
||
for future in futures:
|
||
result = future.result()
|
||
results.append(result)
|
||
# 立即显示结果
|
||
self.print_result(result)
|
||
|
||
# 计算总耗时
|
||
total_time = time.time() - start_time
|
||
|
||
# 打印统计信息
|
||
success_count = sum(1 for r in results if r['success'])
|
||
print(f"\n{'=' * 60}")
|
||
print(f"下载完成! 成功: {success_count}/{len(results)}")
|
||
print(f"总耗时: {total_time:.2f}秒")
|
||
if self.base_save_directory:
|
||
print(f"下载目录: {self.base_save_directory} (分为{self.subfolder_count}个子文件夹)")
|
||
else:
|
||
print(f"下载目录: {self.download_dir}")
|
||
print(f"{'=' * 60}")
|
||
|
||
# 列出所有成功下载的文件
|
||
if success_count > 0:
|
||
print("\n成功下载的文件:")
|
||
for result in results:
|
||
if result['success']:
|
||
print(f"- {os.path.basename(result['save_path'])}")
|
||
|
||
# 如果有失败的下载，创建失败列表
|
||
failed_downloads = [r for r in results if not r['success']]
|
||
if failed_downloads:
|
||
print("\n下载失败的文件:")
|
||
for result in failed_downloads:
|
||
print(f"- {result['material_name']} (ID: {result['material_id']}) - 原因: {result['message']}")
|
||
|
||
return results
|
||
|
||
except Exception as e:
|
||
# 避免使用特殊Unicode字符
|
||
print(f"\n[错误] 批量下载失败: {str(e)}")
|
||
traceback.print_exc() # 打印详细的错误堆栈
|
||
return []
|
||
|
||
def print_result(self, result):
|
||
"""打印单个下载结果"""
|
||
material_info = f"{result['material_id']} - {result['material_name']}"
|
||
save_path = result['save_path']
|
||
message = result['message']
|
||
|
||
# 下载成功的特殊格式
|
||
if result['success']:
|
||
print(f"\n[成功] {material_info} - {message}")
|
||
else:
|
||
print(f"\n[失败] {material_info} - {message}")
|
||
|
||
# 输出保存路径
|
||
print(f" 保存位置: {save_path}")
|
||
if 'url' in result and result['url'] != 'unknown':
|
||
print(f" 下载地址: {result['url']}")
|
||
|
||
|
||
# 使用示例
|
||
if __name__ == "__main__":
|
||
downloader = MaterialDownloader()
|
||
|
||
# 设置分层文件夹下载结构
|
||
# 在该目录下创建30个子文件夹（1-30），每个文件夹最多存放2个文件
|
||
if not downloader.setup_subfolder_download(
|
||
base_directory="D:\\创量\\A素材\\80个文件夹素材",
|
||
subfolder_count=80,
|
||
files_per_folder=1
|
||
):
|
||
print("无法创建下载目录结构")
|
||
exit(1)
|
||
|
||
print("+" * 60)
|
||
print("开始素材批量下载程序")
|
||
print(f"下载基础目录: {downloader.base_save_directory}")
|
||
print(f"子文件夹: 1-{downloader.subfolder_count}")
|
||
print(f"每个文件夹最多存放: {downloader.files_per_folder}个文件")
|
||
print("+" * 60)
|
||
print()
|
||
# 使用自定义查询下载
|
||
custom_sql = """
|
||
SELECT material_id, material_name,file_uri
|
||
FROM `material`
|
||
USE INDEX (idx_cTime_mainUserId)
|
||
WHERE main_user_id = '12000016859'
|
||
AND create_time BETWEEN '2025-07-14 17:53:37' AND '2025-07-18 17:53:37' -- 限定时间范围
|
||
AND video_duration BETWEEN 10 AND 30
|
||
AND material_type = 'video'
|
||
AND file_uri LIKE 'tos_beijing/%'
|
||
AND material_name NOT LIKE '%-衍生'
|
||
LIMIT 80;
|
||
"""
|
||
# 设置并行下载数量
|
||
max_workers = 5
|
||
print(f"使用 {max_workers} 个并发线程下载")
|
||
|
||
results = downloader.batch_download(sql_query=custom_sql, max_workers=max_workers) |