From 418171e6ddbe1f36584e1cacb736be6a9cdfb12b Mon Sep 17 00:00:00 2001 From: Sun Yuhan <1085481446@qq.com> Date: Wed, 14 May 2025 16:26:30 +0800 Subject: [PATCH 1/5] fix: Added transaction support for saveAll in JdbcChatMemoryRepository Signed-off-by: Sun Yuhan <1085481446@qq.com> --- .../jdbc/JdbcChatMemoryRepository.java | 35 +++++++++++++++---- 1 file changed, 28 insertions(+), 7 deletions(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java index 85d8c1c3265..dba49723e87 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java @@ -16,13 +16,11 @@ package org.springframework.ai.chat.memory.repository.jdbc; -import java.sql.PreparedStatement; -import java.sql.ResultSet; -import java.sql.SQLException; -import java.sql.Timestamp; +import java.sql.*; import java.time.Instant; import java.util.ArrayList; import java.util.List; +import java.util.Optional; import java.util.concurrent.atomic.AtomicLong; import org.springframework.ai.chat.memory.ChatMemoryRepository; @@ -35,6 +33,7 @@ import org.springframework.jdbc.core.BatchPreparedStatementSetter; import org.springframework.jdbc.core.JdbcTemplate; import org.springframework.jdbc.core.RowMapper; +import org.springframework.jdbc.datasource.DataSourceUtils; import org.springframework.lang.Nullable; import org.springframework.util.Assert; @@ -83,9 +82,31 @@ public void saveAll(String conversationId, List messages) { Assert.hasText(conversationId, "conversationId cannot be null or empty"); Assert.notNull(messages, "messages cannot be null"); Assert.noNullElements(messages, "messages cannot contain null elements"); - this.deleteByConversationId(conversationId); - this.jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), - new AddBatchPreparedStatement(conversationId, messages)); + + Connection connection = null; + Assert.notNull(jdbcTemplate.getDataSource(), "jdbcTemplate.getDataSource() cannot be null"); + try { + connection = DataSourceUtils.getConnection(jdbcTemplate.getDataSource()); + connection.setAutoCommit(false); + this.deleteByConversationId(conversationId); + this.jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), + new AddBatchPreparedStatement(conversationId, messages)); + connection.commit(); + } + catch (SQLException ex) { + try { + connection.rollback(); + } + catch (SQLException e) { + throw new RuntimeException("Transaction rollback exception", e); + } + throw new RuntimeException("save messages failed", ex); + } + finally { + Optional.ofNullable(connection) + .ifPresent(conn -> DataSourceUtils.releaseConnection(conn, jdbcTemplate.getDataSource())); + } + } @Override From 0a78b48619e34dc601bcc8be35e56f0af446bc01 Mon Sep 17 00:00:00 2001 From: Sun Yuhan <1085481446@qq.com> Date: Wed, 14 May 2025 17:30:47 +0800 Subject: [PATCH 2/5] fix: Adjust the transaction implementation method in JdbcChatMemoryRepository. Signed-off-by: Sun Yuhan <1085481446@qq.com> --- .../jdbc/JdbcChatMemoryRepository.java | 38 ++++++++----------- 1 file changed, 15 insertions(+), 23 deletions(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java index dba49723e87..b066c671203 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java @@ -20,7 +20,6 @@ import java.time.Instant; import java.util.ArrayList; import java.util.List; -import java.util.Optional; import java.util.concurrent.atomic.AtomicLong; import org.springframework.ai.chat.memory.ChatMemoryRepository; @@ -33,8 +32,9 @@ import org.springframework.jdbc.core.BatchPreparedStatementSetter; import org.springframework.jdbc.core.JdbcTemplate; import org.springframework.jdbc.core.RowMapper; -import org.springframework.jdbc.datasource.DataSourceUtils; +import org.springframework.jdbc.datasource.DataSourceTransactionManager; import org.springframework.lang.Nullable; +import org.springframework.transaction.support.TransactionTemplate; import org.springframework.util.Assert; /** @@ -83,30 +83,22 @@ public void saveAll(String conversationId, List messages) { Assert.notNull(messages, "messages cannot be null"); Assert.noNullElements(messages, "messages cannot contain null elements"); - Connection connection = null; - Assert.notNull(jdbcTemplate.getDataSource(), "jdbcTemplate.getDataSource() cannot be null"); - try { - connection = DataSourceUtils.getConnection(jdbcTemplate.getDataSource()); - connection.setAutoCommit(false); - this.deleteByConversationId(conversationId); - this.jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), - new AddBatchPreparedStatement(conversationId, messages)); - connection.commit(); - } - catch (SQLException ex) { + Assert.notNull(jdbcTemplate.getDataSource(), "dataSource can not be null"); + TransactionTemplate transactionTemplate = new TransactionTemplate( + new DataSourceTransactionManager(jdbcTemplate.getDataSource())); + + transactionTemplate.execute(status -> { try { - connection.rollback(); + deleteByConversationId(conversationId); + jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), + new AddBatchPreparedStatement(conversationId, messages)); } - catch (SQLException e) { - throw new RuntimeException("Transaction rollback exception", e); + catch (RuntimeException e) { + status.setRollbackOnly(); + throw e; } - throw new RuntimeException("save messages failed", ex); - } - finally { - Optional.ofNullable(connection) - .ifPresent(conn -> DataSourceUtils.releaseConnection(conn, jdbcTemplate.getDataSource())); - } - + return null; + }); } @Override From 8cce1720c575740237f38db6fb211c9de784d682 Mon Sep 17 00:00:00 2001 From: Sun Yuhan <1085481446@qq.com> Date: Wed, 14 May 2025 17:48:33 +0800 Subject: [PATCH 3/5] fix: Adjust the transaction implementation method in JdbcChatMemoryRepository. Signed-off-by: Sun Yuhan <1085481446@qq.com> --- .../memory/repository/jdbc/JdbcChatMemoryRepository.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java index b066c671203..41a2677f467 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java @@ -50,13 +50,17 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository { private final JdbcTemplate jdbcTemplate; + private final TransactionTemplate transactionTemplate; + private final JdbcChatMemoryRepositoryDialect dialect; private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryRepositoryDialect dialect) { Assert.notNull(jdbcTemplate, "jdbcTemplate cannot be null"); Assert.notNull(dialect, "dialect cannot be null"); + Assert.notNull(jdbcTemplate.getDataSource(), "dataSource can not be null"); this.jdbcTemplate = jdbcTemplate; this.dialect = dialect; + transactionTemplate = new TransactionTemplate(new DataSourceTransactionManager(jdbcTemplate.getDataSource())); } @Override @@ -83,10 +87,6 @@ public void saveAll(String conversationId, List messages) { Assert.notNull(messages, "messages cannot be null"); Assert.noNullElements(messages, "messages cannot contain null elements"); - Assert.notNull(jdbcTemplate.getDataSource(), "dataSource can not be null"); - TransactionTemplate transactionTemplate = new TransactionTemplate( - new DataSourceTransactionManager(jdbcTemplate.getDataSource())); - transactionTemplate.execute(status -> { try { deleteByConversationId(conversationId); From a610007b7c8624a10a4f1167a42b2f9810db3022 Mon Sep 17 00:00:00 2001 From: Sun Yuhan <1085481446@qq.com> Date: Wed, 14 May 2025 18:03:30 +0800 Subject: [PATCH 4/5] fix: Adjust the transaction implementation method in JdbcChatMemoryRepository. Signed-off-by: Sun Yuhan <1085481446@qq.com> --- .../repository/jdbc/JdbcChatMemoryRepository.java | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java index 41a2677f467..50317652a6a 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java @@ -60,7 +60,7 @@ private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryReposi Assert.notNull(jdbcTemplate.getDataSource(), "dataSource can not be null"); this.jdbcTemplate = jdbcTemplate; this.dialect = dialect; - transactionTemplate = new TransactionTemplate(new DataSourceTransactionManager(jdbcTemplate.getDataSource())); + this.transactionTemplate = new TransactionTemplate(new DataSourceTransactionManager(jdbcTemplate.getDataSource())); } @Override @@ -88,15 +88,9 @@ public void saveAll(String conversationId, List messages) { Assert.noNullElements(messages, "messages cannot contain null elements"); transactionTemplate.execute(status -> { - try { - deleteByConversationId(conversationId); - jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), - new AddBatchPreparedStatement(conversationId, messages)); - } - catch (RuntimeException e) { - status.setRollbackOnly(); - throw e; - } + deleteByConversationId(conversationId); + jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(), + new AddBatchPreparedStatement(conversationId, messages)); return null; }); } From aefd7cf3dcc68da442556aaadfa9913bad08fc5d Mon Sep 17 00:00:00 2001 From: Sun Yuhan <1085481446@qq.com> Date: Wed, 14 May 2025 19:26:10 +0800 Subject: [PATCH 5/5] fix: Fix format. Signed-off-by: Sun Yuhan <1085481446@qq.com> --- .../chat/memory/repository/jdbc/JdbcChatMemoryRepository.java | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java index 50317652a6a..68df4791949 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java @@ -60,7 +60,8 @@ private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryReposi Assert.notNull(jdbcTemplate.getDataSource(), "dataSource can not be null"); this.jdbcTemplate = jdbcTemplate; this.dialect = dialect; - this.transactionTemplate = new TransactionTemplate(new DataSourceTransactionManager(jdbcTemplate.getDataSource())); + this.transactionTemplate = new TransactionTemplate( + new DataSourceTransactionManager(jdbcTemplate.getDataSource())); } @Override