feat(personal): process scanned PDFs in OCR batches

This commit is contained in:
2026-07-12 21:33:18 +08:00
parent d63532cc90
commit 0e3a68c3c1
4 changed files with 729 additions and 5 deletions
@@ -35,6 +35,7 @@ public class PersonalIngestionWorker {
private final TransactionTemplate transactionTemplate;
private final StoredObjectReader objectReader;
private final TenantRunner tenantRunner;
private final PdfOcrHandler pdfOcrHandler;
private final long maxInputBytes;
private final int chunkSize;
private final int chunkOverlap;
@@ -44,22 +45,25 @@ public class PersonalIngestionWorker {
@Autowired
public PersonalIngestionWorker(JdbcTemplate jdbcTemplate, ISysOssService ossService,
KnowledgeDocumentParser parser, PersonalKnowledgeProperties properties,
PlatformTransactionManager transactionManager) {
PlatformTransactionManager transactionManager,
PersonalPdfOcrService pdfOcrService) {
this(jdbcTemplate, parser, new TransactionTemplate(transactionManager),
defaultReader(ossService, PersonalIngestionWorker::ossClient), configuredMaxBytes(properties),
properties.getChunkSize(), properties.getChunkOverlap(), properties.getParsingLeaseMinutes(),
properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant);
properties.getMaxParseAttempts(), PersonalIngestionWorker::runInTenant, pdfOcrService::enqueue);
}
private PersonalIngestionWorker(JdbcTemplate jdbcTemplate, KnowledgeDocumentParser parser,
TransactionTemplate transactionTemplate, StoredObjectReader objectReader,
long maxInputBytes, int chunkSize, int chunkOverlap,
int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner) {
int parsingLeaseMinutes, int maxParseAttempts, TenantRunner tenantRunner,
PdfOcrHandler pdfOcrHandler) {
this.jdbcTemplate = jdbcTemplate;
this.parser = parser;
this.transactionTemplate = transactionTemplate;
this.objectReader = objectReader;
this.tenantRunner = tenantRunner;
this.pdfOcrHandler = pdfOcrHandler;
this.maxInputBytes = maxInputBytes;
if (chunkSize <= 0 || chunkOverlap < 0 || chunkOverlap >= chunkSize
|| parsingLeaseMinutes <= 0 || maxParseAttempts <= 0) {
@@ -76,7 +80,8 @@ public class PersonalIngestionWorker {
TransactionTemplate transactionTemplate,
StoredObjectReader objectReader) {
return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader,
20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute());
20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute(),
unsupportedPdfOcr());
}
public static PersonalIngestionWorker forTest(JdbcTemplate jdbcTemplate, ISysOssService ossService,
@@ -84,7 +89,17 @@ public class PersonalIngestionWorker {
TransactionTemplate transactionTemplate,
StoredObjectReader objectReader, TenantRunner tenantRunner) {
return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader,
20L * 1024 * 1024, 800, 120, 15, 3, tenantRunner);
20L * 1024 * 1024, 800, 120, 15, 3, tenantRunner, unsupportedPdfOcr());
}
public static PersonalIngestionWorker forTestWithPdfOcr(JdbcTemplate jdbcTemplate, ISysOssService ossService,
KnowledgeDocumentParser parser,
TransactionTemplate transactionTemplate,
StoredObjectReader objectReader,
PdfOcrHandler pdfOcrHandler) {
return new PersonalIngestionWorker(jdbcTemplate, parser, transactionTemplate, objectReader,
20L * 1024 * 1024, 800, 120, 15, 3, (tenantId, operation) -> operation.execute(),
pdfOcrHandler);
}
public static StoredObjectReader objectReaderForTest(ISysOssService ossService,
@@ -173,6 +188,12 @@ public class PersonalIngestionWorker {
ParsedDocument document = parser.parse(stored.fileName(), item.mimeType(), stored.bytes());
List<String> chunks = document.chunks(chunkSize, chunkOverlap);
if (chunks.isEmpty()) {
if (isPdf(item, stored)) {
pdfOcrHandler.enqueue(new PersonalPdfOcrService.OcrItem(item.id(), item.tenantId(),
item.spaceId(), item.ownerUserId(), item.ossId(), item.tagsJson(), attemptVersion),
stored.bytes());
return true;
}
throw new KnowledgeDocumentParser.ParseException(
KnowledgeDocumentParser.Failure.EMPTY, "document contains no text");
}
@@ -308,6 +329,12 @@ public class PersonalIngestionWorker {
&& "PERSONAL_OSS_NOT_PRIVATE".equals(serviceError.getMessage())) {
return new Failure("PERSONAL_OSS_NOT_PRIVATE", "个人资料存储策略不可用");
}
if (candidate instanceof PersonalPdfPageRenderer.PdfPageLimitException) {
return new Failure("PERSONAL_PDF_PAGE_LIMIT", "文件超过200页,请拆分后重新上传");
}
if (candidate instanceof PersonalPdfPageRenderer.PdfRenderException renderError) {
return new Failure(renderError.code(), "扫描PDF预处理失败,请检查文件后重试");
}
candidate = candidate.getCause();
}
return new Failure("PERSONAL_PARSE_FAILED", "资料处理失败,请稍后重试");
@@ -351,6 +378,18 @@ public class PersonalIngestionWorker {
return "personal/" + item.tenantId() + "/" + item.ownerUserId() + "/" + item.id() + "/";
}
private static boolean isPdf(Item item, StoredObject stored) {
return "application/pdf".equalsIgnoreCase(item.mimeType())
|| stored.fileName().toLowerCase(java.util.Locale.ROOT).endsWith(".pdf");
}
private static PdfOcrHandler unsupportedPdfOcr() {
return (item, bytes) -> {
throw new KnowledgeDocumentParser.ParseException(
KnowledgeDocumentParser.Failure.EMPTY, "document contains no text");
};
}
@FunctionalInterface
public interface StoredObjectReader {
StoredObject read(long ossId, String expectedPrefix, long ownerUserId, long maxBytes) throws Exception;
@@ -371,6 +410,11 @@ public class PersonalIngestionWorker {
OssClient get(String configKey);
}
@FunctionalInterface
public interface PdfOcrHandler {
void enqueue(PersonalPdfOcrService.OcrItem item, byte[] pdfBytes);
}
public record StoredObject(String fileName, byte[] bytes) {
public StoredObject {
bytes = bytes == null ? new byte[0] : bytes.clone();
@@ -0,0 +1,459 @@
package org.dromara.aihr.personal.service;
import org.dromara.aihr.knowledge.parse.ParsedDocument;
import org.dromara.aihr.personal.support.PersonalKnowledgeProperties;
import org.dromara.common.core.exception.ServiceException;
import org.dromara.common.oss.core.OssClient;
import org.dromara.common.oss.enums.AccessPolicyType;
import org.dromara.common.oss.factory.OssFactory;
import org.dromara.common.tenant.helper.TenantHelper;
import org.dromara.system.domain.vo.SysOssVo;
import org.dromara.system.service.ISysOssService;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Service;
import org.springframework.transaction.PlatformTransactionManager;
import org.springframework.transaction.support.TransactionTemplate;
import java.io.InputStream;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.StringJoiner;
@Service
public class PersonalPdfOcrService {
private final JdbcTemplate jdbcTemplate;
private final PersonalPdfPageRenderer renderer;
private final PersonalVisionOcrService visionOcr;
private final TransactionTemplate transactions;
private final OcrSourceReader sourceReader;
public PersonalPdfOcrService(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer,
PersonalVisionOcrService visionOcr, PlatformTransactionManager transactionManager,
ISysOssService ossService, PersonalKnowledgeProperties properties) {
this(jdbcTemplate, renderer, visionOcr, new TransactionTemplate(transactionManager),
defaultSourceReader(ossService, maxBytes(properties)));
}
private PersonalPdfOcrService(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer,
PersonalVisionOcrService visionOcr, TransactionTemplate transactions,
OcrSourceReader sourceReader) {
this.jdbcTemplate = jdbcTemplate;
this.renderer = renderer;
this.visionOcr = visionOcr;
this.transactions = transactions;
this.sourceReader = sourceReader;
}
public static PersonalPdfOcrService forTest(JdbcTemplate jdbcTemplate, PersonalPdfPageRenderer renderer,
PersonalVisionOcrService visionOcr,
TransactionTemplate transactions,
OcrSourceReader sourceReader) {
return new PersonalPdfOcrService(jdbcTemplate, renderer, visionOcr, transactions, sourceReader);
}
public void enqueue(OcrItem item, byte[] pdfBytes) {
int totalPages = renderer.pageCount(pdfBytes);
transactions.execute(status -> {
jdbcTemplate.update("""
insert into aihr_personal_ocr_job
(tenant_id, owner_user_id, item_id, status, total_pages, processed_pages,
success_pages, failed_pages, failed_page_numbers_json, next_page,
attempt_count, create_time, update_time)
values (?, ?, ?, 'PENDING', ?, 0, 0, 0, json_array(), 1, 0, now(), now())
on duplicate key update status = 'PENDING', total_pages = values(total_pages),
processed_pages = 0, success_pages = 0, failed_pages = 0,
failed_page_numbers_json = json_array(), next_page = 1,
attempt_count = 0, last_error_code = null, lease_until = null, update_time = now()
""", item.tenantId(), item.ownerUserId(), item.id(), totalPages);
jdbcTemplate.update("""
delete from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ?
""", item.tenantId(), item.ownerUserId(), item.id());
jdbcTemplate.batchUpdate("""
insert into aihr_personal_ocr_page
(tenant_id, owner_user_id, item_id, page_number, status,
attempt_count, create_time, update_time)
values (?, ?, ?, ?, 'PENDING', 0, now(), now())
""", new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement statement, int index) throws SQLException {
statement.setString(1, item.tenantId());
statement.setLong(2, item.ownerUserId());
statement.setLong(3, item.id());
statement.setInt(4, index + 1);
}
@Override
public int getBatchSize() {
return totalPages;
}
});
return null;
});
}
public static List<Integer> batchPageNumbers(int firstPage, int totalPages) {
if (firstPage < 1 || totalPages < firstPage) {
return List.of();
}
int end = Math.min(totalPages, firstPage + PersonalPdfPageRenderer.BATCH_SIZE - 1);
List<Integer> pages = new ArrayList<>(end - firstPage + 1);
for (int page = firstPage; page <= end; page++) {
pages.add(page);
}
return List.copyOf(pages);
}
@Scheduled(fixedDelayString = "${aihr.personal.ocr-delay-ms:2000}", scheduler = "personalTaskScheduler")
public void poll() {
recoverExpiredJobs();
processNextBatch();
}
public void recoverExpiredJobs() {
jdbcTemplate.update("""
update aihr_personal_ocr_job
set status = 'RETRY', lease_until = null, last_error_code = 'PERSONAL_OCR_LEASE_EXPIRED',
update_time = now()
where status = 'RUNNING' and lease_until < now()
""");
}
public boolean processNextBatch() {
List<Map<String, Object>> rows = jdbcTemplate.queryForList("""
select j.id, j.tenant_id, j.owner_user_id, j.item_id, j.status, j.attempt_count,
j.total_pages, j.next_page,
i.space_id, i.oss_id, i.tags_json
from aihr_personal_ocr_job j
join aihr_personal_item i on i.id = j.item_id
and binary i.tenant_id = binary j.tenant_id and i.owner_user_id = j.owner_user_id
where j.status in ('PENDING','RETRY') and i.status = 'PARSING'
order by j.update_time, j.id
limit 1
""");
if (rows.isEmpty()) {
return false;
}
OcrJob job = job(rows.get(0));
int claimed = jdbcTemplate.update("""
update aihr_personal_ocr_job
set status = 'RUNNING', attempt_count = attempt_count + 1,
lease_until = date_add(now(), interval 15 minute), update_time = now()
where id = ? and status in ('PENDING','RETRY')
""", job.id());
if (claimed != 1) {
return false;
}
List<Integer> pageNumbers = retryJob(job) ? retryPageNumbers(job)
: batchPageNumbers(job.nextPage(), job.totalPages());
int success = 0;
int failed = 0;
try {
byte[] pdfBytes = sourceReader.read(job);
List<PersonalPdfPageRenderer.RenderedPage> rendered = renderPages(pdfBytes, pageNumbers, retryJob(job));
for (PersonalPdfPageRenderer.RenderedPage page : rendered) {
int pageNumber = page.pageIndex() + 1;
try {
String text = visionOcr.recognize(page.bytes(), page.mimeType(), pageNumber);
if (text.isBlank()) {
throw new PersonalVisionOcrService.OcrUnavailableException("PERSONAL_OCR_PAGE_EMPTY");
}
jdbcTemplate.update("""
update aihr_personal_ocr_page
set status = 'SUCCESS', ocr_text = ?, attempt_count = attempt_count + 1,
last_error_code = null, update_time = now()
where tenant_id = ? and owner_user_id = ? and item_id = ? and page_number = ?
""", text, job.tenantId(), job.ownerUserId(), job.itemId(), pageNumber);
success++;
} catch (PersonalVisionOcrService.OcrUnavailableException exception) {
jdbcTemplate.update("""
update aihr_personal_ocr_page
set status = 'FAILED', ocr_text = null, attempt_count = attempt_count + 1,
last_error_code = ?, update_time = now()
where tenant_id = ? and owner_user_id = ? and item_id = ? and page_number = ?
""", exception.code(), job.tenantId(), job.ownerUserId(), job.itemId(), pageNumber);
failed++;
}
}
int nextPage = job.nextPage() + rendered.size();
if (retryJob(job)) {
Integer remaining = jdbcTemplate.queryForObject("""
select count(*) from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ?
and status = 'FAILED' and attempt_count < 3
""", Integer.class, job.tenantId(), job.ownerUserId(), job.itemId());
if (remaining != null && remaining > 0) {
refreshRetryState(job);
} else {
publishTerminal(job);
}
} else if (nextPage <= job.totalPages()) {
jdbcTemplate.update("""
update aihr_personal_ocr_job
set processed_pages = processed_pages + ?, success_pages = success_pages + ?,
failed_pages = failed_pages + ?, next_page = ?, status = 'PENDING',
lease_until = null, update_time = now()
where id = ? and status = 'RUNNING'
""", rendered.size(), success, failed, nextPage, job.id());
} else if (failed > 0) {
jdbcTemplate.update("""
update aihr_personal_ocr_job
set processed_pages = processed_pages + ?, success_pages = success_pages + ?,
failed_pages = failed_pages + ?, next_page = total_pages + 1,
status = 'RETRY', lease_until = null, update_time = now()
where id = ? and status = 'RUNNING'
""", rendered.size(), success, failed, job.id());
} else {
publishTerminal(job);
}
} catch (Exception exception) {
if (job.attemptCount() + 1 >= 3) {
transactions.execute(status -> {
jdbcTemplate.update("""
update aihr_personal_ocr_job
set status = 'FAILED', last_error_code = 'PERSONAL_OCR_BATCH_FAILED',
lease_until = null, update_time = now()
where id = ? and status = 'RUNNING'
""", job.id());
jdbcTemplate.update("""
update aihr_personal_item
set status = 'FAILED', error_code = 'PERSONAL_OCR_BATCH_FAILED',
error_message = '扫描PDF识别任务连续失败,请稍后重试', update_time = now()
where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING'
""", job.tenantId(), job.ownerUserId(), job.itemId());
return null;
});
} else {
jdbcTemplate.update("""
update aihr_personal_ocr_job
set status = 'RETRY', last_error_code = 'PERSONAL_OCR_BATCH_FAILED',
lease_until = null, update_time = now()
where id = ? and status = 'RUNNING'
""", job.id());
}
}
return true;
}
private List<Integer> retryPageNumbers(OcrJob job) {
return jdbcTemplate.queryForList("""
select page_number from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ?
and status = 'FAILED' and attempt_count < 3
order by page_number limit ?
""", Integer.class, job.tenantId(), job.ownerUserId(), job.itemId(),
PersonalPdfPageRenderer.BATCH_SIZE);
}
private List<PersonalPdfPageRenderer.RenderedPage> renderPages(byte[] pdfBytes, List<Integer> pageNumbers,
boolean retry) {
if (!retry) {
if (pageNumbers.isEmpty()) return List.of();
return renderer.render(pdfBytes, pageNumbers.get(0) - 1, pageNumbers.size());
}
List<PersonalPdfPageRenderer.RenderedPage> rendered = new ArrayList<>(pageNumbers.size());
for (int pageNumber : pageNumbers) {
List<PersonalPdfPageRenderer.RenderedPage> page = renderer.render(pdfBytes, pageNumber - 1, 1);
if (!page.isEmpty()) rendered.add(page.get(0));
}
return List.copyOf(rendered);
}
private void refreshRetryState(OcrJob job) {
jdbcTemplate.update("""
update aihr_personal_ocr_job j
set j.processed_pages = (select count(*) from aihr_personal_ocr_page p
where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id
and p.item_id = j.item_id and (p.status = 'SUCCESS' or p.attempt_count >= 3)),
j.success_pages = (select count(*) from aihr_personal_ocr_page p
where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id
and p.item_id = j.item_id and p.status = 'SUCCESS'),
j.failed_pages = (select count(*) from aihr_personal_ocr_page p
where p.tenant_id = j.tenant_id and p.owner_user_id = j.owner_user_id
and p.item_id = j.item_id and p.status = 'FAILED'),
j.status = 'RETRY', j.lease_until = null, j.update_time = now()
where j.id = ? and j.status = 'RUNNING'
""", job.id());
}
private static boolean retryJob(OcrJob job) {
return "RETRY".equals(job.status());
}
private void publishTerminal(OcrJob job) {
List<String> pageTexts = jdbcTemplate.queryForList("""
select ocr_text from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ? and status = 'SUCCESS'
order by page_number
""", String.class, job.tenantId(), job.ownerUserId(), job.itemId());
List<Integer> failedPages = jdbcTemplate.queryForList("""
select page_number from aihr_personal_ocr_page
where tenant_id = ? and owner_user_id = ? and item_id = ? and status = 'FAILED'
order by page_number
""", Integer.class, job.tenantId(), job.ownerUserId(), job.itemId());
transactions.execute(status -> {
if (pageTexts.isEmpty()) {
jdbcTemplate.update("""
update aihr_personal_item
set status = 'FAILED', error_code = 'PERSONAL_OCR_EMPTY',
error_message = '扫描PDF未识别到可用文字', update_time = now()
where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING'
""", job.tenantId(), job.ownerUserId(), job.itemId());
} else {
publishFragments(job, String.join("\n\n", pageTexts), failedPages);
}
jdbcTemplate.update("""
update aihr_personal_ocr_job
set processed_pages = ?, success_pages = ?, failed_pages = ?,
failed_page_numbers_json = ?, next_page = total_pages + 1,
status = 'COMPLETED', lease_until = null, update_time = now()
where id = ? and status = 'RUNNING'
""", job.totalPages(), pageTexts.size(), failedPages.size(), jsonNumbers(failedPages), job.id());
return null;
});
}
private void publishFragments(OcrJob job, String text, List<Integer> failedPages) {
List<String> chunks = new ParsedDocument(text, "application/pdf", Map.of()).chunks(800, 120);
jdbcTemplate.update("""
delete from aihr_personal_fragment
where tenant_id = ? and owner_user_id = ? and item_id = ?
""", job.tenantId(), job.ownerUserId(), job.itemId());
jdbcTemplate.batchUpdate("""
insert into aihr_personal_fragment
(tenant_id, space_id, owner_user_id, item_id, idx, content, token_count,
embedding_json, embedding_model, embedding_time, create_time)
values (?, ?, ?, ?, ?, ?, ?, null, null, null, now())
""", new BatchPreparedStatementSetter() {
@Override
public void setValues(PreparedStatement statement, int index) throws SQLException {
String chunk = chunks.get(index);
statement.setString(1, job.tenantId());
statement.setLong(2, job.spaceId());
statement.setLong(3, job.ownerUserId());
statement.setLong(4, job.itemId());
statement.setInt(5, index);
statement.setString(6, chunk);
statement.setInt(7, Math.max(1, (chunk.codePointCount(0, chunk.length()) + 1) / 2));
}
@Override
public int getBatchSize() {
return chunks.size();
}
});
if (failedPages.isEmpty()) {
jdbcTemplate.update("""
update aihr_personal_item
set status = 'READY', parsed_at = now(), summary = ?, tags_json = ?,
error_code = null, error_message = null, update_time = now()
where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING'
""", summary(text), job.tagsJson(), job.tenantId(), job.ownerUserId(), job.itemId());
} else {
jdbcTemplate.update("""
update aihr_personal_item
set status = 'READY', parsed_at = now(), summary = ?, tags_json = ?,
error_code = 'PERSONAL_OCR_PARTIAL', error_message = ?, update_time = now()
where tenant_id = ? and owner_user_id = ? and id = ? and status = 'PARSING'
""", summary(text), job.tagsJson(), "已收录," + failedPages.size() + "页识别失败",
job.tenantId(), job.ownerUserId(), job.itemId());
}
}
private static String jsonNumbers(List<Integer> values) {
StringJoiner joiner = new StringJoiner(",", "[", "]");
values.forEach(value -> joiner.add(String.valueOf(value)));
return joiner.toString();
}
private static String summary(String text) {
int[] codePoints = text.codePoints().limit(300).toArray();
return new String(codePoints, 0, codePoints.length).trim();
}
private static OcrJob job(Map<String, Object> row) {
return new OcrJob(number(row, "id"), String.valueOf(row.get("tenant_id")),
number(row, "owner_user_id"), number(row, "item_id"), number(row, "space_id"),
number(row, "oss_id"), row.get("tags_json") == null ? "[]" : String.valueOf(row.get("tags_json")),
String.valueOf(row.get("status")), intValue(row, "attempt_count", 0),
Math.toIntExact(number(row, "total_pages")),
Math.toIntExact(number(row, "next_page")));
}
private static long number(Map<String, Object> row, String key) {
Object value = row.get(key);
if (!(value instanceof Number number)) {
throw new IllegalStateException("personal OCR metadata is incomplete");
}
return number.longValue();
}
private static int intValue(Map<String, Object> row, String key, int fallback) {
Object value = row.get(key);
return value instanceof Number number ? Math.toIntExact(number.longValue()) : fallback;
}
private static OcrSourceReader defaultSourceReader(ISysOssService ossService, long maxBytes) {
return job -> runInTenant(job.tenantId(), () -> {
SysOssVo object = ossService.getById(job.ossId());
if (object == null || object.getCreateBy() == null || object.getCreateBy() != job.ownerUserId()
|| object.getService() == null || object.getFileName() == null
|| !object.getFileName().startsWith(ownerPrefix(job))) {
throw new ServiceException("PERSONAL_OCR_SOURCE_UNAVAILABLE");
}
OssClient client = OssFactory.instance(object.getService());
if (client.getAccessPolicy() != AccessPolicyType.PRIVATE) {
throw new ServiceException("PERSONAL_OSS_NOT_PRIVATE");
}
try (InputStream input = client.getObjectContent(object.getFileName())) {
int limit = (int) Math.min(Integer.MAX_VALUE - 1L, maxBytes);
byte[] bytes = input.readNBytes(limit + 1);
if (bytes.length > maxBytes) throw new ServiceException("PERSONAL_PARSE_TOO_LARGE");
return bytes;
}
});
}
private static byte[] runInTenant(String tenantId, OcrSourceOperation operation) throws Exception {
String previous = TenantHelper.getDynamic();
TenantHelper.setDynamic(tenantId);
try {
return operation.execute();
} finally {
TenantHelper.clearDynamic();
if (previous != null && !previous.isBlank()) TenantHelper.setDynamic(previous);
}
}
private static long maxBytes(PersonalKnowledgeProperties properties) {
return Math.multiplyExact(properties.getMaxFileSizeMb(), 1024L * 1024L);
}
private static String ownerPrefix(OcrJob job) {
return "personal/" + job.tenantId() + "/" + job.ownerUserId() + "/" + job.itemId() + "/";
}
@FunctionalInterface
public interface OcrSourceReader {
byte[] read(OcrJob job) throws Exception;
}
@FunctionalInterface
private interface OcrSourceOperation {
byte[] execute() throws Exception;
}
public record OcrItem(long id, String tenantId, long spaceId, long ownerUserId, long ossId,
String tagsJson, int attemptVersion) {
}
public record OcrJob(long id, String tenantId, long ownerUserId, long itemId, long spaceId,
long ossId, String tagsJson, String status, int attemptCount, int totalPages, int nextPage) {
}
}
@@ -37,6 +37,28 @@ import static org.mockito.Mockito.when;
@Tag("dev")
class PersonalIngestionWorkerTest {
@Test
void emptyPdfHandsOffToOcrInsteadOfMarkingParseEmpty() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
KnowledgeDocumentParser parser = mock(KnowledgeDocumentParser.class);
when(jdbc.queryForList(contains("status = 'QUEUED'"))).thenReturn(List.of(item()));
when(jdbc.update(contains("status = 'PARSING'"), eq("000000"), eq(101L), eq(9L), eq(0))).thenReturn(1);
when(parser.parse(eq("sample.pdf"), eq("text/plain"), any(byte[].class)))
.thenReturn(new ParsedDocument("", "application/pdf", Map.of()));
AtomicReference<org.dromara.aihr.personal.service.PersonalPdfOcrService.OcrItem> handedOff =
new AtomicReference<>();
PersonalIngestionWorker worker = PersonalIngestionWorker.forTestWithPdfOcr(
jdbc, mock(ISysOssService.class), parser, immediateTransactions(),
(ossId, prefix, ownerUserId, maxBytes) ->
new PersonalIngestionWorker.StoredObject("sample.pdf", new byte[] {1, 2, 3}),
(item, bytes) -> handedOff.set(item));
assertTrue(worker.processNext());
assertEquals(9L, handedOff.get().id());
verify(jdbc, never()).update(contains("PERSONAL_PARSE_EMPTY"), any(), any(), any(), any(), any());
}
@Test
void workerReadsPrivateObjectInsideItemTenantScope() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
@@ -0,0 +1,199 @@
package org.dromara.aihr.personal;
import org.dromara.aihr.personal.service.PersonalPdfOcrService;
import org.dromara.aihr.personal.service.PersonalPdfOcrService.OcrItem;
import org.dromara.aihr.personal.service.PersonalPdfPageRenderer;
import org.dromara.aihr.personal.service.PersonalVisionOcrService;
import org.junit.jupiter.api.Tag;
import org.junit.jupiter.api.Test;
import org.mockito.ArgumentCaptor;
import org.springframework.jdbc.core.BatchPreparedStatementSetter;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.transaction.support.TransactionCallback;
import org.springframework.transaction.support.TransactionTemplate;
import java.util.List;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.contains;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
@Tag("dev")
class PersonalPdfOcrServiceTest {
@Test
void enqueueCreatesOwnerScopedJobAndOnePendingRowPerPage() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class);
when(renderer.pageCount(any(byte[].class))).thenReturn(8);
TransactionTemplate transactions = immediateTransactions();
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, renderer, mock(PersonalVisionOcrService.class), transactions,
job -> new byte[] {1});
OcrItem item = new OcrItem(9L, "000000", 7L, 101L, 88L, "[]", 1);
service.enqueue(item, new byte[] {1, 2, 3});
verify(jdbc).update(contains("insert into aihr_personal_ocr_job"),
eq("000000"), eq(101L), eq(9L), eq(8));
ArgumentCaptor<BatchPreparedStatementSetter> pages =
ArgumentCaptor.forClass(BatchPreparedStatementSetter.class);
verify(jdbc).batchUpdate(contains("insert into aihr_personal_ocr_page"), pages.capture());
assertEquals(8, pages.getValue().getBatchSize());
}
@Test
void batchPageNumbersNeverExceedTwentyAndContinueWherePreviousBatchEnded() {
assertEquals(List.of(1, 2, 3), PersonalPdfOcrService.batchPageNumbers(1, 3));
assertEquals(20, PersonalPdfOcrService.batchPageNumbers(1, 86).size());
assertEquals(21, PersonalPdfOcrService.batchPageNumbers(21, 86).get(0));
assertEquals(40, PersonalPdfOcrService.batchPageNumbers(21, 86).get(19));
assertEquals(List.of(81, 82, 83, 84, 85, 86), PersonalPdfOcrService.batchPageNumbers(81, 86));
}
@Test
void processNextBatchRendersAndRecognizesAtMostTwentyPages() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class);
PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class);
Map<String, Object> job = Map.of(
"id", 5L, "tenant_id", "000000", "owner_user_id", 101L, "item_id", 9L,
"space_id", 7L, "oss_id", 88L, "tags_json", "[]", "total_pages", 86, "next_page", 1);
when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job));
when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1);
List<PersonalPdfPageRenderer.RenderedPage> rendered = java.util.stream.IntStream.range(0, 20)
.mapToObj(index -> new PersonalPdfPageRenderer.RenderedPage(index, new byte[] {1}, "image/jpeg"))
.toList();
when(renderer.render(any(byte[].class), eq(0), eq(20))).thenReturn(rendered);
when(vision.recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class)))
.thenAnswer(invocation -> "第" + invocation.<Integer>getArgument(2) + "页");
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9, 8, 7});
assertEquals(true, service.processNextBatch());
verify(renderer).render(any(byte[].class), eq(0), eq(20));
verify(vision, times(20)).recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class));
verify(jdbc).update(contains("next_page = ?"), eq(20), eq(20), eq(0), eq(21), eq(5L));
}
@Test
void finalBatchPublishesFragmentsAndMarksItemReady() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class);
PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class);
Map<String, Object> job = Map.of(
"id", 5L, "tenant_id", "000000", "owner_user_id", 101L, "item_id", 9L,
"space_id", 7L, "oss_id", 88L, "tags_json", "[]", "total_pages", 2, "next_page", 1);
when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job));
when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1);
when(renderer.render(any(byte[].class), eq(0), eq(2))).thenReturn(List.of(
new PersonalPdfPageRenderer.RenderedPage(0, new byte[] {1}, "image/jpeg"),
new PersonalPdfPageRenderer.RenderedPage(1, new byte[] {2}, "image/jpeg")));
when(vision.recognize(any(byte[].class), eq("image/jpeg"), eq(1))).thenReturn("第一页内容");
when(vision.recognize(any(byte[].class), eq("image/jpeg"), eq(2))).thenReturn("第二页内容");
when(jdbc.queryForList(contains("select ocr_text"), eq(String.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of("第一页内容", "第二页内容"));
when(jdbc.queryForList(contains("select page_number"), eq(Integer.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of());
when(jdbc.update(contains("set status = 'READY'"), any(), eq("[]"),
eq("000000"), eq(101L), eq(9L))).thenReturn(1);
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9, 8, 7});
assertEquals(true, service.processNextBatch());
verify(jdbc).batchUpdate(contains("insert into aihr_personal_fragment"),
any(BatchPreparedStatementSetter.class));
verify(jdbc).update(contains("set status = 'READY'"), any(), eq("[]"),
eq("000000"), eq(101L), eq(9L));
verify(jdbc).update(contains("status = 'COMPLETED'"), eq(2), eq(2), eq(0), any(), eq(5L));
}
@Test
void retryJobProcessesOnlyFailedPagesThenPublishes() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalPdfPageRenderer renderer = mock(PersonalPdfPageRenderer.class);
PersonalVisionOcrService vision = mock(PersonalVisionOcrService.class);
Map<String, Object> job = Map.ofEntries(
Map.entry("id", 5L), Map.entry("tenant_id", "000000"), Map.entry("owner_user_id", 101L),
Map.entry("item_id", 9L), Map.entry("space_id", 7L), Map.entry("oss_id", 88L),
Map.entry("tags_json", "[]"), Map.entry("total_pages", 30), Map.entry("next_page", 31),
Map.entry("status", "RETRY"));
when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job));
when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1);
when(jdbc.queryForList(contains("attempt_count < 3"), eq(Integer.class),
eq("000000"), eq(101L), eq(9L), eq(20))).thenReturn(List.of(4, 27));
when(renderer.render(any(byte[].class), eq(3), eq(1))).thenReturn(List.of(
new PersonalPdfPageRenderer.RenderedPage(3, new byte[] {4}, "image/jpeg")));
when(renderer.render(any(byte[].class), eq(26), eq(1))).thenReturn(List.of(
new PersonalPdfPageRenderer.RenderedPage(26, new byte[] {27}, "image/jpeg")));
when(vision.recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class)))
.thenReturn("重试成功");
when(jdbc.queryForObject(contains("attempt_count < 3"), eq(Integer.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(0);
when(jdbc.queryForList(contains("select ocr_text"), eq(String.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of("已有页面", "重试成功", "重试成功"));
when(jdbc.queryForList(contains("select page_number"), eq(Integer.class),
eq("000000"), eq(101L), eq(9L))).thenReturn(List.of());
when(jdbc.update(contains("set status = 'READY'"), any(), eq("[]"),
eq("000000"), eq(101L), eq(9L))).thenReturn(1);
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, renderer, vision, immediateTransactions(), ignored -> new byte[] {9});
assertEquals(true, service.processNextBatch());
verify(renderer).render(any(byte[].class), eq(3), eq(1));
verify(renderer).render(any(byte[].class), eq(26), eq(1));
verify(vision, times(2)).recognize(any(byte[].class), eq("image/jpeg"), any(Integer.class));
verify(jdbc).update(contains("set status = 'READY'"), any(), eq("[]"),
eq("000000"), eq(101L), eq(9L));
}
@Test
void recoversExpiredRunningLease() {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, mock(PersonalPdfPageRenderer.class), mock(PersonalVisionOcrService.class),
immediateTransactions(), ignored -> new byte[] {1});
service.recoverExpiredJobs();
verify(jdbc).update(contains("lease_until < now()"));
}
@Test
void thirdBatchInfrastructureFailureStopsRetryLoop() throws Exception {
JdbcTemplate jdbc = mock(JdbcTemplate.class);
Map<String, Object> job = Map.ofEntries(
Map.entry("id", 5L), Map.entry("tenant_id", "000000"), Map.entry("owner_user_id", 101L),
Map.entry("item_id", 9L), Map.entry("space_id", 7L), Map.entry("oss_id", 88L),
Map.entry("tags_json", "[]"), Map.entry("total_pages", 2), Map.entry("next_page", 1),
Map.entry("status", "PENDING"), Map.entry("attempt_count", 2));
when(jdbc.queryForList(contains("from aihr_personal_ocr_job"))).thenReturn(List.of(job));
when(jdbc.update(contains("set status = 'RUNNING'"), eq(5L))).thenReturn(1);
PersonalPdfOcrService service = PersonalPdfOcrService.forTest(
jdbc, mock(PersonalPdfPageRenderer.class), mock(PersonalVisionOcrService.class),
immediateTransactions(), ignored -> { throw new IllegalStateException("storage unavailable"); });
assertEquals(true, service.processNextBatch());
verify(jdbc).update(contains("status = 'FAILED'"), eq(5L));
verify(jdbc).update(contains("PERSONAL_OCR_BATCH_FAILED"), eq("000000"), eq(101L), eq(9L));
}
private TransactionTemplate immediateTransactions() {
TransactionTemplate transactions = mock(TransactionTemplate.class);
when(transactions.execute(any())).thenAnswer(invocation -> {
TransactionCallback<?> callback = invocation.getArgument(0);
return callback.doInTransaction(null);
});
return transactions;
}
}