import logging import os import queue import threading import time from datetime import datetime import random import pymysql from tqdm import tqdm from save_page_as_pdf import PDFSaver from save_remote_as_mhtml import RemoteMHTMLSaver from save_page_as_mhtml import MHTMLSaver import tldextract # 配置日志 from save_remote_as_pdf import RemotePDFSaver logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', handlers=[ logging.StreamHandler(), logging.FileHandler('pdf_downloader.log') ] ) logger = logging.getLogger(__name__) # =============== MySQL 配置 =============== MYSQL_CONFIG = { 'host': '47.113.231.200', 'port': 28089, 'user': 'root', 'password': 'passok123A', 'database': 'dsp', 'charset': 'utf8mb4', 'autocommit': False # 手动控制事务 } # ========================================= # 配置参数 BATCH_SIZE = 500 MAX_WORKERS = 1 TIMEOUT = 10 PDF_OUTPUT_DIR = 'D:/data/output/pdf' MIN_PDF_SIZE = 5 * 1024 # 80KB MHTML_OUTPUT_DIR = 'D:/data/output/mhtml' os.makedirs(PDF_OUTPUT_DIR, exist_ok=True) running = True running_interval_seconds = 10 skip_host_name = [ 'epochtimes.com', 'secretchina.com', # 'rodong.rep.kp', # 'kcna.kp' ] class PDFDownloader: def __init__(self): self.db_lock = threading.Lock() self.db_connection = None self.task_queue = queue.Queue(maxsize=MAX_WORKERS * 3) self.processed_count = 0 self.success_count = 0 self.fail_count = 0 self.small_file_count = 0 # 新增：统计小文件数量 self.last_loadtime = self.get_last_loadtime() self.total_rows = self.get_total_rows() self.start_time = time.time() self.skip_hosts = [] self.local_handler = None self.remote_handler = None # 替换 MYSQL_CONFIG 中的连接方式 def get_db_connection(self): self.db_connection = pymysql.connect( host=MYSQL_CONFIG['host'], port=MYSQL_CONFIG['port'], user=MYSQL_CONFIG['user'], password=MYSQL_CONFIG['password'], database=MYSQL_CONFIG['database'], charset='utf8mb4', autocommit=False ) def get_total_rows(self): """获取总记录数""" if self.db_connection is None: self.get_db_connection() cursor = self.db_connection.cursor() cursor.execute( "SELECT COUNT(*) FROM indeximos " "WHERE (es_video IS NULL OR es_video IN ('-1')) " "AND es_loadtime > %s", self.last_loadtime ) return cursor.fetchone()[0] def get_last_loadtime(self): """获取上次导出数据的时间""" if self.db_connection is None: self.get_db_connection() cursor = self.db_connection.cursor() cursor.execute( "SELECT config_value FROM config " "WHERE config_name = 'last_loadtime' " ) return cursor.fetchone()[0] def use_remote_selenium(self, url): for host in skip_host_name: if host in url: return True return False def format_pdf_filename(self, row): """格式化PDF文件名""" es_urltitle = row[2] or 'untitled' es_urltime = str(row[3]) or '19700101_000000' es_sitename = row[4] or 'anonymous' def clean_filename(text): if not text: return '' invalid_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*'] for char in invalid_chars: text = text.replace(char, '_') return text.strip()[:100] try: dt = datetime.strptime(es_urltime, '%Y-%m-%d %H:%M:%S') es_urltime_fix = dt.strftime('%Y%m%d_%H%M%S') except: es_urltime_fix = '19700101_000000' filename = f"{clean_filename(es_urltitle)}_{es_urltime_fix}_{es_sitename}.pdf" return os.path.join(PDF_OUTPUT_DIR, filename) def format_mhtml_filename(self, row): """格式化PDF文件名""" es_urltitle = row[2] or 'untitled' es_urltime = str(row[3]) or '19700101_000000' es_sitename = row[4] or 'anonymous' def clean_filename(text): if not text: return '' invalid_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*'] for char in invalid_chars: text = text.replace(char, '_') return text.strip()[:100] try: dt = datetime.strptime(es_urltime, '%Y-%m-%d %H:%M:%S') es_urltime_fix = dt.strftime('%Y%m%d_%H%M%S') except: es_urltime_fix = '19700101_000000' filename = f"{clean_filename(es_urltitle)}_{es_urltime_fix}_{es_sitename}.mhtml" return os.path.join(PDF_OUTPUT_DIR, filename) def fetch_data_batch(self, offset): """分页获取数据""" if self.db_connection is None: self.get_db_connection() cursor = self.db_connection.cursor() cursor.execute( "SELECT es_sid, es_urlname, es_urltitle, es_urltime, es_sitename, es_authors FROM indeximos " "WHERE (es_video IS NULL OR es_video IN ('-1')) " "AND es_loadtime > %s " "ORDER BY es_urltime LIMIT %s OFFSET %s", (self.last_loadtime, BATCH_SIZE, offset) ) return cursor.fetchall() def update_file_status(self, es_sid, status, retry=3): """更新数据库状态""" for attempt in range(retry): try: with self.db_lock: if self.db_connection is None: self.get_db_connection() cursor = self.db_connection.cursor() cursor.execute( "UPDATE indeximos SET es_video = %s WHERE es_sid = %s", (status, es_sid)) self.db_connection.commit() return True except Exception as e: if attempt == retry - 1: logger.error(f"更新数据库失败(es_sid={es_sid}): {e}") return False time.sleep(1) def extract_main_domain(self, url): extracted = tldextract.extract(url) # 组合注册域名（主域名） main_domain = f"{extracted.domain}.{extracted.suffix}" return main_domain def download_worker(self): """工作线程函数""" while True: try: task = self.task_queue.get(timeout=1) if task is None: break row = task url = row[1] if self.extract_main_domain(url) in self.skip_hosts: self.small_file_count += 1 self.processed_count += 1 self.task_queue.task_done() print(f"小文件规避，暂时跳过URL：{url}") continue output_file = self.format_pdf_filename(row) # 获取格式化后的文件名 try: os.makedirs(os.path.dirname(output_file), exist_ok=True) # 调用下载函数 if self.use_remote_selenium(url): self.processed_count += 1 self.task_queue.task_done() continue # if self.remote_handler is None: # self.remote_handler = RemotePDFSaver() # success = self.remote_handler.save_as_pdf( # url=url, # output_path=output_file, # timeout=TIMEOUT # ) else: if self.local_handler is None: self.local_handler = PDFSaver(headless=False) success = self.local_handler.save_as_pdf( url=url, output_path=output_file, timeout=TIMEOUT, wait_time=5 ) # 验证下载结果 if success and os.path.exists(output_file): file_size = os.path.getsize(output_file) if file_size >= MIN_PDF_SIZE: # 文件大小合格 self.update_file_status(row[0], output_file) self.success_count += 1 else: # 文件太小 self.update_file_status(row[0], '-2') self.small_file_count += 1 logger.warning(f"文件过小({file_size}字节): {output_file}") try: os.remove(output_file) self.skip_hosts.append(self.extract_main_domain(url)) except: pass else: # 下载失败 self.update_file_status(row[0], '0') self.fail_count += 1 if os.path.exists(output_file): try: os.remove(output_file) except: pass except Exception as e: logger.error(f"下载出现异常(es_sid={row[0]}, url={url}): {str(e)}") self.update_file_status(row[0], '-1') self.fail_count += 1 self.processed_count += 1 self.task_queue.task_done() except queue.Empty: continue def run(self): """启动下载任务""" threads = [] # 创建工作线程 for _ in range(MAX_WORKERS): t = threading.Thread(target=self.download_worker) t.start() threads.append(t) # 使用进度条显示进度 with tqdm(total=self.total_rows, desc="处理进度", unit="条") as pbar: offset = 0 while True: batch = self.fetch_data_batch(offset) if not batch: break batch_list = list(batch) random.shuffle(batch_list) batch = tuple(batch_list) for row in batch: self.task_queue.put(row) pbar.update(len(batch)) pbar.set_postfix({ '成功': self.success_count, '失败': self.fail_count, '小文件': self.small_file_count, '速度': f"{self.processed_count / (time.time() - self.start_time):.1f}条/秒" }) offset += BATCH_SIZE self.task_queue.join() for _ in range(MAX_WORKERS): self.task_queue.put(None) for t in threads: t.join() total_time = time.time() - self.start_time print(f"\n处理完成! 总计: {self.total_rows}条") print(f"成功: {self.success_count}条, 失败: {self.fail_count}条, 小文件: {self.small_file_count}条") print(f"总耗时: {total_time:.2f}秒, 平均速度: {self.total_rows / total_time:.2f}条/秒") def terminate(self): if self.local_handler: self.local_handler.quit() if self.remote_handler: self.remote_handler.quit() self.db_connection.close() if __name__ == "__main__": while running: try: print(f"开始处理，总记录数: {PDFDownloader().get_total_rows()}") downloader = PDFDownloader() downloader.run() print(f"运行完成，暂停{running_interval_seconds}秒后开始下一次运行...") downloader.terminate() time.sleep(running_interval_seconds) except Exception as e: print(repr(e))