Skip to content

Commit 118c40a

Browse files
KDB-WindKDB-Wind
committed
fix(knowledge): 文档上传失败时清理存储文件
对象写入成功后，文档插入异常或返回 0 会留下没有数据库记录的文件。插入失败时按已生成 ID 复核记录，仅在确认不存在时删除；无法确认提交结果或清理失败时保留原异常。
1 parent b7d0c8d commit 118c40a

2 files changed

Lines changed: 367 additions & 1 deletion

File tree

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

Lines changed: 38 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -177,7 +177,15 @@ public KnowledgeDocumentVO upload(String kbId, KnowledgeDocumentUploadRequest re
177177
.createdBy(UserContext.getUsername())
178178
.updatedBy(UserContext.getUsername())
179179
.build();
180-
documentMapper.insert(documentDO);
180+
try {
181+
int inserted = documentMapper.insert(documentDO);
182+
if (inserted <= 0) {
183+
throw new ClientException("文档保存失败");
184+
}
185+
} catch (RuntimeException e) {
186+
deleteStoredFileIfDocumentAbsent(documentDO, stored.getUrl());
187+
throw e;
188+
}
181189
bizChangeLogContext.put(String.valueOf(documentDO.getId()), null, documentDO);
182190
bizChangeLogContext.putName(documentDO.getDocName());
183191

@@ -853,4 +861,33 @@ private void deleteStoredFileQuietly(KnowledgeDocumentDO documentDO) {
853861
log.warn("删除文档存储文件失败, docId={}, fileUrl={}", documentDO.getId(), documentDO.getFileUrl(), e);
854862
}
855863
}
864+
865+
private void deleteStoredFileIfDocumentAbsent(KnowledgeDocumentDO documentDO, String fileUrl) {
866+
String docId = documentDO.getId();
867+
if (!StringUtils.hasText(docId)) {
868+
log.warn("无法确认文档插入结果，保留已上传文件, fileUrl={}", fileUrl);
869+
return;
870+
}
871+
try {
872+
if (documentMapper.selectById(docId) != null) {
873+
log.warn("文档插入结果异常但记录已存在，保留已上传文件, docId={}, fileUrl={}", docId, fileUrl);
874+
return;
875+
}
876+
} catch (Exception e) {
877+
log.warn("复核文档插入结果失败，保留已上传文件, docId={}, fileUrl={}", docId, fileUrl, e);
878+
return;
879+
}
880+
deleteStoredFileQuietly(fileUrl);
881+
}
882+
883+
private void deleteStoredFileQuietly(String fileUrl) {
884+
if (!StringUtils.hasText(fileUrl)) {
885+
return;
886+
}
887+
try {
888+
fileStorageService.deleteByUrl(fileUrl);
889+
} catch (Exception e) {
890+
log.warn("补偿删除上传文件失败, fileUrl={}", fileUrl, e);
891+
}
892+
}
856893
}

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

Lines changed: 329 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,329 @@
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.ingest.IngestionKernel;
23+
import com.nageoffer.ai.ragent.core.ingest.sink.ChunkIndexWriter;
24+
import com.nageoffer.ai.ragent.core.parser.registry.ParserRegistry;
25+
import com.nageoffer.ai.ragent.framework.exception.ClientException;
26+
import com.nageoffer.ai.ragent.framework.mq.producer.MessageQueueProducer;
27+
import com.nageoffer.ai.ragent.ingestion.dao.mapper.IngestionPipelineMapper;
28+
import com.nageoffer.ai.ragent.ingestion.engine.IngestionEngine;
29+
import com.nageoffer.ai.ragent.ingestion.service.IngestionPipelineService;
30+
import com.nageoffer.ai.ragent.knowledge.config.KnowledgeScheduleProperties;
31+
import com.nageoffer.ai.ragent.knowledge.controller.request.KnowledgeDocumentUploadRequest;
32+
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeBaseDO;
33+
import com.nageoffer.ai.ragent.knowledge.dao.entity.KnowledgeDocumentDO;
34+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeBaseMapper;
35+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeChunkMapper;
36+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeDocumentChunkLogMapper;
37+
import com.nageoffer.ai.ragent.knowledge.dao.mapper.KnowledgeDocumentMapper;
38+
import com.nageoffer.ai.ragent.knowledge.handler.RemoteFileFetcher;
39+
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeChunkService;
40+
import com.nageoffer.ai.ragent.knowledge.service.KnowledgeDocumentScheduleService;
41+
import com.nageoffer.ai.ragent.knowledge.support.IngestionSpecCodec;
42+
import com.nageoffer.ai.ragent.knowledge.support.VectorTargetResolver;
43+
import com.nageoffer.ai.ragent.rag.core.vector.VectorStoreService;
44+
import com.nageoffer.ai.ragent.rag.dto.StoredFileDTO;
45+
import com.nageoffer.ai.ragent.rag.service.FileStorageService;
46+
import org.junit.jupiter.api.BeforeEach;
47+
import org.junit.jupiter.api.Test;
48+
import org.junit.jupiter.api.extension.ExtendWith;
49+
import org.mockito.Mock;
50+
import org.mockito.junit.jupiter.MockitoExtension;
51+
import org.springframework.transaction.support.TransactionOperations;
52+
import org.springframework.web.multipart.MultipartFile;
53+
54+
import static org.junit.jupiter.api.Assertions.assertSame;
55+
import static org.junit.jupiter.api.Assertions.assertThrows;
56+
import static org.mockito.ArgumentMatchers.any;
57+
import static org.mockito.Mockito.doAnswer;
58+
import static org.mockito.Mockito.doThrow;
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 ParserRegistry parserRegistry;
74+
@Mock private IngestionKernel ingestionKernel;
75+
@Mock private ChunkIndexWriter chunkIndexWriter;
76+
@Mock private IngestionSpecCodec ingestionSpecCodec;
77+
@Mock private FileStorageService fileStorageService;
78+
@Mock private VectorStoreService vectorStoreService;
79+
@Mock private KnowledgeChunkService knowledgeChunkService;
80+
@Mock private KnowledgeDocumentScheduleService scheduleService;
81+
@Mock private IngestionPipelineService ingestionPipelineService;
82+
@Mock private IngestionPipelineMapper ingestionPipelineMapper;
83+
@Mock private IngestionEngine ingestionEngine;
84+
@Mock private KnowledgeDocumentChunkLogMapper chunkLogMapper;
85+
@Mock private KnowledgeChunkMapper chunkMapper;
86+
@Mock private TransactionOperations transactionOperations;
87+
@Mock private MessageQueueProducer messageQueueProducer;
88+
@Mock private KnowledgeScheduleProperties scheduleProperties;
89+
@Mock private RemoteFileFetcher remoteFileFetcher;
90+
@Mock private VectorTargetResolver vectorTargetResolver;
91+
@Mock private BizChangeLogContext bizChangeLogContext;
92+
@Mock private MultipartFile file;
93+
94+
private KnowledgeDocumentServiceImpl service;
95+
96+
@BeforeEach
97+
void setUp() {
98+
service = new KnowledgeDocumentServiceImpl(
99+
knowledgeBaseMapper,
100+
documentMapper,
101+
parserRegistry,
102+
ingestionKernel,
103+
chunkIndexWriter,
104+
ingestionSpecCodec,
105+
fileStorageService,
106+
vectorStoreService,
107+
knowledgeChunkService,
108+
new ObjectMapper(),
109+
scheduleService,
110+
ingestionPipelineService,
111+
ingestionPipelineMapper,
112+
ingestionEngine,
113+
chunkLogMapper,
114+
chunkMapper,
115+
transactionOperations,
116+
messageQueueProducer,
117+
scheduleProperties,
118+
remoteFileFetcher,
119+
vectorTargetResolver,
120+
bizChangeLogContext
121+
);
122+
when(knowledgeBaseMapper.selectById(KB_ID)).thenReturn(KnowledgeBaseDO.builder()
123+
.id(KB_ID)
124+
.collectionName(COLLECTION_NAME)
125+
.build());
126+
}
127+
128+
@Test
129+
void invalidProcessModeConfigurationDoesNotUploadFile() {
130+
KnowledgeDocumentUploadRequest request = pipelineRequestWithoutId();
131+
132+
assertThrows(ClientException.class, () -> service.upload(KB_ID, request, file));
133+
134+
verify(fileStorageService, never()).upload(COLLECTION_NAME, file);
135+
}
136+
137+
@Test
138+
void invalidUrlProcessModeConfigurationDoesNotFetchRemoteFile() {
139+
KnowledgeDocumentUploadRequest request = pipelineRequestWithoutId();
140+
request.setSourceType("url");
141+
request.setSourceLocation("https://example.com/document.pdf");
142+
143+
assertThrows(ClientException.class, () -> service.upload(KB_ID, request, null));
144+
145+
verifyNoInteractions(remoteFileFetcher);
146+
}
147+
148+
@Test
149+
void insertFailureDeletesUploadedFileAndRethrowsOriginalException() {
150+
RuntimeException original = new RuntimeException("insert failed");
151+
stubSuccessfulUploadPreparation();
152+
stubInsertFailureWithAssignedId(original);
153+
when(documentMapper.selectById("doc-1")).thenReturn(null);
154+
155+
RuntimeException thrown = assertThrows(
156+
RuntimeException.class,
157+
() -> service.upload(KB_ID, chunkRequest(), file)
158+
);
159+
160+
assertSame(original, thrown);
161+
verify(fileStorageService).deleteByUrl(FILE_URL);
162+
verify(documentMapper).selectById("doc-1");
163+
}
164+
165+
@Test
166+
void cleanupFailureDoesNotReplaceInsertFailure() {
167+
RuntimeException original = new RuntimeException("insert failed");
168+
stubSuccessfulUploadPreparation();
169+
stubInsertFailureWithAssignedId(original);
170+
when(documentMapper.selectById("doc-1")).thenReturn(null);
171+
doThrow(new RuntimeException("delete failed")).when(fileStorageService).deleteByUrl(FILE_URL);
172+
173+
RuntimeException thrown = assertThrows(
174+
RuntimeException.class,
175+
() -> service.upload(KB_ID, chunkRequest(), file)
176+
);
177+
178+
assertSame(original, thrown);
179+
verify(fileStorageService).deleteByUrl(FILE_URL);
180+
verify(documentMapper).selectById("doc-1");
181+
}
182+
183+
@Test
184+
void insertExceptionDoesNotDeleteFileWhenDocumentExists() {
185+
RuntimeException original = new RuntimeException("insert result unknown");
186+
stubSuccessfulUploadPreparation();
187+
stubInsertFailureWithAssignedId(original);
188+
when(documentMapper.selectById("doc-1")).thenReturn(KnowledgeDocumentDO.builder().id("doc-1").build());
189+
190+
RuntimeException thrown = assertThrows(
191+
RuntimeException.class,
192+
() -> service.upload(KB_ID, chunkRequest(), file)
193+
);
194+
195+
assertSame(original, thrown);
196+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
197+
}
198+
199+
@Test
200+
void insertExceptionWithoutAssignedIdKeepsFileAndOriginalException() {
201+
RuntimeException original = new RuntimeException("insert failed before id assignment");
202+
stubSuccessfulUploadPreparation();
203+
when(documentMapper.insert(any(KnowledgeDocumentDO.class))).thenThrow(original);
204+
205+
RuntimeException thrown = assertThrows(
206+
RuntimeException.class,
207+
() -> service.upload(KB_ID, chunkRequest(), file)
208+
);
209+
210+
assertSame(original, thrown);
211+
verify(documentMapper, never()).selectById(any());
212+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
213+
}
214+
215+
@Test
216+
void verificationFailureDoesNotDeleteFileOrReplaceInsertFailure() {
217+
RuntimeException original = new RuntimeException("insert result unknown");
218+
stubSuccessfulUploadPreparation();
219+
stubInsertFailureWithAssignedId(original);
220+
when(documentMapper.selectById("doc-1")).thenThrow(new RuntimeException("query failed"));
221+
222+
RuntimeException thrown = assertThrows(
223+
RuntimeException.class,
224+
() -> service.upload(KB_ID, chunkRequest(), file)
225+
);
226+
227+
assertSame(original, thrown);
228+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
229+
}
230+
231+
@Test
232+
void successfulUploadDoesNotDeleteStoredFile() {
233+
stubSuccessfulUploadPreparation();
234+
doAnswer(invocation -> {
235+
KnowledgeDocumentDO document = invocation.getArgument(0);
236+
document.setId("doc-1");
237+
return 1;
238+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
239+
240+
service.upload(KB_ID, chunkRequest(), file);
241+
242+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
243+
}
244+
245+
@Test
246+
void zeroInsertResultDeletesFileOnlyAfterConfirmingDocumentIsAbsent() {
247+
stubSuccessfulUploadPreparation();
248+
doAnswer(invocation -> {
249+
KnowledgeDocumentDO document = invocation.getArgument(0);
250+
document.setId("doc-1");
251+
return 0;
252+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
253+
when(documentMapper.selectById("doc-1")).thenReturn(null);
254+
255+
assertThrows(
256+
ClientException.class,
257+
() -> service.upload(KB_ID, chunkRequest(), file)
258+
);
259+
260+
verify(documentMapper).selectById("doc-1");
261+
verify(fileStorageService).deleteByUrl(FILE_URL);
262+
}
263+
264+
@Test
265+
void auditFailureAfterInsertDoesNotDeleteStoredFile() {
266+
RuntimeException auditFailure = new RuntimeException("audit failed");
267+
stubSuccessfulUploadPreparation();
268+
doAnswer(invocation -> {
269+
KnowledgeDocumentDO document = invocation.getArgument(0);
270+
document.setId("doc-1");
271+
return 1;
272+
}).when(documentMapper).insert(any(KnowledgeDocumentDO.class));
273+
doThrow(auditFailure).when(bizChangeLogContext).put(any(), any(), any());
274+
275+
RuntimeException thrown = assertThrows(
276+
RuntimeException.class,
277+
() -> service.upload(KB_ID, chunkRequest(), file)
278+
);
279+
280+
assertSame(auditFailure, thrown);
281+
verify(fileStorageService, never()).deleteByUrl(FILE_URL);
282+
}
283+
284+
@Test
285+
void unsupportedMimeTypeDeletesStoredFileExactlyOnce() {
286+
when(fileStorageService.upload(COLLECTION_NAME, file)).thenReturn(storedFile());
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(parserRegistry.canParse("application/pdf")).thenReturn(true);
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+
return request;
311+
}
312+
313+
private KnowledgeDocumentUploadRequest pipelineRequestWithoutId() {
314+
KnowledgeDocumentUploadRequest request = new KnowledgeDocumentUploadRequest();
315+
request.setSourceType("file");
316+
request.setProcessMode("pipeline");
317+
return request;
318+
}
319+
320+
private StoredFileDTO storedFile() {
321+
return StoredFileDTO.builder()
322+
.url(FILE_URL)
323+
.detectedType("pdf")
324+
.mimeType("application/pdf")
325+
.size(128L)
326+
.originalFilename("document.pdf")
327+
.build();
328+
}
329+
}

0 commit comments

Comments
 (0)