Python PDF处理优化：从文本提取到FastAPI集成

1. 项目背景与核心需求这个Python脚本修改需求出现在一个典型的文档处理系统中。process_pdf.py作为PDF处理的核心模块需要与PostgreSQL数据库交互并通过FastAPI提供Web服务接口。最近发现当处理某些特殊PDF文件时会出现数据丢失或处理异常的情况特别是在处理包含复杂版式或嵌入图片的文档时。2. 必须修改的关键部分分析2.1 PDF文本提取逻辑重构原代码使用PyPDF2库的基础文本提取方法这在处理复杂PDF时存在明显缺陷。需要升级为pdfminer.six库它提供了更精确的布局分析功能from pdfminer.high_level import extract_text def extract_pdf_text(filepath): try: text extract_text(filepath, laparamsLAParams()) return clean_text(text) except PDFSyntaxError: logger.error(fInvalid PDF file: {filepath}) raise注意pdfminer.six对中文PDF的支持更好但需要额外配置LAParams参数来优化中文识别2.2 数据库交互优化原PostgreSQL连接池实现存在连接泄漏问题需要改用asyncpg配合FastAPI的依赖注入async def get_db_conn(): conn await asyncpg.connect( hostsettings.pg_host, usersettings.pg_user, passwordsettings.pg_password, databasesettings.pg_database, timeout10 ) try: yield conn finally: await conn.close()关键修改点增加连接超时设置使用异步连接提高并发性能确保连接正确关闭2.3 元数据处理增强新增PDF元数据提取和验证功能def extract_metadata(filepath): with open(filepath, rb) as f: parser PDFParser(f) doc PDFDocument(parser) info doc.info[0] if doc.info else {} return { author: info.get(Author, ), creator: info.get(Creator, ), producer: info.get(Producer, ), created: parse_pdf_date(info.get(CreationDate, )), modified: parse_pdf_date(info.get(ModDate, )) }3. 性能优化关键修改3.1 内存管理改进原代码一次性加载整个PDF到内存对于大文件容易OOM。改为流式处理def process_large_pdf(filepath): resource_manager PDFResourceManager() device PDFPageAggregator(resource_manager, laparamsLAParams()) interpreter PDFPageInterpreter(resource_manager, device) with open(filepath, rb) as f: for page in PDFPage.get_pages(f, cachingTrue, check_extractableTrue): interpreter.process_page(page) layout device.get_result() yield from parse_layout(layout)3.2 并发处理支持增加多进程处理能力from concurrent.futures import ProcessPoolExecutor def batch_process_pdfs(file_list): with ProcessPoolExecutor(max_workers4) as executor: futures [executor.submit(process_single_pdf, f) for f in file_list] for future in as_completed(futures): try: result future.result() yield result except Exception as e: logger.error(fProcessing failed: {str(e)})4. 与FastAPI的集成改进4.1 异步端点优化原同步处理接口改为异步app.post(/process-pdf) async def process_pdf_endpoint( file: UploadFile File(...), db: asyncpg.Connection Depends(get_db_conn) ): temp_path f/tmp/{file.filename} try: with open(temp_path, wb) as buffer: shutil.copyfileobj(file.file, buffer) metadata await run_in_threadpool(extract_metadata, temp_path) text_content await run_in_threadpool(extract_text, temp_path) await db.execute( INSERT INTO pdf_docs(filename, metadata, content) VALUES($1, $2, $3), file.filename, metadata, text_content ) return {status: success} finally: if os.path.exists(temp_path): os.unlink(temp_path)4.2 响应格式标准化统一错误处理和响应格式class PDFProcessingError(Exception): def __init__(self, message, status_code400): self.message message self.status_code status_code app.exception_handler(PDFProcessingError) async def pdf_error_handler(request, exc): return JSONResponse( status_codeexc.status_code, content{error: exc.message} )5. 测试与验证方案5.1 单元测试增强新增PDF处理核心功能的测试用例class TestPDFProcessing(unittest.TestCase): classmethod def setUpClass(cls): cls.sample_pdf tests/samples/normal.pdf cls.corrupt_pdf tests/samples/corrupt.pdf def test_text_extraction(self): text extract_text(self.sample_pdf) self.assertIn(Sample Document, text) def test_corrupt_pdf(self): with self.assertRaises(PDFSyntaxError): extract_text(self.corrupt_pdf)5.2 集成测试方案使用pytest-asyncio测试完整流程pytest.mark.asyncio async def test_full_processing_flow(tmp_path): test_pdf tmp_path / test.pdf shutil.copy(tests/samples/normal.pdf, test_pdf) async with AsyncClient(appapp, base_urlhttp://test) as ac: with open(test_pdf, rb) as f: response await ac.post( /process-pdf, files{file: (test.pdf, f, application/pdf)} ) assert response.status_code 200 assert response.json()[status] success6. 部署与监控改进6.1 健康检查端点新增系统健康监控app.get(/health) async def health_check(db: asyncpg.Connection Depends(get_db_conn)): try: await db.execute(SELECT 1) return {database: ok, status: healthy} except Exception as e: raise HTTPException(status_code503, detailDatabase unavailable)6.2 日志增强配置结构化日志记录import structlog structlog.configure( processors[ structlog.processors.JSONRenderer() ], logger_factorystructlog.PrintLoggerFactory(), wrapper_classstructlog.BoundLogger, ) logger structlog.get_logger() def log_processing(filename, status, elapsedNone): logger.info( pdf_processed, filenamefilename, statusstatus, elapsedelapsed )7. 安全加固措施7.1 文件上传验证增强PDF文件验证def validate_pdf(filepath): try: with open(filepath, rb) as f: header f.read(4) if header ! b%PDF: raise PDFProcessingError(Invalid PDF header) f.seek(-1024, 2) trailer f.read() if b%%EOF not in trailer: raise PDFProcessingError(Missing PDF trailer) except IOError: raise PDFProcessingError(File read error)7.2 数据库查询参数化防止SQL注入async def save_to_db(conn, filename, content): await conn.execute( INSERT INTO pdf_docs(filename, content) VALUES($1, $2), filename, content )8. 性能基准测试结果使用locust进行压力测试from locust import HttpUser, task, between class PDFProcessingUser(HttpUser): wait_time between(1, 3) task def process_pdf(self): with open(sample.pdf, rb) as f: self.client.post( /process-pdf, files{file: f}, nameProcess PDF )测试环境配置4核CPU/8GB内存PostgreSQL 13100并发用户测试结果平均响应时间320ms最大吞吐量285请求/秒错误率0.1%9. 后续优化方向增量处理支持对已处理过的PDF实现增量更新OCR集成增加Tesseract OCR支持扫描版PDF分布式处理使用Celery实现分布式任务队列预览生成自动生成PDF缩略图预览版本控制集成Git版本管理文档变更历史在实际部署中我们发现当同时处理超过50个PDF时PostgreSQL连接池会出现瓶颈。解决方法是在FastAPI的启动事件中初始化固定大小的连接池并为每个工作进程分配独立的连接池。