fix: Added transaction support for saveAll in JdbcChatMemoryRepository

- Add optional PlatformTransactionManager
- Add tests

Signed-off-by: Sun Yuhan <[email protected]>
Signed-off-by: minsoo.nam <[email protected]>

Author: sunyuhan1998

memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/main/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepository.java

@@ -16,10 +16,7 @@
 
 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;
@@ -35,7 +32,10 @@
 import org.springframework.jdbc.core.BatchPreparedStatementSetter;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.core.RowMapper;
+import org.springframework.jdbc.datasource.DataSourceTransactionManager;
 import org.springframework.lang.Nullable;
+import org.springframework.transaction.PlatformTransactionManager;
+import org.springframework.transaction.support.TransactionTemplate;
 import org.springframework.util.Assert;
 
 /**
@@ -51,13 +51,18 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository {
 
 	private final JdbcTemplate jdbcTemplate;
 
+	private final TransactionTemplate transactionTemplate;
+
 	private final JdbcChatMemoryRepositoryDialect dialect;
 
-	private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryRepositoryDialect dialect) {
+	private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryRepositoryDialect dialect,
+			PlatformTransactionManager txManager) {
 		Assert.notNull(jdbcTemplate, "jdbcTemplate cannot be null");
 		Assert.notNull(dialect, "dialect cannot be null");
 		this.jdbcTemplate = jdbcTemplate;
 		this.dialect = dialect;
+		this.transactionTemplate = new TransactionTemplate(
+				txManager != null ? txManager : new DataSourceTransactionManager(jdbcTemplate.getDataSource()));
 	}
 
 	@Override
@@ -83,9 +88,13 @@ public void saveAll(String conversationId, List<Message> 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));
+
+		transactionTemplate.execute(status -> {
+			deleteByConversationId(conversationId);
+			jdbcTemplate.batchUpdate(dialect.getInsertMessageSql(),
+					new AddBatchPreparedStatement(conversationId, messages));
+			return null;
+		});
 	}
 
 	@Override
@@ -148,6 +157,8 @@ public static class Builder {
 
 		private JdbcChatMemoryRepositoryDialect dialect;
 
+		private PlatformTransactionManager platformTransactionManager;
+
 		private Builder() {
 		}
 
@@ -161,10 +172,15 @@ public Builder dialect(JdbcChatMemoryRepositoryDialect dialect) {
 			return this;
 		}
 
+		public Builder transactionManager(PlatformTransactionManager txManager) {
+			this.platformTransactionManager = txManager;
+			return this;
+		}
+
 		public JdbcChatMemoryRepository build() {
 			if (this.dialect == null)
 				throw new IllegalStateException("Dialect must be set");
-			return new JdbcChatMemoryRepository(this.jdbcTemplate, this.dialect);
+			return new JdbcChatMemoryRepository(this.jdbcTemplate, this.dialect, this.platformTransactionManager);
 		}
 
 	}

memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryPostgresqlIT.java

@@ -16,13 +16,19 @@
 
 package org.springframework.ai.chat.memory.repository.jdbc;
 
+import java.sql.Timestamp;
+import java.util.List;
+import java.util.UUID;
+import javax.sql.DataSource;
+
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
+
 import org.springframework.ai.chat.memory.ChatMemoryRepository;
+import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.Message;
 import org.springframework.ai.chat.messages.MessageType;
-import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.SystemMessage;
 import org.springframework.ai.chat.messages.UserMessage;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -33,14 +39,10 @@
 import org.springframework.boot.test.context.SpringBootTest;
 import org.springframework.context.annotation.Bean;
 import org.springframework.jdbc.core.JdbcTemplate;
+import org.springframework.jdbc.datasource.DataSourceTransactionManager;
 import org.springframework.test.context.TestPropertySource;
 import org.springframework.test.context.jdbc.Sql;
-
-import java.sql.Timestamp;
-import java.util.List;
-import java.util.UUID;
-
-import javax.sql.DataSource;
+import org.springframework.transaction.support.TransactionTemplate;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -156,6 +158,34 @@ void deleteMessagesByConversationId() {
 		assertThat(count).isZero();
 	}
 
+	@Test
+	void repositoryWithExplicitTransactionManager() {
+		// Get the repository with explicit transaction manager
+		ChatMemoryRepository repositoryWithTxManager = TestConfiguration
+			.chatMemoryRepositoryWithTransactionManager(jdbcTemplate, jdbcTemplate.getDataSource());
+
+		var conversationId = UUID.randomUUID().toString();
+		var messages = List.<Message>of(new AssistantMessage("Message with transaction manager - " + conversationId),
+				new UserMessage("User message with transaction manager - " + conversationId));
+
+		// Save messages using the repository with explicit transaction manager
+		repositoryWithTxManager.saveAll(conversationId, messages);
+
+		// Verify messages were saved correctly
+		var savedMessages = repositoryWithTxManager.findByConversationId(conversationId);
+		assertThat(savedMessages).hasSize(2);
+		assertThat(savedMessages).isEqualTo(messages);
+
+		// Verify transaction works by updating and checking atomicity
+		var newMessages = List.<Message>of(new SystemMessage("New system message - " + conversationId));
+		repositoryWithTxManager.saveAll(conversationId, newMessages);
+
+		// The old messages should be deleted and only the new one should exist
+		var updatedMessages = repositoryWithTxManager.findByConversationId(conversationId);
+		assertThat(updatedMessages).hasSize(1);
+		assertThat(updatedMessages).isEqualTo(newMessages);
+	}
+
 	@SpringBootConfiguration
 	@ImportAutoConfiguration({ DataSourceAutoConfiguration.class, JdbcTemplateAutoConfiguration.class })
 	static class TestConfiguration {
@@ -168,6 +198,20 @@ ChatMemoryRepository chatMemoryRepository(JdbcTemplate jdbcTemplate, DataSource
 				.build();
 		}
 
+		@Bean
+		ChatMemoryRepository chatMemoryRepositoryWithTxManager(JdbcTemplate jdbcTemplate, DataSource dataSource) {
+			return chatMemoryRepositoryWithTransactionManager(jdbcTemplate, dataSource);
+		}
+
+		static ChatMemoryRepository chatMemoryRepositoryWithTransactionManager(JdbcTemplate jdbcTemplate,
+				DataSource dataSource) {
+			return JdbcChatMemoryRepository.builder()
+				.jdbcTemplate(jdbcTemplate)
+				.dialect(JdbcChatMemoryRepositoryDialect.from(dataSource))
+				.transactionManager(new DataSourceTransactionManager(dataSource))
+				.build();
+		}
+
 	}
 
 }