From b83a162506ca273fb2e1a1dadf78571226ae569e Mon Sep 17 00:00:00 2001 From: Yanming Zhou Date: Fri, 16 May 2025 10:40:07 +0800 Subject: [PATCH] JdbcChatMemoryRepository should use the provided JdbcTemplate Before this commit, the underlying `JdbcTemplate` is created like `new JdbcTemplate(providedJdbcTemplate.getDataSource())`, it means that settings on provided `JdbcTemplate` will lose. Signed-off-by: Yanming Zhou --- .../AbstractJdbcChatMemoryRepositoryIT.java | 35 +++++---- .../jdbc/JdbcChatMemoryRepositoryMysqlIT.java | 12 +--- .../JdbcChatMemoryRepositoryPostgresqlIT.java | 71 +------------------ 3 files changed, 22 insertions(+), 96 deletions(-) diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/AbstractJdbcChatMemoryRepositoryIT.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/AbstractJdbcChatMemoryRepositoryIT.java index 000e828fc..269d91c67 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/AbstractJdbcChatMemoryRepositoryIT.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/AbstractJdbcChatMemoryRepositoryIT.java @@ -39,6 +39,14 @@ import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration; import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration; import org.springframework.context.annotation.Bean; import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.test.context.ContextConfiguration; + +import java.sql.Timestamp; +import java.util.List; +import java.util.UUID; +import java.util.stream.Collectors; + +import javax.sql.DataSource; import static org.assertj.core.api.Assertions.assertThat; @@ -46,20 +54,17 @@ import static org.assertj.core.api.Assertions.assertThat; * Base class for integration tests for {@link JdbcChatMemoryRepository}. * * @author Mark Pollack + * @author Yanming Zhou */ +@ContextConfiguration(classes = AbstractJdbcChatMemoryRepositoryIT.TestConfiguration.class) public abstract class AbstractJdbcChatMemoryRepositoryIT { @Autowired - protected ChatMemoryRepository chatMemoryRepository; + protected JdbcChatMemoryRepository chatMemoryRepository; @Autowired protected JdbcTemplate jdbcTemplate; - @Test - void correctChatMemoryRepositoryInstance() { - assertThat(this.chatMemoryRepository).isInstanceOf(ChatMemoryRepository.class); - } - @ParameterizedTest @CsvSource({ "Message from assistant,ASSISTANT", "Message from user,USER", "Message from system,SYSTEM" }) void saveMessagesSingleMessage(String content, MessageType messageType) { @@ -163,11 +168,6 @@ public abstract class AbstractJdbcChatMemoryRepositoryIT { @Test void testMessageOrder() { - // Create a repository using the from method to detect the dialect - JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder() - .jdbcTemplate(this.jdbcTemplate) - .dialect(JdbcChatMemoryRepositoryDialect.from(this.jdbcTemplate.getDataSource())) - .build(); var conversationId = UUID.randomUUID().toString(); @@ -179,10 +179,10 @@ public abstract class AbstractJdbcChatMemoryRepositoryIT { // Save messages in the expected order List orderedMessages = List.of(firstMessage, secondMessage, thirdMessage, fourthMessage); - repository.saveAll(conversationId, orderedMessages); + chatMemoryRepository.saveAll(conversationId, orderedMessages); // Retrieve messages using the repository - List retrievedMessages = repository.findByConversationId(conversationId); + List retrievedMessages = chatMemoryRepository.findByConversationId(conversationId); assertThat(retrievedMessages).hasSize(4); // Get the actual order from the retrieved messages @@ -197,14 +197,11 @@ public abstract class AbstractJdbcChatMemoryRepositoryIT { * Base configuration for all integration tests. */ @ImportAutoConfiguration({ DataSourceAutoConfiguration.class, JdbcTemplateAutoConfiguration.class }) - static abstract class BaseTestConfiguration { + static class TestConfiguration { @Bean - ChatMemoryRepository chatMemoryRepository(JdbcTemplate jdbcTemplate, DataSource dataSource) { - return JdbcChatMemoryRepository.builder() - .jdbcTemplate(jdbcTemplate) - .dialect(JdbcChatMemoryRepositoryDialect.from(dataSource)) - .build(); + ChatMemoryRepository chatMemoryRepository(DataSource dataSource) { + return JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); } } diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryMysqlIT.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryMysqlIT.java index 62707e160..3ee33e42b 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryMysqlIT.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryMysqlIT.java @@ -16,7 +16,6 @@ package org.springframework.ai.chat.memory.repository.jdbc; -import org.springframework.boot.SpringBootConfiguration; import org.springframework.boot.test.context.SpringBootTest; import org.springframework.test.context.TestPropertySource; import org.springframework.test.context.jdbc.Sql; @@ -27,16 +26,11 @@ import org.springframework.test.context.jdbc.Sql; * @author Jonathan Leijendekker * @author Thomas Vitale * @author Mark Pollack + * @author Yanming Zhou */ -@SpringBootTest(classes = JdbcChatMemoryRepositoryMysqlIT.TestConfiguration.class) -@TestPropertySource(properties = { "spring.datasource.url=jdbc:tc:mariadb:10.3.39:///", - "spring.datasource.hikari.maximum-pool-size=20", "spring.datasource.hikari.minimum-idle=5" }) +@SpringBootTest +@TestPropertySource(properties = { "spring.datasource.url=jdbc:tc:mariadb:10.3.39:///" }) @Sql(scripts = "classpath:org/springframework/ai/chat/memory/repository/jdbc/schema-mariadb.sql") class JdbcChatMemoryRepositoryMysqlIT extends AbstractJdbcChatMemoryRepositoryIT { - @SpringBootConfiguration - static class TestConfiguration extends BaseTestConfiguration { - - } - } diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryPostgresqlIT.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryPostgresqlIT.java index ae22c45fb..3f288190c 100644 --- a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryPostgresqlIT.java +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryPostgresqlIT.java @@ -16,86 +16,21 @@ package org.springframework.ai.chat.memory.repository.jdbc; -import java.util.List; -import java.util.UUID; - -import javax.sql.DataSource; - -import org.junit.jupiter.api.Test; - -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.SystemMessage; -import org.springframework.ai.chat.messages.UserMessage; -import org.springframework.boot.SpringBootConfiguration; 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 static org.assertj.core.api.Assertions.assertThat; - /** * Integration tests for {@link JdbcChatMemoryRepository} with PostgreSQL. * * @author Jonathan Leijendekker * @author Thomas Vitale * @author Mark Pollack + * @author Yanming Zhou */ -@SpringBootTest(classes = JdbcChatMemoryRepositoryPostgresqlIT.TestConfiguration.class) -@TestPropertySource(properties = { "spring.datasource.url=jdbc:tc:postgresql:17:///", - "spring.datasource.hikari.maximum-pool-size=20", "spring.datasource.hikari.minimum-idle=5" }) +@SpringBootTest +@TestPropertySource(properties = "spring.datasource.url=jdbc:tc:postgresql:17:///") @Sql(scripts = "classpath:org/springframework/ai/chat/memory/repository/jdbc/schema-postgresql.sql") class JdbcChatMemoryRepositoryPostgresqlIT extends AbstractJdbcChatMemoryRepositoryIT { - @Test - void repositoryWithExplicitTransactionManager() { - // Get the repository with explicit transaction manager - ChatMemoryRepository repositoryWithTxManager = TestConfiguration - .chatMemoryRepositoryWithTransactionManager(this.jdbcTemplate, this.jdbcTemplate.getDataSource()); - - var conversationId = UUID.randomUUID().toString(); - var messages = List.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.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 - static class TestConfiguration extends BaseTestConfiguration { - - @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(); - } - - } - }