856 lines
35 KiB
Python
856 lines
35 KiB
Python
"""
|
||
品牌爬虫蓝图：展开文件夹、运行任务、任务列表/详情、下载结果
|
||
"""
|
||
import os
|
||
import json
|
||
import threading
|
||
import datetime
|
||
import traceback
|
||
import zipfile
|
||
import io
|
||
import time
|
||
import tempfile
|
||
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
|
||
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()
|
||
|
||
|
||
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)
|
||
|
||
|
||
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/{int(time.time() * 1000)}/{user_id}_{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_files(file_paths, task_id=None, check_cancelled=None, strategy="Terms"):
|
||
"""对多个 xlsx 执行 single_file_handle（线程池多线程），返回 (result_paths, error_message)。
|
||
check_cancelled: 可调用对象，返回 True 时中止并返回 (result_paths, 'cancelled')。
|
||
task_id 存在时会在每处理完一个文件后更新 progress_current。
|
||
strategy: 品牌匹配方式，默认 Terms，可选 Simple。"""
|
||
os.makedirs(BRAND_OUTPUT_DIR, exist_ok=True)
|
||
subdir = str(task_id) if task_id else ('immediate_' + datetime.datetime.now().strftime('%Y%m%d_%H%M%S'))
|
||
out_dir = os.path.join(BRAND_OUTPUT_DIR, subdir)
|
||
os.makedirs(out_dir, exist_ok=True)
|
||
|
||
# 先筛出有效路径并生成 (fp, out_path) 列表
|
||
tasks = []
|
||
for fp in file_paths:
|
||
if not fp or not isinstance(fp, str) or not os.path.isfile(fp):
|
||
continue
|
||
base = os.path.splitext(os.path.basename(fp))[0]
|
||
safe_base = "".join(c if c.isalnum() or c in '-_' else '_' for c in base)
|
||
out_path = os.path.join(out_dir, safe_base + '_result.xlsx')
|
||
tasks.append((fp, out_path))
|
||
|
||
if not tasks:
|
||
return [], None
|
||
|
||
result_paths = []
|
||
result_lock = threading.Lock()
|
||
current = [0] # 用列表以便在闭包中修改
|
||
|
||
def process_one(fp, out_path, file_index, files_total):
|
||
# 行级进度回调：由 single_file_handle 在循环中调用，直接写入内存字典，不落库
|
||
def progress_callback(current_line, total_lines, extra=None):
|
||
if not task_id:
|
||
return
|
||
info = {
|
||
'file_index': file_index,
|
||
'file_total': files_total,
|
||
'file_path': fp,
|
||
'current_line': int(current_line or 0),
|
||
'total_lines': int(total_lines or 0),
|
||
}
|
||
if isinstance(extra, dict):
|
||
info.update(extra)
|
||
with _task_line_progress_lock:
|
||
_task_line_progress[task_id] = info
|
||
|
||
single_file_handle(fp, out_path, strategy, check_cancelled=check_cancelled, progress_callback=progress_callback)
|
||
return out_path
|
||
|
||
max_workers = min(8, len(tasks))
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
future_to_out = {
|
||
executor.submit(process_one, fp, out_path, idx + 1, len(tasks)): out_path
|
||
for idx, (fp, out_path) in enumerate(tasks)
|
||
}
|
||
try:
|
||
for future in as_completed(future_to_out):
|
||
if check_cancelled and check_cancelled():
|
||
for f in future_to_out:
|
||
f.cancel()
|
||
return result_paths, 'cancelled'
|
||
try:
|
||
out_path = future.result()
|
||
with result_lock:
|
||
result_paths.append(out_path)
|
||
current[0] += 1
|
||
if task_id:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET progress_current = %s WHERE id = %s",
|
||
(current[0], task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
print(traceback.format_exc())
|
||
except TaskCancelledError:
|
||
for f in future_to_out:
|
||
f.cancel()
|
||
return result_paths, 'cancelled'
|
||
except Exception as e:
|
||
print(traceback.format_exc())
|
||
print(e)
|
||
for f in future_to_out:
|
||
f.cancel()
|
||
return result_paths, str(e)
|
||
except Exception as e:
|
||
print(traceback.format_exc())
|
||
return result_paths, str(e)
|
||
return result_paths, None
|
||
|
||
|
||
def _upload_results_to_oss(source_paths, local_result_paths, task_id):
|
||
"""
|
||
将源文件和结果文件上传到 OSS，并打包成 zip（分两个文件夹：源文件、结果文件）上传。
|
||
返回 (result_data, error_message)，result_data 为 {"urls": [url1, ...], "zip_url": "..."}，出错时为 (None, err)。
|
||
"""
|
||
if not local_result_paths or not task_id:
|
||
return None, "无结果文件"
|
||
try:
|
||
from ali_oss import upload_file as oss_upload_file
|
||
except ImportError:
|
||
return None, "OSS 模块未配置"
|
||
urls = []
|
||
subdir = f"{OSS_PREFIX}/{task_id}"
|
||
for fp in local_result_paths:
|
||
if not fp or not os.path.isfile(fp):
|
||
continue
|
||
name = os.path.basename(fp)
|
||
key = f"{bucket_path}{subdir}/{name}"
|
||
try:
|
||
with open(fp, "rb") as f:
|
||
content = f.read()
|
||
url = oss_upload_file(content, key)
|
||
urls.append(url)
|
||
except Exception as e:
|
||
return None, f"上传结果文件失败: {e}"
|
||
if not urls:
|
||
return None, "没有可上传的结果文件"
|
||
# 打包成 zip：源文件 -> 源文件/，结果文件 -> 结果文件/
|
||
buf = io.BytesIO()
|
||
with zipfile.ZipFile(buf, "w", zipfile.ZIP_DEFLATED) as zf:
|
||
for fp in (source_paths or []):
|
||
if fp and os.path.isfile(fp):
|
||
zf.write(fp, os.path.join("源文件", os.path.basename(fp)))
|
||
for fp in local_result_paths:
|
||
if fp and os.path.isfile(fp):
|
||
zf.write(fp, os.path.join("结果文件", os.path.basename(fp)))
|
||
buf.seek(0)
|
||
zip_key = f"{bucket_path}{subdir}/result.zip"
|
||
try:
|
||
zip_url = oss_upload_file(buf.getvalue(), zip_key)
|
||
except Exception as e:
|
||
return {"urls": urls, "zip_url": None}, None # 单文件链接已保存，zip 失败不报错
|
||
return {"urls": urls, "zip_url": zip_url}, None
|
||
|
||
|
||
def _background_brand_task(task_id):
|
||
"""后台执行单个品牌爬虫任务并更新数据库"""
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"SELECT id, user_id, file_paths, strategy, `desc` FROM brand_crawl_tasks WHERE id = %s AND status = 'pending'",
|
||
(task_id,)
|
||
)
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row:
|
||
return
|
||
file_paths = row.get('file_paths')
|
||
if isinstance(file_paths, str):
|
||
file_paths = json.loads(file_paths) if file_paths else []
|
||
# 优先使用表里存的 strategy 字段，兼容旧数据则从 desc 解析
|
||
strategy = (row.get('strategy') or '').strip() or None
|
||
if not strategy:
|
||
desc_val = (row.get('desc') or '').strip()
|
||
strategy = "Simple" if desc_val.startswith('[Simple]') else "Terms"
|
||
if strategy not in ('Terms', 'Simple'):
|
||
strategy = 'Terms'
|
||
if not file_paths:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'failed', error_message = %s WHERE id = %s",
|
||
('没有有效文件路径', task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': '没有有效文件路径'})
|
||
return
|
||
# 将 file_paths 解析为本地路径：URL 下载到 BRAND_OUTPUT_DIR 对应任务文件夹下（兼容旧数据中的本地路径）
|
||
os.makedirs(BRAND_OUTPUT_DIR, exist_ok=True)
|
||
task_subdir = str(task_id) if task_id else ('immediate_' + datetime.datetime.now().strftime('%Y%m%d_%H%M%S'))
|
||
task_dir = os.path.join(BRAND_OUTPUT_DIR, task_subdir)
|
||
os.makedirs(task_dir, exist_ok=True)
|
||
|
||
local_paths = []
|
||
downloaded_files_to_cleanup_on_fail = []
|
||
for idx, p in enumerate(file_paths):
|
||
if not p or not isinstance(p, str):
|
||
continue
|
||
p = p.strip()
|
||
if _is_url(p):
|
||
try:
|
||
parsed = urlparse(p)
|
||
url_name = os.path.basename(parsed.path or '')
|
||
if not url_name:
|
||
url_name = f"input_{idx + 1}.xlsx"
|
||
safe_name = "".join(c if c.isalnum() or c in '-_.' else '_' for c in url_name)
|
||
if not safe_name.lower().endswith('.xlsx'):
|
||
safe_name = safe_name + '.xlsx'
|
||
|
||
base_no_ext, ext = os.path.splitext(safe_name)
|
||
dest = os.path.join(task_dir, safe_name)
|
||
n = 1
|
||
while os.path.exists(dest):
|
||
dest = os.path.join(task_dir, f"{base_no_ext}_{n}{ext}")
|
||
n += 1
|
||
|
||
r = requests.get(p, timeout=120, stream=True)
|
||
r.raise_for_status()
|
||
with open(dest, "wb") as f:
|
||
for chunk in r.iter_content(chunk_size=1024 * 1024):
|
||
if chunk:
|
||
f.write(chunk)
|
||
local_paths.append(dest)
|
||
downloaded_files_to_cleanup_on_fail.append(dest)
|
||
except Exception as e:
|
||
for tf in downloaded_files_to_cleanup_on_fail:
|
||
try:
|
||
os.unlink(tf)
|
||
except Exception:
|
||
pass
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'failed', error_message = %s WHERE id = %s",
|
||
(f'下载文件失败: {str(e)}', task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': f'下载文件失败: {str(e)}'})
|
||
return
|
||
elif os.path.isfile(p):
|
||
local_paths.append(p)
|
||
if not local_paths:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'failed', error_message = %s WHERE id = %s",
|
||
('没有有效的本地或可下载文件', task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': '没有有效的本地或可下载文件'})
|
||
return
|
||
try:
|
||
# 原子地将 pending 改为 running，避免用户先点取消时仍被覆盖为 running
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'running', progress_total = %s, progress_current = 0 WHERE id = %s AND status = 'pending'",
|
||
(len(local_paths), task_id)
|
||
)
|
||
n = cur.rowcount
|
||
conn.commit()
|
||
conn.close()
|
||
if n == 0:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT status FROM brand_crawl_tasks WHERE id = %s", (task_id,))
|
||
r = cur.fetchone()
|
||
conn.close()
|
||
if r and (r.get('status') or '') == 'cancelled':
|
||
_push_task_event(task_id, {'status': 'cancelled'})
|
||
return
|
||
except Exception:
|
||
pass
|
||
|
||
def check_cancelled():
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute("SELECT status FROM brand_crawl_tasks WHERE id = %s", (task_id,))
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
return row and (row.get('status') or '') == 'cancelled'
|
||
except Exception:
|
||
return False
|
||
|
||
result_paths, err = _run_brand_files(local_paths, task_id=task_id, check_cancelled=check_cancelled, strategy=strategy)
|
||
if err:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
st = 'cancelled' if err == 'cancelled' else 'failed'
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = %s, error_message = %s WHERE id = %s",
|
||
(st, err if st == 'failed' else None, task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception as e:
|
||
traceback.print_exc()
|
||
print(e)
|
||
pass
|
||
_push_task_event(task_id, {'status': st, 'error_message': err if st == 'failed' else None})
|
||
return
|
||
result_data, upload_err = _upload_results_to_oss(local_paths, result_paths, task_id)
|
||
if upload_err:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'failed', error_message = %s WHERE id = %s",
|
||
(upload_err, task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': upload_err})
|
||
return
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'success', result_paths = %s WHERE id = %s",
|
||
(json.dumps(result_data), task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'success'})
|
||
finally:
|
||
pass
|
||
except Exception as e:
|
||
try:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'failed', error_message = %s WHERE id = %s",
|
||
(str(e), task_id)
|
||
)
|
||
conn.commit()
|
||
conn.close()
|
||
except Exception:
|
||
pass
|
||
_push_task_event(task_id, {'status': 'failed', 'error_message': str(e)})
|
||
|
||
|
||
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:
|
||
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
|
||
|
||
|
||
@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 = session['user_id']
|
||
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
|
||
urls.append(url)
|
||
desc_core = ', '.join(os.path.basename(p) for p in paths)[:500] if paths else ''
|
||
desc = f'[{strategy}] ' + (desc_core or '')
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"INSERT INTO brand_crawl_tasks (user_id, file_paths, `desc`, strategy, status, task_type) VALUES (%s, %s, %s, %s, 'pending', %s)",
|
||
(user_id, json.dumps(urls), desc or None, strategy, task_type)
|
||
)
|
||
task_id = cur.lastrowid
|
||
conn.commit()
|
||
conn.close()
|
||
threading.Thread(target=_background_brand_task, args=(task_id,), daemon=True).start()
|
||
return jsonify({'success': True, 'task_id': task_id})
|
||
except Exception as e:
|
||
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:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"""SELECT id, file_paths, `desc`, strategy, status, result_paths, error_message,
|
||
COALESCE(progress_current, 0) AS progress_current,
|
||
COALESCE(progress_total, 0) AS progress_total,
|
||
created_at, updated_at
|
||
FROM brand_crawl_tasks WHERE user_id = %s ORDER BY id DESC LIMIT 100""",
|
||
(session['user_id'],)
|
||
)
|
||
rows = cur.fetchall()
|
||
conn.close()
|
||
items = []
|
||
for r in rows:
|
||
items.append({
|
||
'id': r['id'],
|
||
'file_paths': _parse_json(r['file_paths'], []),
|
||
'desc': (r.get('desc') or '').strip() or None,
|
||
'strategy': (r.get('strategy') or 'Terms').strip() or 'Terms',
|
||
'status': r.get('status') or 'pending',
|
||
'result_paths': _parse_json(r['result_paths']),
|
||
'error_message': (r.get('error_message') or '').strip() or None,
|
||
'progress_current': int(r.get('progress_current') or 0),
|
||
'progress_total': int(r.get('progress_total') or 0),
|
||
'created_at': r['created_at'].strftime('%Y-%m-%d %H:%M') if r.get('created_at') else '',
|
||
'updated_at': r['updated_at'].strftime('%Y-%m-%d %H:%M') if r.get('updated_at') else '',
|
||
})
|
||
return jsonify({'success': True, 'items': items})
|
||
except Exception as 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 = session['user_id']
|
||
urls = []
|
||
for p in paths:
|
||
if _is_url(p):
|
||
urls.append(p)
|
||
elif os.path.isfile(p):
|
||
url, err = _upload_local_xlsx_to_oss(p, user_id)
|
||
if err:
|
||
return jsonify({'success': False, 'error': f'上传文件失败: {os.path.basename(p)} - {err}'}), 500
|
||
urls.append(url)
|
||
else:
|
||
return jsonify({'success': False, 'error': f'无效路径或链接: {p[:80]}'}), 400
|
||
desc_core = ', '.join(os.path.basename(urlparse(u).path or u) for u in urls)[:500] if urls else ''
|
||
desc = f'[{strategy}] ' + (desc_core or '')
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"INSERT INTO brand_crawl_tasks (user_id, file_paths, `desc`, strategy, status, task_type) VALUES (%s, %s, %s, %s, 'pending', %s)",
|
||
(user_id, json.dumps(urls), desc or None, strategy, task_type)
|
||
)
|
||
task_id = cur.lastrowid
|
||
conn.commit()
|
||
conn.close()
|
||
# threading.Thread(target=_background_brand_task, args=(task_id,), daemon=True).start()
|
||
return jsonify({'success': True, 'task_id': 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:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"""SELECT id, file_paths, `desc`, strategy, status, result_paths, error_message,
|
||
COALESCE(progress_current, 0) AS progress_current,
|
||
COALESCE(progress_total, 0) AS progress_total,
|
||
created_at, updated_at
|
||
FROM brand_crawl_tasks WHERE id = %s AND user_id = %s""",
|
||
(task_id, session['user_id'])
|
||
)
|
||
row = cur.fetchone()
|
||
conn.close()
|
||
if not row:
|
||
return jsonify({'success': False, 'error': '任务不存在'}), 404
|
||
return jsonify({
|
||
'success': True,
|
||
'task': {
|
||
'id': row['id'],
|
||
'file_paths': _parse_json(row['file_paths'], []),
|
||
'desc': (row.get('desc') or '').strip() or None,
|
||
'strategy': (row.get('strategy') or 'Terms').strip() or 'Terms',
|
||
'status': row.get('status') or 'pending',
|
||
'result_paths': _parse_json(row['result_paths']),
|
||
'error_message': (row.get('error_message') or '').strip() or None,
|
||
'progress_current': int(row.get('progress_current') or 0),
|
||
'progress_total': int(row.get('progress_total') or 0),
|
||
'created_at': row['created_at'].strftime('%Y-%m-%d %H:%M') if row.get('created_at') else '',
|
||
'updated_at': row['updated_at'].strftime('%Y-%m-%d %H:%M') if row.get('updated_at') else '',
|
||
}
|
||
})
|
||
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:
|
||
with _task_line_progress_lock:
|
||
info = _task_line_progress.get(task_id)
|
||
if not info:
|
||
return jsonify({'success': True, 'has_progress': False})
|
||
# 为了避免暴露完整路径，只返回文件名
|
||
file_name = os.path.basename(info.get('file_path') or '') if info.get('file_path') else ''
|
||
resp = {
|
||
'file_index': int(info.get('file_index') or 0),
|
||
'file_total': int(info.get('file_total') or 0),
|
||
'file_name': 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})
|
||
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：任务完成时后端主动推送事件，前端监听后隐藏进度条与取消按钮"""
|
||
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, session['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:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"UPDATE brand_crawl_tasks SET status = 'cancelled' WHERE id = %s AND user_id = %s AND status IN ('pending', 'running')",
|
||
(task_id, session['user_id'])
|
||
)
|
||
n = cur.rowcount
|
||
conn.commit()
|
||
conn.close()
|
||
if n == 0:
|
||
return jsonify({'success': False, 'error': '任务不存在或无法取消'}), 400
|
||
_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:
|
||
conn = get_db()
|
||
with conn.cursor() as cur:
|
||
cur.execute(
|
||
"DELETE FROM brand_crawl_tasks WHERE id = %s AND user_id = %s AND status != 'running'",
|
||
(task_id, session['user_id'])
|
||
)
|
||
n = cur.rowcount
|
||
conn.commit()
|
||
conn.close()
|
||
if n == 0:
|
||
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:
|
||
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, session['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
|
||
|
||
|