Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

JdbcChatMemoryRepository should use the provided JdbcTemplate

Browse Source 
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 <zhouyanming@gmail.com>
This commit is contained in:
Yanming Zhou
2025-05-16 10:40:07 +08:00
committed by Mark Pollack
parent bbc2bbf661
commit b83a162506
3 changed files with 22 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

35
memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/AbstractJdbcChatMemoryRepositoryIT.java
View File

@@ -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<Message> orderedMessages = List.of(firstMessage, secondMessage, thirdMessage, fourthMessage);
repository.saveAll(conversationId, orderedMessages);
chatMemoryRepository.saveAll(conversationId, orderedMessages);
// Retrieve messages using the repository
List<Message> retrievedMessages = repository.findByConversationId(conversationId);
List<Message> 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();
}
}

12
memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryMysqlIT.java
View File

@@ -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 {
}
}

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

@@ -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.<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
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();
}
}
}