729 lines
27 KiB
Python
729 lines
27 KiB
Python
"""
|
||
品牌爬虫蓝图：展开文件夹、运行任务、任务列表/详情、下载结果
|
||
"""
|
||
import os
|
||
import json
|
||
import threading
|
||
import traceback
|
||
import zipfile
|
||
import io
|
||
import time
|
||
import requests
|
||
from queue import Queue, Empty
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from urllib.parse import urlparse
|
||
|
||
from flask import Blueprint, request, jsonify, session, send_file, redirect, Response
|
||
|
||
from app_common import get_db, login_required, BASE_DIR
|
||
from config import bucket_path,JAVA_API_BASE
|
||
from brand_spider.main import single_file_handle, TaskCancelledError
|
||
|
||
brand_bp = Blueprint('brand', __name__)
|
||
BRAND_OUTPUT_DIR = os.path.join(BASE_DIR, 'brand_output')
|
||
OSS_PREFIX = "brand_results"
|
||
|
||
# 任务完成时推送给前端的 SSE 队列：task_id -> [Queue, ...]
|
||
_task_event_queues = {}
|
||
_task_event_lock = threading.Lock()
|
||
|
||
# 任务行级进度（当前处理文件的行数/总行数），仅存内存，不落库：
|
||
# task_id -> {
|
||
# 'file_index': int, # 当前处理第几个文件
|
||
# 'file_total': int, # 总文件数
|
||
# 'file_path': str, # 当前文件路径
|
||
# 'current_line': int, # 当前行号（或当前品牌序号）
|
||
# 'total_lines': int, # 总行数（或总品牌数）
|
||
# }
|
||
_task_line_progress = {}
|
||
_task_line_progress_lock = threading.Lock()
|
||
_cancelled_brand_tasks = set()
|
||
_cancelled_brand_tasks_lock = threading.Lock()
|
||
|
||
|
||
def _mark_task_cancelled(task_id):
|
||
with _cancelled_brand_tasks_lock:
|
||
_cancelled_brand_tasks.add(task_id)
|
||
|
||
|
||
def _is_task_cancelled(task_id):
|
||
with _cancelled_brand_tasks_lock:
|
||
return task_id in _cancelled_brand_tasks
|
||
|
||
|
||
def _get_uid_from_request_headers():
|
||
"""
|
||
从请求头读取 uid。
|
||
前端会在 headers 里携带 uid，用于确定本次请求的归属用户。
|
||
"""
|
||
uid = request.headers.get('uid')
|
||
if uid is None:
|
||
return None
|
||
uid = str(uid).strip()
|
||
if not uid:
|
||
return None
|
||
try:
|
||
return int(uid)
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _resolve_user_id():
|
||
"""
|
||
优先使用请求头 uid；没有请求头 uid 时回退到 session['user_id']。
|
||
若二者同时存在但不一致，则视为无效请求并返回 None。
|
||
"""
|
||
req_uid = _get_uid_from_request_headers()
|
||
if req_uid is not None:
|
||
sess_uid = session.get('user_id')
|
||
if sess_uid is not None:
|
||
try:
|
||
if int(sess_uid) != int(req_uid):
|
||
return None
|
||
except Exception:
|
||
if str(sess_uid) != str(req_uid):
|
||
return None
|
||
return req_uid
|
||
return session.get('user_id')
|
||
|
||
|
||
def _get_user_id_or_error():
|
||
user_id = _resolve_user_id()
|
||
if not user_id:
|
||
return None, (jsonify({'success': False, 'error': '未登录或缺少uid'}), 401)
|
||
try:
|
||
return int(user_id), None
|
||
except Exception:
|
||
return None, (jsonify({'success': False, 'error': 'uid无效'}), 400)
|
||
|
||
|
||
def _push_task_event(task_id, event):
|
||
"""向订阅了该任务的所有 SSE 连接推送事件，并移除该任务的队列列表。
|
||
同时清理内存中的行级进度，不通过数据库中转行级进度。"""
|
||
with _task_event_lock:
|
||
queues = _task_event_queues.pop(task_id, [])
|
||
for q in queues:
|
||
try:
|
||
q.put_nowait(event)
|
||
except Exception:
|
||
pass
|
||
# 任务进入终态时，顺便清理行级进度缓存
|
||
with _task_line_progress_lock:
|
||
_task_line_progress.pop(task_id, None)
|
||
status = (event or {}).get('status') if isinstance(event, dict) else None
|
||
if status in ('success', 'failed', 'cancelled'):
|
||
with _cancelled_brand_tasks_lock:
|
||
_cancelled_brand_tasks.discard(task_id)
|
||
|
||
|
||
def _expand_folder_xlsx(folder_path):
|
||
"""返回文件夹下所有 .xlsx 文件的绝对路径列表"""
|
||
if not folder_path or not os.path.isdir(folder_path):
|
||
return []
|
||
paths = []
|
||
for name in os.listdir(folder_path):
|
||
if name.endswith('.xlsx') or name.endswith('.XLSX'):
|
||
paths.append(os.path.normpath(os.path.join(folder_path, name)))
|
||
return paths
|
||
|
||
|
||
def _upload_local_xlsx_to_oss(local_path, user_id):
|
||
"""将本地 xlsx 文件上传到 OSS，返回 (url, None) 或 (None, error_message)。"""
|
||
if not local_path or not os.path.isfile(local_path):
|
||
return None, "文件不存在"
|
||
try:
|
||
from ali_oss import upload_file as oss_upload_file
|
||
except ImportError:
|
||
return None, "OSS 模块未配置"
|
||
base = os.path.basename(local_path)
|
||
safe_base = "".join(c if c.isalnum() or c in '-_.' else '_' for c in base)
|
||
if not safe_base.endswith('.xlsx'):
|
||
safe_base = safe_base + '.xlsx'
|
||
key = f"{bucket_path}brand_input/{user_id}/{int(time.time() * 1000)}/{safe_base}"
|
||
try:
|
||
with open(local_path, "rb") as f:
|
||
content = f.read()
|
||
url = oss_upload_file(content, key)
|
||
return url, None
|
||
except Exception as e:
|
||
return None, str(e)
|
||
|
||
|
||
def _run_brand_single(taskid,brand_ls, fileUrl,strategy,totalLines,chunkTotal,chunkIndex):
|
||
""""""
|
||
invalidBrands = [] # 不符合品牌的数据
|
||
queryFailedBrands = [] # 查询失败的数据
|
||
keptRows = []
|
||
for brand in brand_ls:
|
||
if _is_task_cancelled(taskid):
|
||
raise TaskCancelledError("任务已取消")
|
||
faild_data,query_faild_data = single_file_handle(brand,strategy)
|
||
invalidBrands.extend(faild_data)
|
||
queryFailedBrands.extend(query_faild_data)
|
||
if len(faild_data) == 0 and len(query_faild_data) == 0:
|
||
keptRows.append(brand)
|
||
|
||
data = { "strategy": strategy ,
|
||
"files": [
|
||
{
|
||
"fileUrl": fileUrl,
|
||
"originalFilename": "",
|
||
"relativePath": "",
|
||
"mainSheetName": "",
|
||
"chunkIndex": chunkIndex,
|
||
"chunkTotal": chunkTotal,
|
||
"totalLines": totalLines,
|
||
"keptRows": keptRows,
|
||
"invalidBrands": invalidBrands,
|
||
"queryFailedBrands": queryFailedBrands
|
||
}
|
||
]
|
||
}
|
||
resp = requests.post(f"{JAVA_API_BASE}/api/brand/tasks/{taskid}/result",headers={"accept":"application/json"},
|
||
# json={ "strategy": strategy ,
|
||
# "files": [
|
||
# { "fileUrl": fileUrl,"originalFilename": "","relativePath": "",
|
||
# "mainSheetName": "","columns": [],"keptRows": [],
|
||
# "invalidBrands": invalidBrands,"queryFailedBrands": queryFailedBrands
|
||
# }]})
|
||
json=data)
|
||
print(data)
|
||
print(taskid,brand_ls, fileUrl)
|
||
print("提交结果",data,"\n-->",resp.text)
|
||
return True
|
||
|
||
|
||
|
||
def _background_brand_task(task_id,data):
|
||
"""
|
||
后台执行爬虫任务
|
||
参数示例：
|
||
data : {
|
||
"taskId": 348,
|
||
"strategy": "Terms",
|
||
"files": [
|
||
{
|
||
"fileIndex": 1,
|
||
"fileUrl": "ab8cce4878754bfd89b4cd3ed20e1395",
|
||
"originalFilename": "品牌样例.xlsx",
|
||
"relativePath": "店铺A/品牌样例.xlsx",
|
||
"sheetName": "Sheet1",
|
||
"columns": [
|
||
"品牌",
|
||
"ASIH",
|
||
"状态",
|
||
"时间"
|
||
],
|
||
"rows": [
|
||
{
|
||
"品牌": "YQAUTEC",
|
||
"ASIH": "B0F28NZ752",
|
||
"状态": "",
|
||
"时间": "",
|
||
"__rowIndex": 2
|
||
}
|
||
],
|
||
"uniqueBrands": [
|
||
"YQAUTEC"
|
||
]
|
||
}
|
||
]
|
||
}
|
||
"""
|
||
try:
|
||
file_ls = data.get("files")
|
||
strategy = data.get("strategy")
|
||
if not isinstance(file_ls, list):
|
||
file_ls = []
|
||
tasks = []
|
||
for file_data in file_ls:
|
||
rows = file_data.get("rows") or []
|
||
brand_ls = [i.get("品牌") for i in rows if i.get("品牌")]
|
||
fileUrl = file_data.get("fileUrl")
|
||
if not brand_ls:
|
||
continue
|
||
file_chunks = [brand_ls[i:i + 5] for i in range(0, len(brand_ls), 5)]
|
||
chunk_total = len(file_chunks)
|
||
total_lines = len(brand_ls)
|
||
for chunk_index, chunk in enumerate(file_chunks, start=1):
|
||
tasks.append((chunk, fileUrl, total_lines, chunk_total, chunk_index))
|
||
|
||
if tasks:
|
||
max_workers = min(8, len(tasks))
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
if _is_task_cancelled(task_id):
|
||
_push_task_event(task_id, {'status': 'cancelled'})
|
||
return
|
||
futures = [
|
||
executor.submit(_run_brand_single, task_id, chunk, file_url, strategy, total_lines, chunk_total, chunk_index)
|
||
for chunk, file_url, total_lines, chunk_total, chunk_index in tasks
|
||
]
|
||
for future in as_completed(futures):
|
||
if _is_task_cancelled(task_id):
|
||
for f in futures:
|
||
f.cancel()
|
||
_push_task_event(task_id, {'status': 'cancelled'})
|
||
return
|
||
future.result()
|
||
_push_task_event(task_id, {'status': 'success'})
|
||
except Exception as e:
|
||
print("执行出错",traceback.format_exc())
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': str(e)})
|
||
print("执行完成")
|
||
|
||
|
||
def _parse_json(val, default=None):
|
||
if val is None:
|
||
return default if default is not None else []
|
||
if isinstance(val, (list, dict)):
|
||
return val
|
||
try:
|
||
return json.loads(val)
|
||
except Exception:
|
||
return default if default is not None else []
|
||
|
||
|
||
def _is_url(s):
|
||
if not isinstance(s, str) or not s.strip():
|
||
return False
|
||
return s.strip().startswith('http://') or s.strip().startswith('https://')
|
||
|
||
|
||
def _normalize_result_paths(raw):
|
||
"""
|
||
将 result_paths 规范化为 (url_list, zip_url)。
|
||
支持旧格式 list 或新格式 dict {"urls": [...], "zip_url": "..."}。
|
||
"""
|
||
if not raw:
|
||
return [], None
|
||
if isinstance(raw, dict):
|
||
urls = raw.get('urls') or []
|
||
if not isinstance(urls, list):
|
||
urls = []
|
||
zip_url = raw.get('zip_url')
|
||
if zip_url and not isinstance(zip_url, str):
|
||
zip_url = None
|
||
return urls, zip_url
|
||
if isinstance(raw, list):
|
||
return raw, None
|
||
return [], None
|
||
|
||
|
||
@brand_bp.route('/api/brand/expand-folder', methods=['POST'])
|
||
@login_required
|
||
def api_brand_expand_folder():
|
||
"""展开文件夹，返回其下所有 .xlsx 文件路径列表"""
|
||
try:
|
||
_, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
data = request.get_json() or {}
|
||
folder = (data.get('folder') or '').strip()
|
||
if not folder:
|
||
return jsonify({'success': False, 'error': '请提供 folder 路径'}), 400
|
||
paths = _expand_folder_xlsx(folder)
|
||
return jsonify({'success': True, 'paths': paths})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
def _expand_folder_xlsx_recursive_items(folder_path):
|
||
"""返回文件夹下所有 .xlsx 文件的绝对路径和相对路径列表（递归）"""
|
||
if not folder_path or not os.path.isdir(folder_path):
|
||
return []
|
||
items = []
|
||
root = os.path.normpath(folder_path)
|
||
for current_root, _, filenames in os.walk(root):
|
||
for name in filenames:
|
||
if not (name.endswith('.xlsx') or name.endswith('.XLSX')):
|
||
continue
|
||
absolute_path = os.path.normpath(os.path.join(current_root, name))
|
||
relative_path = os.path.relpath(absolute_path, root).replace('\\', '/')
|
||
items.append({
|
||
'absolutePath': absolute_path,
|
||
'relativePath': relative_path,
|
||
})
|
||
items.sort(key=lambda item: item['relativePath'])
|
||
return items
|
||
|
||
@brand_bp.route('/api/brand/expand-folder-recursive', methods=['POST'])
|
||
@login_required
|
||
def api_brand_expand_folder_recursive():
|
||
"""递归展开文件夹，返回 .xlsx 文件绝对路径及相对路径列表"""
|
||
try:
|
||
_, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
data = request.get_json() or {}
|
||
folder = (data.get('folder') or '').strip()
|
||
if not folder:
|
||
return jsonify({'success': False, 'error': '请提供 folder 路径'}), 400
|
||
items = _expand_folder_xlsx_recursive_items(folder)
|
||
return jsonify({'success': True, 'items': items})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/run', methods=['POST'])
|
||
# @login_required
|
||
def api_brand_run():
|
||
"""立即运行：创建任务并后台执行，立即返回 task_id，前端可轮询进度与取消"""
|
||
try:
|
||
data = request.get_json() or {}
|
||
paths = data.get('paths') or []
|
||
# task_type: 1=立即执行，2=添加任务；此接口默认 1
|
||
try:
|
||
task_type = int(data.get('task_type') or 1)
|
||
except Exception:
|
||
task_type = 1
|
||
if task_type not in (1, 2):
|
||
task_type = 1
|
||
strategy = (data.get('strategy') or 'Terms').strip()
|
||
if strategy not in ('Terms', 'Simple'):
|
||
strategy = 'Terms'
|
||
if not isinstance(paths, list):
|
||
paths = []
|
||
paths = [p.strip() for p in paths if p and isinstance(p, str) and os.path.isfile(p.strip())]
|
||
if not paths:
|
||
return jsonify({'success': False, 'error': '没有有效的 xlsx 文件路径'}), 400
|
||
# 先上传到 OSS，获取链接再入库
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
urls = []
|
||
for p in paths:
|
||
url, err = _upload_local_xlsx_to_oss(p, user_id)
|
||
if err:
|
||
return jsonify({'success': False, 'error': f'上传文件失败: {os.path.basename(p)} - {err}'}), 500
|
||
base_name = os.path.basename(p)
|
||
urls.append({"fileUrl": url,"originalFilename": base_name,"relativePath": p })
|
||
|
||
# 请求提交
|
||
resp = requests.post(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id}",headers={
|
||
"content-type":"application/json",
|
||
},data=json.dumps({ "files": urls, "strategy": strategy,"taskType": 1,"archiveName": ""}))
|
||
print({ "files": urls, "strategy": strategy,"taskType": 1,"archiveName": ""})
|
||
print(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id}")
|
||
resp_data = resp.json()
|
||
# print(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id} 返回：",resp_data,"状态：",resp.status_code)
|
||
if resp_data.get("success"):
|
||
task_id = resp_data.get("data").get("taskId")
|
||
data = resp_data.get("data")
|
||
threading.Thread(target=_background_brand_task, args=(task_id,data), daemon=True).start()
|
||
return jsonify({'success': True, 'task_id': task_id})
|
||
else:
|
||
print(resp_data)
|
||
return jsonify({'success': False, 'error': "请求接口失败"}), 500
|
||
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks', methods=['GET', 'POST'])
|
||
# @login_required
|
||
def api_brand_tasks():
|
||
"""GET: 获取当前用户的任务列表；POST: 添加任务（后台执行）"""
|
||
if request.method == 'GET':
|
||
try:
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
resp = requests.get(f"{JAVA_API_BASE}/api/brand/tasks", params={
|
||
"userId": user_id
|
||
})
|
||
resp_data = resp.json()
|
||
print("获取列表",resp.text)
|
||
if resp_data.get("success"):
|
||
items = resp_data.get("data").get("items")
|
||
return jsonify({'success': True, 'items': items})
|
||
return jsonify({'success': True, 'items': []})
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
print("获取列表失败",e)
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
# POST
|
||
try:
|
||
data = request.get_json() or {}
|
||
paths = data.get('paths') or []
|
||
# task_type: 1=立即执行，2=添加任务；此接口默认 2
|
||
try:
|
||
task_type = int(data.get('task_type') or 2)
|
||
except Exception:
|
||
task_type = 2
|
||
if task_type not in (1, 2):
|
||
task_type = 2
|
||
strategy = (data.get('strategy') or 'Terms').strip()
|
||
if strategy not in ('Terms', 'Simple'):
|
||
strategy = 'Terms'
|
||
if not isinstance(paths, list):
|
||
paths = []
|
||
paths = [p.strip() for p in paths if p and isinstance(p, str)]
|
||
if not paths:
|
||
return jsonify({'success': False, 'error': '请提供至少一个文件路径或链接'}), 400
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
urls = []
|
||
for p in paths:
|
||
url, err = _upload_local_xlsx_to_oss(p, user_id)
|
||
if err:
|
||
return jsonify({'success': False, 'error': f'上传文件失败: {os.path.basename(p)} - {err}'}), 500
|
||
base_name = os.path.basename(p)
|
||
urls.append({"fileUrl": url, "originalFilename": base_name, "relativePath": p})
|
||
params = {
|
||
"userId": user_id
|
||
}
|
||
# resp = requests.post(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id}",headers={
|
||
# "content-type":"application/json",
|
||
# },data=json.dumps({ "files": urls, "strategy": strategy,"taskType": 1,"archiveName": ""}))
|
||
req_data = {"files": urls, "strategy": strategy, "taskType": 2, "archiveName": ""}
|
||
resp = requests.post(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id}", headers={"content-type": "application/json"},
|
||
data=json.dumps(req_data))
|
||
print(req_data)
|
||
# print(f"{JAVA_API_BASE}/api/brand/tasks?userId={user_id}，返回", resp.text)
|
||
resp_data = resp.json()
|
||
if resp_data.get("success"):
|
||
task_id = resp_data.get("data").get("taskId")
|
||
data = resp_data.get("data")
|
||
return jsonify({'success': True, 'task_id': task_id})
|
||
return jsonify({'success': False, 'task_id': "请求异常"})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks/<int:task_id>')
|
||
# @login_required
|
||
def api_brand_task_detail(task_id):
|
||
"""获取单个任务详情（用于轮询状态）"""
|
||
try:
|
||
resp = requests.get(f"{JAVA_API_BASE}/api/brand/tasks/{task_id}")
|
||
resp_data = resp.json()
|
||
if resp_data.get("success"):
|
||
task = resp_data["data"]["task"]
|
||
return jsonify({
|
||
'success': True,
|
||
'task': task
|
||
})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks/<int:task_id>/line-progress')
|
||
# @login_required
|
||
def api_brand_task_line_progress(task_id):
|
||
"""
|
||
获取任务的“当前文件行级进度”（当前行数/总行数）。
|
||
该信息仅存放在内存字典 _task_line_progress 中，不写入数据库。
|
||
"""
|
||
try:
|
||
resp = requests.get(f"{JAVA_API_BASE}/api/brand/tasks/{task_id}")
|
||
resp_data = resp.json()
|
||
if resp_data.get("success"):
|
||
if resp_data["data"]["line_progress"]["has_progress"]:
|
||
info = resp_data["data"]["line_progress"]["info"]
|
||
resp = {
|
||
'file_index': int(info.get('file_index') or 0),
|
||
'file_total': int(info.get('file_total') or 0),
|
||
'file_name': info.get("file_name",""),
|
||
'current_line': int(info.get('current_line') or 0),
|
||
'total_lines': int(info.get('total_lines') or 0),
|
||
}
|
||
return jsonify({'success': True, 'has_progress': True, 'info': resp})
|
||
raise RuntimeError("获取进度失败")
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks/<int:task_id>/events')
|
||
# @login_required
|
||
def api_brand_task_events(task_id):
|
||
"""SSE：任务完成时后端主动推送事件，前端监听后隐藏进度条与取消按钮"""
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
|
||
def _task_status_event():
|
||
conn = get_db()
|
||
try:
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT status, error_message FROM brand_crawl_tasks WHERE id = %s AND user_id = %s",
|
||
(task_id, user_id)
|
||
)
|
||
row = cur.fetchone()
|
||
finally:
|
||
conn.close()
|
||
if not row:
|
||
return None
|
||
st = (row.get('status') or '').lower()
|
||
if st in ('success', 'failed', 'cancelled'):
|
||
return {'status': st, 'error_message': (row.get('error_message') or '').strip() or None}
|
||
return None
|
||
|
||
# 若任务已处于终态，直接返回一条事件后结束
|
||
ev = _task_status_event()
|
||
if ev is not None:
|
||
def _one_shot():
|
||
yield "data: " + json.dumps(ev, ensure_ascii=False) + "\n\n"
|
||
return Response(
|
||
_one_shot(),
|
||
mimetype='text/event-stream',
|
||
headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
|
||
)
|
||
|
||
# 否则注册队列，等待后台任务完成时推送
|
||
q = Queue()
|
||
with _task_event_lock:
|
||
_task_event_queues.setdefault(task_id, []).append(q)
|
||
|
||
def _stream():
|
||
try:
|
||
while True:
|
||
try:
|
||
event = q.get(timeout=20)
|
||
yield "data: " + json.dumps(event, ensure_ascii=False) + "\n\n"
|
||
return
|
||
except Empty:
|
||
yield ": keepalive\n\n"
|
||
finally:
|
||
with _task_event_lock:
|
||
lst = _task_event_queues.get(task_id, [])
|
||
if q in lst:
|
||
lst.remove(q)
|
||
if not lst:
|
||
_task_event_queues.pop(task_id, None)
|
||
|
||
return Response(
|
||
_stream(),
|
||
mimetype='text/event-stream',
|
||
headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
|
||
)
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks/<int:task_id>/cancel', methods=['POST'])
|
||
# @login_required
|
||
def api_brand_task_cancel(task_id):
|
||
"""取消正在执行或等待中的任务（pending/running 均可取消）"""
|
||
try:
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
|
||
|
||
resp = requests.post(
|
||
f"{JAVA_API_BASE}/api/brand/tasks/{task_id}/cancel",
|
||
headers={"accept": "application/json"},
|
||
timeout=20
|
||
)
|
||
try:
|
||
resp_data = resp.json()
|
||
except Exception:
|
||
return jsonify({'success': False, 'error': '取消接口返回非 JSON'}), 500
|
||
if not resp_data.get('success'):
|
||
return jsonify({'success': False, 'error': '第三方取消任务失败'}), 400
|
||
|
||
_mark_task_cancelled(task_id)
|
||
_push_task_event(task_id, {'status': 'cancelled'})
|
||
return jsonify({'success': True})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/tasks/<int:task_id>', methods=['DELETE'])
|
||
@login_required
|
||
def api_brand_task_delete(task_id):
|
||
"""删除任务（仅限非 running 状态）"""
|
||
try:
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
|
||
resp = requests.delete(
|
||
f"{JAVA_API_BASE}/api/brand/tasks/{task_id}",
|
||
headers={"accept": "application/json"},
|
||
timeout=20
|
||
)
|
||
try:
|
||
resp_data = resp.json()
|
||
except Exception:
|
||
return jsonify({'success': False, 'error': '取消接口返回非 JSON'}), 500
|
||
if not resp_data.get('success'):
|
||
return jsonify({'success': False, 'error': '第三方取消任务失败'}), 400
|
||
return jsonify({'success': True})
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|
||
@brand_bp.route('/api/brand/download/<int:task_id>')
|
||
@login_required
|
||
def api_brand_download(task_id):
|
||
"""下载任务结果：优先使用 OSS zip_url 重定向；否则本地/单链接/多链接按原逻辑处理"""
|
||
try:
|
||
user_id, err = _get_user_id_or_error()
|
||
if err:
|
||
return err
|
||
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT result_paths FROM brand_crawl_tasks WHERE id = %s AND user_id = %s",
|
||
(task_id, user_id)
|
||
)
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row or not row.get('result_paths'):
|
||
return jsonify({'success': False, 'error': '无结果可下载'}), 404
|
||
raw = row['result_paths']
|
||
if isinstance(raw, str):
|
||
raw = _parse_json(raw)
|
||
url_list, zip_url = _normalize_result_paths(raw)
|
||
# 新格式：有 zip_url 直接重定向到 OSS
|
||
if zip_url and _is_url(zip_url):
|
||
return redirect(zip_url, code=302)
|
||
|
||
# 兼容旧格式：result_paths 可能为 list（本地路径或 url）
|
||
if not url_list and isinstance(raw, list):
|
||
url_list = raw
|
||
local_paths = [p for p in url_list if isinstance(p, str) and not _is_url(p) and os.path.isfile(p)]
|
||
url_paths = [p.strip() for p in url_list if isinstance(p, str) and _is_url(p)]
|
||
|
||
if len(url_paths) == 1 and not local_paths:
|
||
return redirect(url_paths[0], code=302)
|
||
|
||
if len(local_paths) == 1 and not url_paths:
|
||
return send_file(
|
||
local_paths[0],
|
||
as_attachment=True,
|
||
download_name=os.path.basename(local_paths[0])
|
||
)
|
||
|
||
if local_paths or url_paths:
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, 'w', zipfile.ZIP_DEFLATED) as zf:
|
||
for p in local_paths:
|
||
zf.write(p, os.path.basename(p))
|
||
for url in url_paths:
|
||
try:
|
||
import requests as req
|
||
r = req.get(url, timeout=30, stream=True)
|
||
r.raise_for_status()
|
||
name = os.path.basename(urlparse(url).path) or ('file_%s' % (url_paths.index(url)))
|
||
if not name or name == 'file_%s' % url_paths.index(url):
|
||
name = 'download_%s' % url_paths.index(url)
|
||
zf.writestr(name, r.content)
|
||
except Exception:
|
||
pass
|
||
buf.seek(0)
|
||
return send_file(
|
||
buf,
|
||
mimetype='application/zip',
|
||
as_attachment=True,
|
||
download_name='brand_task_%s.zip' % task_id
|
||
)
|
||
|
||
return jsonify({'success': False, 'error': '结果文件不存在或链接不可用'}), 404
|
||
except Exception as e:
|
||
return jsonify({'success': False, 'error': str(e)}), 500
|
||
|
||
|