From 008a760fa76201f860b85a49af4b1db61bb49d3c Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Thu, 15 May 2025 16:32:22 -0400 Subject: [PATCH] Add dataSource() method to JdbcChatMemoryRepository.Builder and improve dialect detection - Added a dataSource(DataSource) method to the builder - Now, if no dialect is specified, the implementation will attempt to detect it from the effective DataSource (either set directly or obtained from the JdbcTemplate). - Updated builder logic to default to the DataSource from the JdbcTemplate if dataSource() is not called, ensuring dialect detection works out-of-the-box. - The builder now prefers a directly provided DataSource, but remains backwards compatible with JdbcTemplate-based configuration. - Added tests Fixes #3148 Signed-off-by: Mark Pollack --- .../jdbc/JdbcChatMemoryRepository.java | 76 +++++- .../JdbcChatMemoryRepositoryBuilderTests.java | 226 ++++++++++++++++++ 2 files changed, 295 insertions(+), 7 deletions(-) create mode 100644 memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryBuilderTests.java 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 166501790..b24254deb 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 @@ -22,6 +22,8 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.atomic.AtomicLong; +import javax.sql.DataSource; + import org.springframework.ai.chat.memory.ChatMemoryRepository; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; @@ -35,8 +37,15 @@ 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.TransactionDefinition; +import org.springframework.transaction.TransactionException; +import org.springframework.transaction.TransactionStatus; +import org.springframework.transaction.support.AbstractPlatformTransactionManager; +import org.springframework.transaction.support.DefaultTransactionStatus; import org.springframework.transaction.support.TransactionTemplate; import org.springframework.util.Assert; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; /** * An implementation of {@link ChatMemoryRepository} for JDBC. @@ -55,14 +64,16 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository { private final JdbcChatMemoryRepositoryDialect dialect; - private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryRepositoryDialect dialect, + private static final Logger logger = LoggerFactory.getLogger(JdbcChatMemoryRepository.class); + + private JdbcChatMemoryRepository(DataSource dataSource, JdbcChatMemoryRepositoryDialect dialect, PlatformTransactionManager txManager) { - Assert.notNull(jdbcTemplate, "jdbcTemplate cannot be null"); + Assert.notNull(dataSource, "dataSource cannot be null"); Assert.notNull(dialect, "dialect cannot be null"); - this.jdbcTemplate = jdbcTemplate; + this.jdbcTemplate = new JdbcTemplate(dataSource); this.dialect = dialect; this.transactionTemplate = new TransactionTemplate( - txManager != null ? txManager : new DataSourceTransactionManager(jdbcTemplate.getDataSource())); + txManager != null ? txManager : new DataSourceTransactionManager(dataSource)); } @Override @@ -157,8 +168,12 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository { private JdbcChatMemoryRepositoryDialect dialect; + private DataSource dataSource; + private PlatformTransactionManager platformTransactionManager; + private static final Logger logger = LoggerFactory.getLogger(Builder.class); + private Builder() { } @@ -172,15 +187,62 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository { return this; } + public Builder dataSource(DataSource dataSource) { + this.dataSource = dataSource; + 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, this.platformTransactionManager); + DataSource effectiveDataSource = resolveDataSource(); + JdbcChatMemoryRepositoryDialect effectiveDialect = resolveDialect(effectiveDataSource); + return new JdbcChatMemoryRepository(effectiveDataSource, effectiveDialect, this.platformTransactionManager); + } + + private DataSource resolveDataSource() { + if (this.dataSource != null) { + return this.dataSource; + } + if (this.jdbcTemplate != null && this.jdbcTemplate.getDataSource() != null) { + return this.jdbcTemplate.getDataSource(); + } + throw new IllegalArgumentException("DataSource must be set (either via dataSource() or jdbcTemplate())"); + } + + private JdbcChatMemoryRepositoryDialect resolveDialect(DataSource dataSource) { + if (this.dialect == null) { + try { + return JdbcChatMemoryRepositoryDialect.from(dataSource); + } + catch (Exception ex) { + throw new IllegalStateException("Could not detect dialect from datasource", ex); + } + } + else { + warnIfDialectMismatch(dataSource, this.dialect); + return this.dialect; + } + } + + /** + * Logs a warning if the explicitly set dialect differs from the dialect detected + * from the DataSource. + */ + private void warnIfDialectMismatch(DataSource dataSource, JdbcChatMemoryRepositoryDialect explicitDialect) { + try { + JdbcChatMemoryRepositoryDialect detected = JdbcChatMemoryRepositoryDialect.from(dataSource); + if (!detected.getClass().equals(explicitDialect.getClass())) { + logger.warn("Explicitly set dialect {} will be used instead of detected dialect {} from datasource", + explicitDialect.getClass().getSimpleName(), detected.getClass().getSimpleName()); + } + } + catch (Exception ex) { + logger.debug("Could not detect dialect from datasource", ex); + } } } diff --git a/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryBuilderTests.java b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryBuilderTests.java new file mode 100644 index 000000000..c6eb00c4a --- /dev/null +++ b/memory/repository/spring-ai-model-chat-memory-repository-jdbc/src/test/java/org/springframework/ai/chat/memory/repository/jdbc/JdbcChatMemoryRepositoryBuilderTests.java @@ -0,0 +1,226 @@ +/* + * Copyright 2023-2025 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.ai.chat.memory.repository.jdbc; + +import java.sql.Connection; +import java.sql.DatabaseMetaData; +import java.sql.SQLException; +import javax.sql.DataSource; + +import org.junit.jupiter.api.Test; + +import org.springframework.transaction.PlatformTransactionManager; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Tests for {@link JdbcChatMemoryRepository.Builder}. + * + * @author Mark Pollack + */ +public class JdbcChatMemoryRepositoryBuilderTests { + + @Test + void testBuilderWithExplicitDialect() { + DataSource dataSource = mock(DataSource.class); + JdbcChatMemoryRepositoryDialect dialect = mock(JdbcChatMemoryRepositoryDialect.class); + + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder() + .dataSource(dataSource) + .dialect(dialect) + .build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithExplicitDialectAndTransactionManager() { + DataSource dataSource = mock(DataSource.class); + JdbcChatMemoryRepositoryDialect dialect = mock(JdbcChatMemoryRepositoryDialect.class); + PlatformTransactionManager txManager = mock(PlatformTransactionManager.class); + + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder() + .dataSource(dataSource) + .dialect(dialect) + .transactionManager(txManager) + .build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithDialectFromDataSource() throws SQLException { + // Setup mocks + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:postgresql://localhost:5432/testdb"); + + // Test with dialect from datasource + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithMysqlDialectFromDataSource() throws SQLException { + // Setup mocks for MySQL + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:mysql://localhost:3306/testdb"); + + // Test with dialect from datasource + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithSqlServerDialectFromDataSource() throws SQLException { + // Setup mocks for SQL Server + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:sqlserver://localhost:1433;databaseName=testdb"); + + // Test with dialect from datasource + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithHsqldbDialectFromDataSource() throws SQLException { + // Setup mocks for HSQLDB + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:hsqldb:mem:testdb"); + + // Test with dialect from datasource + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithUnknownDialectFromDataSource() throws SQLException { + // Setup mocks for unknown database + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:unknown://localhost:1234/testdb"); + + // Test with dialect from datasource - should default to PostgreSQL + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithExceptionInDataSourceConnection() throws SQLException { + // Setup mocks with exception + DataSource dataSource = mock(DataSource.class); + when(dataSource.getConnection()).thenThrow(new SQLException("Connection failed")); + + // Test with dialect from datasource - should default to PostgreSQL + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder().dataSource(dataSource).build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithNullDataSource() { + assertThatThrownBy(() -> JdbcChatMemoryRepository.builder().build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("DataSource must be set (either via dataSource() or jdbcTemplate())"); + } + + @Test + void testBuilderWithNullDataSourceButExplicitDialect() { + DataSource dataSource = mock(DataSource.class); + JdbcChatMemoryRepositoryDialect dialect = mock(JdbcChatMemoryRepositoryDialect.class); + + // Should work because dialect is explicitly set + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder() + .dataSource(dataSource) + .dialect(dialect) + .build(); + + assertThat(repository).isNotNull(); + } + + @Test + void testBuilderWithNullDataSourceAndDialect() { + assertThatThrownBy(() -> JdbcChatMemoryRepository.builder().build()) + .isInstanceOf(IllegalArgumentException.class) + .hasMessage("DataSource must be set (either via dataSource() or jdbcTemplate())"); + } + + /** + * Verifies that when an explicit dialect is provided to the builder, it takes + * precedence over any dialect detected from the DataSource. If the explicit dialect + * differs from the detected one, the explicit dialect is used and a warning is + * logged. This ensures that user intent (explicit configuration) always overrides + * automatic detection. + */ + @Test + void testBuilderPreferenceForExplicitDialect() throws SQLException { + // Setup mocks for PostgreSQL + DataSource dataSource = mock(DataSource.class); + Connection connection = mock(Connection.class); + DatabaseMetaData metaData = mock(DatabaseMetaData.class); + + when(dataSource.getConnection()).thenReturn(connection); + when(connection.getMetaData()).thenReturn(metaData); + when(metaData.getURL()).thenReturn("jdbc:postgresql://localhost:5432/testdb"); + + // Create an explicit MySQL dialect + JdbcChatMemoryRepositoryDialect mysqlDialect = new MysqlChatMemoryRepositoryDialect(); + + // Test with explicit dialect - should use MySQL dialect even though PostgreSQL is + // detected + JdbcChatMemoryRepository repository = JdbcChatMemoryRepository.builder() + .dataSource(dataSource) + .dialect(mysqlDialect) + .build(); + + assertThat(repository).isNotNull(); + // Verify warning was logged (would need to use a logging framework test utility + // for this) + } + +}