|
|
@@ -0,0 +1,275 @@
|
|
|
+package com.zsjz.ai.module.aiclean.exec;
|
|
|
+
|
|
|
+import com.zsjz.ai.module.aiclean.entity.AiCleanField;
|
|
|
+import com.zsjz.ai.module.aiclean.rule.CompiledColumn;
|
|
|
+import com.zaxxer.hikari.HikariConfig;
|
|
|
+import com.zaxxer.hikari.HikariDataSource;
|
|
|
+import org.junit.jupiter.api.AfterEach;
|
|
|
+import org.junit.jupiter.api.BeforeEach;
|
|
|
+import org.junit.jupiter.api.DisplayName;
|
|
|
+import org.junit.jupiter.api.Test;
|
|
|
+import org.junit.jupiter.api.io.TempDir;
|
|
|
+
|
|
|
+import java.nio.file.Path;
|
|
|
+import java.sql.Connection;
|
|
|
+import java.sql.PreparedStatement;
|
|
|
+import java.sql.Statement;
|
|
|
+import java.util.ArrayList;
|
|
|
+import java.util.LinkedHashMap;
|
|
|
+import java.util.List;
|
|
|
+import java.util.Map;
|
|
|
+
|
|
|
+import static org.junit.jupiter.api.Assertions.assertEquals;
|
|
|
+import static org.junit.jupiter.api.Assertions.assertNotNull;
|
|
|
+import static org.junit.jupiter.api.Assertions.assertNull;
|
|
|
+
|
|
|
+/**
|
|
|
+ * 预览表 {@code raw_{sheetId}} 作为 AI 取数源的读取契约，跑在真 DuckDB 上。
|
|
|
+ *
|
|
|
+ * <p>这段逻辑承担「不再重读源文件」的正确性：表头以上必须切掉、空行必须跳过、
|
|
|
+ * 物理行号必须原样带到 AI 动态表、分段读取不能丢行或重复。
|
|
|
+ * mock 一层 ResultSet 证明不了这些，所以在临时库里真建表真查。
|
|
|
+ *
|
|
|
+ * @author cc
|
|
|
+ * @since 2026/10/22
|
|
|
+ */
|
|
|
+class AiDuckDbRawTest {
|
|
|
+
|
|
|
+ @TempDir
|
|
|
+ Path dir;
|
|
|
+
|
|
|
+ private HikariDataSource ds;
|
|
|
+
|
|
|
+ @BeforeEach
|
|
|
+ void setUp() {
|
|
|
+ HikariConfig cfg = new HikariConfig();
|
|
|
+ cfg.setJdbcUrl("jdbc:duckdb:" + dir.resolve("case.duckdb"));
|
|
|
+ cfg.setDriverClassName("org.duckdb.DuckDBDriver");
|
|
|
+ // 写入器与预览表读取器各占一条连接（模拟真实入库：一边读 raw、一边写 ai_clean）
|
|
|
+ cfg.setMaximumPoolSize(4);
|
|
|
+ ds = new HikariDataSource(cfg);
|
|
|
+ AiDuckDb.useDataSource(ds);
|
|
|
+ }
|
|
|
+
|
|
|
+ @AfterEach
|
|
|
+ void tearDown() {
|
|
|
+ AiDuckDb.useDataSource(null);
|
|
|
+ try {
|
|
|
+ ds.close();
|
|
|
+ } catch (Exception ignored) {
|
|
|
+ // 临时库释放失败不影响断言
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ /** 建预览表：file_id / sheet_id / row_no + col_0..col_{n-1}，全 VARCHAR（与 PreDataListener 一致） */
|
|
|
+ private void createRaw(long sheetId, int colCount, boolean withRowNo) throws Exception {
|
|
|
+ StringBuilder ddl = new StringBuilder("CREATE TABLE raw_" + sheetId
|
|
|
+ + " (file_id VARCHAR, sheet_id VARCHAR");
|
|
|
+ if (withRowNo) {
|
|
|
+ ddl.append(", row_no INTEGER");
|
|
|
+ }
|
|
|
+ for (int i = 0; i < colCount; i++) {
|
|
|
+ ddl.append(", col_").append(i).append(" VARCHAR");
|
|
|
+ }
|
|
|
+ ddl.append(")");
|
|
|
+ try (Connection c = ds.getConnection(); Statement st = c.createStatement()) {
|
|
|
+ st.execute(ddl.toString());
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ private void insertRaw(long sheetId, int colCount, Integer rowNo, String... values) throws Exception {
|
|
|
+ StringBuilder sql = new StringBuilder("INSERT INTO raw_" + sheetId + " VALUES (?, ?");
|
|
|
+ if (rowNo != null) {
|
|
|
+ sql.append(", ?");
|
|
|
+ }
|
|
|
+ for (int i = 0; i < colCount; i++) {
|
|
|
+ sql.append(", ?");
|
|
|
+ }
|
|
|
+ sql.append(")");
|
|
|
+ try (Connection c = ds.getConnection(); PreparedStatement ps = c.prepareStatement(sql.toString())) {
|
|
|
+ int idx = 1;
|
|
|
+ ps.setString(idx++, "1001");
|
|
|
+ ps.setString(idx++, String.valueOf(sheetId));
|
|
|
+ if (rowNo != null) {
|
|
|
+ ps.setInt(idx++, rowNo);
|
|
|
+ }
|
|
|
+ for (int i = 0; i < colCount; i++) {
|
|
|
+ ps.setString(idx++, i < values.length ? values[i] : null);
|
|
|
+ }
|
|
|
+ ps.executeUpdate();
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("老结构预览表（没有 row_no）不可用 → 调用方回退读源文件")
|
|
|
+ void rejectsOldSchema() throws Exception {
|
|
|
+ createRaw(1L, 2, false);
|
|
|
+ insertRaw(1L, 2, null, "a", "b");
|
|
|
+
|
|
|
+ assertNull(AiDuckDb.rawSheet(1L, 1));
|
|
|
+ assertNull(AiDuckDb.readRawSample(1L, 0, 10, 1));
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("空预览表不可用：预览写入失败不能被当成「源文件没有数据」")
|
|
|
+ void rejectsEmptyTable() throws Exception {
|
|
|
+ createRaw(2L, 2, true);
|
|
|
+
|
|
|
+ assertNull(AiDuckDb.rawSheet(2L, 1));
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("表不存在返回 null，不抛异常")
|
|
|
+ void rejectsMissingTable() {
|
|
|
+ assertNull(AiDuckDb.rawSheet(404L, 1));
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("半截预览表不可用：写库中断后行数与上传记录对不上，不能当全量")
|
|
|
+ void rejectsPartialTable() throws Exception {
|
|
|
+ long sheetId = 7L;
|
|
|
+ createRaw(sheetId, 1, true);
|
|
|
+ insertRaw(sheetId, 1, 1, "表头");
|
|
|
+ insertRaw(sheetId, 1, 2, "张三");
|
|
|
+
|
|
|
+ assertNull(AiDuckDb.rawSheet(sheetId, 5), "只写进去 2 行、上传记录 5 行 → 回退读源文件");
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("切掉表头以上、跳过空行、物理行号原样回填（含空行造成的行号缺口）")
|
|
|
+ void skipsHeaderAndBlankRowsKeepingPhysicalRowNo() throws Exception {
|
|
|
+ long sheetId = 3L;
|
|
|
+ createRaw(sheetId, 3, true);
|
|
|
+ insertRaw(sheetId, 3, 1, "某某银行流水", null, null); // 标题行
|
|
|
+ insertRaw(sheetId, 3, 2, "姓名", "金额", "账号"); // 表头
|
|
|
+ insertRaw(sheetId, 3, 3, "张三", "100.00", "6222"); // 数据
|
|
|
+ insertRaw(sheetId, 3, 5, null, null, null); // 空行（Fesod 不会跳的话也会存在）
|
|
|
+ insertRaw(sheetId, 3, 6, "李四", "2,300.50", "6223"); // 数据（4 号空行没有落库 → 行号缺口）
|
|
|
+ insertRaw(sheetId, 3, 10, "王五", "5.00", "6224");
|
|
|
+
|
|
|
+ AiDuckDb.RawSheet raw = AiDuckDb.rawSheet(sheetId, 6);
|
|
|
+ assertNotNull(raw);
|
|
|
+ assertEquals(3, raw.colCount());
|
|
|
+ assertEquals(10L, raw.maxRowNo());
|
|
|
+
|
|
|
+ List<Integer> lineNos = new ArrayList<>();
|
|
|
+ List<String> firstCells = new ArrayList<>();
|
|
|
+ AiDuckDb.forEachRawRow(raw, 2, (row, lineNo) -> {
|
|
|
+ lineNos.add(lineNo);
|
|
|
+ firstCells.add(row.get(0));
|
|
|
+ });
|
|
|
+
|
|
|
+ assertEquals(List.of(3, 6, 10), lineNos, "表头以上与整空行都要跳过，行号保持原文件位置");
|
|
|
+ assertEquals(List.of("张三", "李四", "王五"), firstCells);
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("分段读取覆盖全部行：不丢、不重、顺序稳定")
|
|
|
+ void pagingCoversAllRows() throws Exception {
|
|
|
+ long sheetId = 4L;
|
|
|
+ int rows = 12_345;
|
|
|
+ createRaw(sheetId, 2, true);
|
|
|
+ try (Connection c = ds.getConnection();
|
|
|
+ PreparedStatement ps = c.prepareStatement(
|
|
|
+ "INSERT INTO raw_" + sheetId + " VALUES (?, ?, ?, ?, ?)")) {
|
|
|
+ for (int i = 1; i <= rows; i++) {
|
|
|
+ ps.setString(1, "1001");
|
|
|
+ ps.setString(2, String.valueOf(sheetId));
|
|
|
+ ps.setInt(3, i);
|
|
|
+ ps.setString(4, "v" + i);
|
|
|
+ ps.setString(5, "x" + i);
|
|
|
+ ps.addBatch();
|
|
|
+ if (i % 2000 == 0) {
|
|
|
+ ps.executeBatch();
|
|
|
+ }
|
|
|
+ }
|
|
|
+ ps.executeBatch();
|
|
|
+ }
|
|
|
+
|
|
|
+ AiDuckDb.RawSheet raw = AiDuckDb.rawSheet(sheetId, rows);
|
|
|
+ assertNotNull(raw);
|
|
|
+ List<Integer> lineNos = new ArrayList<>();
|
|
|
+ AiDuckDb.forEachRawRow(raw, 0, (row, lineNo) -> lineNos.add(lineNo));
|
|
|
+
|
|
|
+ assertEquals(rows, lineNos.size(), "分段读取不能丢行");
|
|
|
+ assertEquals(1, lineNos.getFirst());
|
|
|
+ assertEquals(rows, lineNos.getLast());
|
|
|
+ for (int i = 1; i < lineNos.size(); i++) {
|
|
|
+ assertEquals(lineNos.get(i - 1) + 1, lineNos.get(i), "行序必须严格递增且不重复");
|
|
|
+ }
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("抽样只取前 N 条非空行")
|
|
|
+ void sampleStopsAtLimit() throws Exception {
|
|
|
+ long sheetId = 5L;
|
|
|
+ createRaw(sheetId, 1, true);
|
|
|
+ for (int i = 1; i <= 10; i++) {
|
|
|
+ insertRaw(sheetId, 1, i, "v" + i);
|
|
|
+ }
|
|
|
+
|
|
|
+ List<Map<Integer, String>> sample = AiDuckDb.readRawSample(sheetId, 0, 3, 10);
|
|
|
+ assertNotNull(sample);
|
|
|
+ assertEquals(3, sample.size());
|
|
|
+ assertEquals("v1", sample.getFirst().get(0));
|
|
|
+ }
|
|
|
+
|
|
|
+ @Test
|
|
|
+ @DisplayName("边读预览表边写 AI 表：清洗结果与行号都落库")
|
|
|
+ void loadFromRawWhileWriting() throws Exception {
|
|
|
+ long sheetId = 6L;
|
|
|
+ createRaw(sheetId, 2, true);
|
|
|
+ insertRaw(sheetId, 2, 1, "姓名", "金额");
|
|
|
+ insertRaw(sheetId, 2, 2, "张三", "100.00");
|
|
|
+ insertRaw(sheetId, 2, 4, "李四", "200.00");
|
|
|
+
|
|
|
+ AiCleanField f = new AiCleanField();
|
|
|
+ f.setColumnName("c0");
|
|
|
+ f.setFieldNameCn("姓名");
|
|
|
+ f.setFieldNameEn("c0");
|
|
|
+ f.setFieldType("TEXT");
|
|
|
+ f.setRequired(0);
|
|
|
+ f.setFieldLen(0);
|
|
|
+ f.setFieldSort(1);
|
|
|
+ f.setColIndex(0);
|
|
|
+ List<AiCleanField> fields = List.of(f);
|
|
|
+ assertNull(AiDuckDb.createAiTable("ai_clean_raw_test", fields));
|
|
|
+
|
|
|
+ CompiledColumn col = CompiledColumn.of(0, "c0", "c0", "TEXT", false, 0, List.of());
|
|
|
+ AiDuckDb.RawSheet raw = AiDuckDb.rawSheet(sheetId, 3);
|
|
|
+ assertNotNull(raw);
|
|
|
+
|
|
|
+ List<Integer> storedRowNos = new ArrayList<>();
|
|
|
+ long written;
|
|
|
+ try (AiDuckDb.BatchWriter writer = AiDuckDb.openWriter("ai_clean_raw_test", List.of("c0"))) {
|
|
|
+ AiDuckDb.forEachRawRow(raw, 1, (row, lineNo) -> {
|
|
|
+ Map<String, Object> cleaned = AiRecordCleaner.cleanRow(row, List.of(col), new LinkedHashMap<>());
|
|
|
+ assertNotNull(cleaned);
|
|
|
+ Map<String, Object> toWrite = new LinkedHashMap<>(cleaned);
|
|
|
+ toWrite.put("id", (long) lineNo);
|
|
|
+ toWrite.put("case_id", 9L);
|
|
|
+ toWrite.put("file_id", 1001L);
|
|
|
+ toWrite.put("sheet_id", sheetId);
|
|
|
+ toWrite.put("block_no", 0);
|
|
|
+ toWrite.put("row_no", lineNo);
|
|
|
+ toWrite.put("ai_job_id", 88L);
|
|
|
+ toWrite.put("ai_rule_ver", 1);
|
|
|
+ toWrite.put("category", "PERSON");
|
|
|
+ storedRowNos.add(lineNo);
|
|
|
+ writer.add(toWrite);
|
|
|
+ });
|
|
|
+ written = writer.finish();
|
|
|
+ }
|
|
|
+
|
|
|
+ assertEquals(2L, written, "两行数据都要落库（整空行跳过，但没有空行）");
|
|
|
+ assertEquals(List.of(2, 4), storedRowNos, "原文件物理行号要原样传给 AI 表");
|
|
|
+ assertEquals(2L, AiDuckDb.count("ai_clean_raw_test"));
|
|
|
+
|
|
|
+ List<Map<String, Object>> page = AiDuckDb.selectPage("ai_clean_raw_test",
|
|
|
+ List.of("c0", "row_no"), "row_no", true, 1, 10);
|
|
|
+ assertEquals(2, page.size());
|
|
|
+ assertEquals(2, ((Number) page.get(0).get("row_no")).intValue());
|
|
|
+ assertEquals("张三", page.get(0).get("c0"));
|
|
|
+ assertEquals(4, ((Number) page.get(1).get("row_no")).intValue());
|
|
|
+ }
|
|
|
+}
|
