Skip to content

Commit d9ee046

Browse files
KDB-WindKDB-Wind
committed
fix(knowledge): 文档上传失败时清理存储文件
1 parent dc0d001 commit d9ee046

2 files changed

Lines changed: 369 additions & 2 deletions

File tree

‎bootstrap/src/main/java/com/nageoffer/ai/ragent/knowledge/service/impl/KnowledgeDocumentServiceImpl.java‎

Lines changed: 39 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -147,13 +147,13 @@ public KnowledgeDocumentVO upload(String kbId, KnowledgeDocumentUploadRequest re
147147

148148
SourceType sourceType = SourceType.normalize(requestParam.getSourceType());
149149
validateSourceAndSchedule(sourceType, requestParam);
150+
ProcessModeConfig modeConfig = resolveProcessModeConfig(requestParam);
150151
StoredFileDTO stored = resolveStoredFile(kbDO.getCollectionName(), sourceType, requestParam.getSourceLocation(), file);
151152
// 前置拦截：与分块阶段同一套 MIME 路由，无解析器的类型直接拒绝，不落库不发 MQ
152153
if (parserSelector.selectByMimeType(stored.getMimeType()) == null) {
153154
fileStorageService.deleteByUrl(stored.getUrl());
154155
throw new ClientException("暂不支持的文件类型：" + stored.getDetectedType());
155156
}
156-
ProcessModeConfig modeConfig = resolveProcessModeConfig(requestParam);
157157

158158
KnowledgeDocumentDO documentDO = KnowledgeDocumentDO.builder()
159159
.kbId(kbId)
@@ -175,7 +175,15 @@ public KnowledgeDocumentVO upload(String kbId, KnowledgeDocumentUploadRequest re
175175
.createdBy(UserContext.getUsername())
176176
.updatedBy(UserContext.getUsername())
177177
.build();
178-
documentMapper.insert(documentDO);
178+
try {
179+
int inserted = documentMapper.insert(documentDO);
180+
if (inserted <= 0) {
181+
throw new ClientException("文档保存失败");
182+
}
183+
} catch (RuntimeException e) {
184+
deleteStoredFileIfDocumentAbsent(documentDO, stored.getUrl());
185+
throw e;
186+
}
179187
bizChangeLogContext.put(String.valueOf(documentDO.getId()), null, documentDO);
180188
bizChangeLogContext.putName(documentDO.getDocName());
181189

@@ -985,4 +993,33 @@ private void deleteStoredFileQuietly(KnowledgeDocumentDO documentDO) {
985993
log.warn("删除文档存储文件失败, docId={}, fileUrl={}", documentDO.getId(), documentDO.getFileUrl(), e);
986994
}
987995
}
996+
997+
private void deleteStoredFileIfDocumentAbsent(KnowledgeDocumentDO documentDO, String fileUrl) {
998+
String docId = documentDO.getId();
999+
if (!StringUtils.hasText(docId)) {
1000+
log.warn("无法确认文档插入结果，保留已上传文件, fileUrl={}", fileUrl);
1001+
return;
1002+
}
1003+
try {
1004+
if (documentMapper.selectById(docId) != null) {
1005+
log.warn("文档插入结果异常但记录已存在，保留已上传文件, docId={}, fileUrl={}", docId, fileUrl);
1006+
return;
1007+
}
1008+
} catch (Exception e) {
1009+
log.warn("复核文档插入结果失败，保留已上传文件, docId={}, fileUrl={}", docId, fileUrl, e);
1010+
return;
1011+
}
1012+
deleteStoredFileQuietly(fileUrl);
1013+
}
1014+
1015+
private void deleteStoredFileQuietly(String fileUrl) {
1016+
if (!StringUtils.hasText(fileUrl)) {
1017+
return;
1018+
}
1019+
try {
1020+
fileStorageService.deleteByUrl(fileUrl);
1021+
} catch (Exception e) {
1022+
log.warn("补偿删除上传文件失败, fileUrl={}", fileUrl, e);
1023+
}
1024+
}
9881025
}

‎bootstrap/src/test/java/com/nageoffer/ai/ragent/knowledge/service/impl/KnowledgeDocumentServiceImplUploadTest.java‎

Lines changed: 330 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,330 @@
1+
/*
2+
* Licensed to the Apache Software Foundation (ASF) under one or more
3+
* contributor license agreements. See the NOTICE file distributed with
4+
* this work for additional information regarding copyright ownership.
5+
* The ASF licenses this file to You under the Apache License, Version 2.0
6+
* (the "License"); you may not use this file except in compliance with
7+
* the License. You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
18+
package com.nageoffer.ai.ragent.knowledge.service.impl;
19+
20+
import com.fasterxml.jackson.databind.ObjectMapper;
21+
import com.nageoffer.ai.ragent.audit.support.BizChangeLogContext;
22+
import com.nageoffer.ai.ragent.core.chunk.ChunkEmbeddingService;
23+
import com.nageoffer.ai.ragent.core.chunk.StructuredChunkingService;
24+
import com.nageoffer.ai.ragent.core.parser.DocumentParser;
25+
import com.nageoffer.ai.ragent.core.parser.DocumentParserSelector;
26+
import com.nageoffer.ai.ragent.framework.exception.ClientException;
27+
import com.nageoffer.ai.ragent.framework.mq.producer.MessageQueueProducer;
28+
import com.nageoffer.ai.ragent.ingestion.dao.mapper.IngestionPipelineMapper;
29+
import com.nageoffer.ai.ragent.ingestion.engine.IngestionEngine;
30+
import com.nageoffer.ai.ragent.ingestion.service.IngestionPipelineService;
31+
import com.nageoffer.ai.ragent.knowledge.config.KnowledgeScheduleProperties;
32+
import com.nageoffer.ai.ragent.knowledge.controller.request.KnowledgeDocumentUploadRequest;
33+
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeBaseDO;
34+
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeDocumentDO;
35+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeBaseMapper;
36+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeChunkMapper;
37+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeDocumentChunkLogMapper;
38+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeDocumentMapper;
39+
import com.nageoffer.ai.ragent.knowledge.handler.RemoteFileFetcher;
40+
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeChunkService;
41+
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeDocumentScheduleService;
42+
import com.nageoffer.ai.ragent.rag.core.vector.VectorStoreService;
43+
import com.nageoffer.ai.ragent.rag.dto.StoredFileDTO;
44+
import com.nageoffer.ai.ragent.rag.service.FileStorageService;
45+
import org.junit.jupiter.api.BeforeEach;
46+
import org.junit.jupiter.api.Test;
47+
import org.junit.jupiter.api.extension.ExtendWith;
48+
import org.mockito.Mock;
49+
import org.mockito.junit.jupiter.MockitoExtension;
50+
import org.springframework.transaction.support.TransactionOperations;
51+
import org.springframework.web.multipart.MultipartFile;
52+
53+
import static org.junit.jupiter.api.Assertions.assertSame;
54+
import static org.junit.jupiter.api.Assertions.assertThrows;
55+
import static org.mockito.ArgumentMatchers.any;
56+
import static org.mockito.Mockito.doAnswer;
57+
import static org.mockito.Mockito.doThrow;
58+
import static org.mockito.Mockito.lenient;
59+
import static org.mockito.Mockito.never;
60+
import static org.mockito.Mockito.verify;
61+
import static org.mockito.Mockito.verifyNoInteractions;
62+
import static org.mockito.Mockito.when;
63+
64+
@ExtendWith(MockitoExtension.class)
65+
class KnowledgeDocumentServiceImplUploadTest {
66+
67+
private static final String KB_ID = "kb-1";
68+
private static final String COLLECTION_NAME = "collection-1";
69+
private static final String FILE_URL = "collection-1/document.pdf";
70+
71+
@Mock private KnowledgeBaseMapper knowledgeBaseMapper;
72+
@Mock private KnowledgeDocumentMapper documentMapper;
73+
@Mock private DocumentParserSelector parserSelector;
74+
@Mock private StructuredChunkingService structuredChunkingService;
75+
@Mock private FileStorageService fileStorageService;
76+
@Mock private VectorStoreService vectorStoreService;
77+
@Mock private KnowledgeChunkService knowledgeChunkService;
78+
@Mock private KnowledgeDocumentScheduleService scheduleService;
79+
@Mock private IngestionPipelineService ingestionPipelineService;
80+
@Mock private IngestionPipelineMapper ingestionPipelineMapper;
81+
@Mock private IngestionEngine ingestionEngine;
82+
@Mock private ChunkEmbeddingService chunkEmbeddingService;
83+
@Mock private KnowledgeDocumentChunkLogMapper chunkLogMapper;
84+
@Mock private KnowledgeChunkMapper chunkMapper;
85+
@Mock private TransactionOperations transactionOperations;
86+
@Mock private MessageQueueProducer messageQueueProducer;
87+
@Mock private KnowledgeScheduleProperties scheduleProperties;
88+
@Mock private RemoteFileFetcher remoteFileFetcher;
89+
@Mock private BizChangeLogContext bizChangeLogContext;
90+
@Mock private MultipartFile file;
91+
@Mock private DocumentParser documentParser;
92+
93+
private KnowledgeDocumentServiceImpl service;
94+
95+
@BeforeEach
96+
void setUp() {
97+
service = new KnowledgeDocumentServiceImpl(
98+
knowledgeBaseMapper,
99+
documentMapper,
100+
parserSelector,
101+
structuredChunkingService,
102+
fileStorageService,
103+
vectorStoreService,
104+
knowledgeChunkService,
105+
new ObjectMapper(),
106+
scheduleService,
107+
ingestionPipelineService,
108+
ingestionPipelineMapper,
109+
ingestionEngine,
110+
chunkEmbeddingService,
111+
chunkLogMapper,
112+
chunkMapper,
113+
transactionOperations,
114+
messageQueueProducer,
115+
scheduleProperties,
116+
remoteFileFetcher,
117+
bizChangeLogContext
118+
);
119+
when(knowledgeBaseMapper.selectById(KB_ID)).thenReturn(KnowledgeBaseDO.builder()
120+
.id(KB_ID)
121+
.collectionName(COLLECTION_NAME)
122+
.build());
123+
}
124+
125+
@Test
126+
void invalidProcessModeConfigurationDoesNotUploadFile() {
127+
KnowledgeDocumentUploadRequest request = pipelineRequestWithoutId();
128+
lenient().when(fileStorageService.upload(COLLECTION_NAME, file)).thenReturn(storedFile());
129+
lenient().when(parserSelector.selectByMimeType("application/pdf")).thenReturn(documentParser);
130+
131+
assertThrows(ClientException.class, () -> service.upload(KB_ID, request, file));
132+
133+
verify(fileStorageService, never()).upload(COLLECTION_NAME, file);
134+
}
135+
136+
@Test
137+
void invalidUrlProcessModeConfigurationDoesNotFetchRemoteFile() {
138+
KnowledgeDocumentUploadRequest request = pipelineRequestWithoutId();
139+
request.setSourceType("url");
140+
request.setSourceLocation("https://example.com/document.pdf");
141+
142+
assertThrows(ClientException.class, () -> service.upload(KB_ID, request, null));
143+
144+
verifyNoInteractions(remoteFileFetcher);
145+
}
146+
147+
@Test
148+
void insertFailureDeletesUploadedFileAndRethrowsOriginalException() {
149+
RuntimeException original = new RuntimeException("insert failed");
150+
stubSuccessfulUploadPreparation();
151+
stubInsertFailureWithAssignedId(original);
152+
when(documentMapper.selectById("doc-1")).thenReturn(null);
153+
154+
RuntimeException thrown = assertThrows(
155+
RuntimeException.class,
156+
() -> service.upload(KB_ID, chunkRequest(), file)
157+
);
158+
159+
assertSame(original, thrown);
160+
verify(fileStorageService).deleteByUrl(FILE_URL);
161+
verify(documentMapper).selectById("doc-1");
162+
}
163+
164+
@Test
165+
void cleanupFailureDoesNotReplaceInsertFailure() {
166+
RuntimeException original = new RuntimeException("insert failed");
167+
stubSuccessfulUploadPreparation();
168+
stubInsertFailureWithAssignedId(original);
169+
when(documentMapper.selectById("doc-1")).thenReturn(null);
170+
doThrow(new RuntimeException("delete failed")).when(fileStorageService).deleteByUrl(FILE_URL);
171+
172+
RuntimeException thrown = assertThrows(
173+
RuntimeException.class,
174+
() -> service.upload(KB_ID, chunkRequest(), file)
175+
);
176+
177+
assertSame(original, thrown);
178+
verify(fileStorageService).deleteByUrl(FILE_URL);
179+
verify(documentMapper).selectById("doc-1");
180+
}
181+
182+
@Test
183+
void insertExceptionDoesNotDeleteFileWhenDocumentExists() {
184+
RuntimeException original = new RuntimeException("insert result unknown");
185+
stubSuccessfulUploadPreparation();
186+
stubInsertFailureWithAssignedId(original);
187+
when(documentMapper.selectById("doc-1")).thenReturn(KnowledgeDocumentDO.builder().id("doc-1").build());
188+
189+
RuntimeException thrown = assertThrows(
190+
RuntimeException.class,
191+
() -> service.upload(KB_ID, chunkRequest(), file)
192+
);
193+
194+
assertSame(original, thrown);
195+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
196+
}
197+
198+
@Test
199+
void insertExceptionWithoutAssignedIdKeepsFileAndOriginalException() {
200+
RuntimeException original = new RuntimeException("insert failed before id assignment");
201+
stubSuccessfulUploadPreparation();
202+
when(documentMapper.insert(any(KnowledgeDocumentDO.class))).thenThrow(original);
203+
204+
RuntimeException thrown = assertThrows(
205+
RuntimeException.class,
206+
() -> service.upload(KB_ID, chunkRequest(), file)
207+
);
208+
209+
assertSame(original, thrown);
210+
verify(documentMapper, never()).selectById(any());
211+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
212+
}
213+
214+
@Test
215+
void verificationFailureDoesNotDeleteFileOrReplaceInsertFailure() {
216+
RuntimeException original = new RuntimeException("insert result unknown");
217+
stubSuccessfulUploadPreparation();
218+
stubInsertFailureWithAssignedId(original);
219+
when(documentMapper.selectById("doc-1")).thenThrow(new RuntimeException("query failed"));
220+
221+
RuntimeException thrown = assertThrows(
222+
RuntimeException.class,
223+
() -> service.upload(KB_ID, chunkRequest(), file)
224+
);
225+
226+
assertSame(original, thrown);
227+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
228+
}
229+
230+
@Test
231+
void successfulUploadDoesNotDeleteStoredFile() {
232+
stubSuccessfulUploadPreparation();
233+
doAnswer(invocation -> {
234+
KnowledgeDocumentDO document = invocation.getArgument(0);
235+
document.setId("doc-1");
236+
return 1;
237+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
238+
239+
service.upload(KB_ID, chunkRequest(), file);
240+
241+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
242+
}
243+
244+
@Test
245+
void zeroInsertResultDeletesFileOnlyAfterConfirmingDocumentIsAbsent() {
246+
stubSuccessfulUploadPreparation();
247+
doAnswer(invocation -> {
248+
KnowledgeDocumentDO document = invocation.getArgument(0);
249+
document.setId("doc-1");
250+
return 0;
251+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
252+
when(documentMapper.selectById("doc-1")).thenReturn(null);
253+
254+
assertThrows(
255+
ClientException.class,
256+
() -> service.upload(KB_ID, chunkRequest(), file)
257+
);
258+
259+
verify(documentMapper).selectById("doc-1");
260+
verify(fileStorageService).deleteByUrl(FILE_URL);
261+
}
262+
263+
@Test
264+
void auditFailureAfterInsertDoesNotDeleteStoredFile() {
265+
RuntimeException auditFailure = new RuntimeException("audit failed");
266+
stubSuccessfulUploadPreparation();
267+
doAnswer(invocation -> {
268+
KnowledgeDocumentDO document = invocation.getArgument(0);
269+
document.setId("doc-1");
270+
return 1;
271+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
272+
doThrow(auditFailure).when(bizChangeLogContext).put(any(), any(), any());
273+
274+
RuntimeException thrown = assertThrows(
275+
RuntimeException.class,
276+
() -> service.upload(KB_ID, chunkRequest(), file)
277+
);
278+
279+
assertSame(auditFailure, thrown);
280+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
281+
}
282+
283+
@Test
284+
void unsupportedMimeTypeDeletesStoredFileExactlyOnce() {
285+
when(fileStorageService.upload(COLLECTION_NAME, file)).thenReturn(storedFile());
286+
when(parserSelector.selectByMimeType("application/pdf")).thenReturn(null);
287+
288+
assertThrows(ClientException.class, () -> service.upload(KB_ID, chunkRequest(), file));
289+
290+
verify(fileStorageService).deleteByUrl(FILE_URL);
291+
}
292+
293+
private void stubSuccessfulUploadPreparation() {
294+
when(fileStorageService.upload(COLLECTION_NAME, file)).thenReturn(storedFile());
295+
when(parserSelector.selectByMimeType("application/pdf")).thenReturn(documentParser);
296+
}
297+
298+
private void stubInsertFailureWithAssignedId(RuntimeException failure) {
299+
when(documentMapper.insert(any(KnowledgeDocumentDO.class))).thenAnswer(invocation -> {
300+
KnowledgeDocumentDO document = invocation.getArgument(0);
301+
document.setId("doc-1");
302+
throw failure;
303+
});
304+
}
305+
306+
private KnowledgeDocumentUploadRequest chunkRequest() {
307+
KnowledgeDocumentUploadRequest request = new KnowledgeDocumentUploadRequest();
308+
request.setSourceType("file");
309+
request.setProcessMode("chunk");
310+
request.setChunkStrategy("fixed_size");
311+
return request;
312+
}
313+
314+
private KnowledgeDocumentUploadRequest pipelineRequestWithoutId() {
315+
KnowledgeDocumentUploadRequest request = new KnowledgeDocumentUploadRequest();
316+
request.setSourceType("file");
317+
request.setProcessMode("pipeline");
318+
return request;
319+
}
320+
321+
private StoredFileDTO storedFile() {
322+
return StoredFileDTO.builder()
323+
.url(FILE_URL)
324+
.detectedType("pdf")
325+
.mimeType("application/pdf")
326+
.size(128L)
327+
.originalFilename("document.pdf")
328+
.build();
329+
}
330+
}

0 commit comments

Comments
 (0)