Skip to content

Commit d6c5337

Browse files
quaffquaff
committed
Align JdbcChatMemoryProperties with other modules in Spring Boot
Such as `BatchProperties.Jdbc` and `IntegrationProperties.Jdbc` 1. Change type of `initialize-schema` from `boolean` to `DatabaseInitializationMode` 2. Allow to customize schema location Signed-off-by: Yanming Zhou <zhouyanming@gmail.com>
1 parent 4fe74d8 commit d6c5337

File tree

6 files changed

+114
-66
lines changed

6 files changed

+114
-66
lines changed

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/main/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryAutoConfiguration.java

Lines changed: 15 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,15 +26,17 @@
2626
import org.springframework.boot.autoconfigure.AutoConfiguration;
2727
import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
2828
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
29-
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
3029
import org.springframework.boot.autoconfigure.jdbc.JdbcTemplateAutoConfiguration;
30+
import org.springframework.boot.autoconfigure.sql.init.OnDatabaseInitializationCondition;
3131
import org.springframework.boot.context.properties.EnableConfigurationProperties;
3232
import org.springframework.context.annotation.Bean;
33+
import org.springframework.context.annotation.Conditional;
3334
import org.springframework.jdbc.core.JdbcTemplate;
3435

3536
/**
3637
* @author Jonathan Leijendekker
3738
* @author Thomas Vitale
39+
* @author Yanming Zhou
3840
* @since 1.0.0
3941
*/
4042
@AutoConfiguration(after = JdbcTemplateAutoConfiguration.class, before = ChatMemoryAutoConfiguration.class)
@@ -52,11 +54,19 @@ JdbcChatMemoryRepository chatMemoryRepository(JdbcTemplate jdbcTemplate) {
5254

5355
@Bean
5456
@ConditionalOnMissingBean
55-
@ConditionalOnProperty(prefix = JdbcChatMemoryProperties.CONFIG_PREFIX, name = "initialize-schema",
56-
havingValue = "true", matchIfMissing = true)
57-
JdbcChatMemoryDataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource) {
57+
@Conditional(OnJdbcChatMemoryDatasourceInitializationCondition.class)
58+
JdbcChatMemoryDataSourceScriptDatabaseInitializer jdbcChatMemoryScriptDatabaseInitializer(DataSource dataSource,
59+
JdbcChatMemoryProperties properties) {
5860
logger.debug("Initializing schema for JdbcChatMemoryRepository");
59-
return new JdbcChatMemoryDataSourceScriptDatabaseInitializer(dataSource);
61+
return new JdbcChatMemoryDataSourceScriptDatabaseInitializer(dataSource, properties);
62+
}
63+
64+
static class OnJdbcChatMemoryDatasourceInitializationCondition extends OnDatabaseInitializationCondition {
65+
66+
OnJdbcChatMemoryDatasourceInitializationCondition() {
67+
super("Jdbc Chat Memory", JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema");
68+
}
69+
6070
}
6171

6272
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/main/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryDataSourceScriptDatabaseInitializer.java

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -22,35 +22,36 @@
2222

2323
import org.springframework.boot.jdbc.init.DataSourceScriptDatabaseInitializer;
2424
import org.springframework.boot.jdbc.init.PlatformPlaceholderDatabaseDriverResolver;
25-
import org.springframework.boot.sql.init.DatabaseInitializationMode;
2625
import org.springframework.boot.sql.init.DatabaseInitializationSettings;
26+
import org.springframework.util.StringUtils;
2727

2828
/**
2929
* Performs database initialization for the JDBC Chat Memory Repository.
3030
*
31+
* @author Jonathan Leijendekker
32+
* @author Yanming Zhou
3133
* @since 1.0.0
3234
*/
3335
class JdbcChatMemoryDataSourceScriptDatabaseInitializer extends DataSourceScriptDatabaseInitializer {
3436

35-
private static final String SCHEMA_LOCATION = "classpath:org/springframework/ai/chat/memory/jdbc/schema-@@platform@@.sql";
36-
37-
JdbcChatMemoryDataSourceScriptDatabaseInitializer(DataSource dataSource) {
38-
super(dataSource, getSettings(dataSource));
37+
JdbcChatMemoryDataSourceScriptDatabaseInitializer(DataSource dataSource, JdbcChatMemoryProperties properties) {
38+
super(dataSource, getSettings(dataSource, properties));
3939
}
4040

41-
static DatabaseInitializationSettings getSettings(DataSource dataSource) {
41+
static DatabaseInitializationSettings getSettings(DataSource dataSource, JdbcChatMemoryProperties properties) {
4242
var settings = new DatabaseInitializationSettings();
43-
settings.setSchemaLocations(resolveSchemaLocations(dataSource));
44-
settings.setMode(DatabaseInitializationMode.ALWAYS);
43+
settings.setSchemaLocations(resolveSchemaLocations(dataSource, properties));
44+
settings.setMode(properties.getInitializeSchema());
4545
settings.setContinueOnError(true);
46-
4746
return settings;
4847
}
4948

50-
static List<String> resolveSchemaLocations(DataSource dataSource) {
49+
static List<String> resolveSchemaLocations(DataSource dataSource, JdbcChatMemoryProperties properties) {
5150
var platformResolver = new PlatformPlaceholderDatabaseDriverResolver();
52-
53-
return platformResolver.resolveAll(dataSource, SCHEMA_LOCATION);
51+
if (StringUtils.hasText(properties.getPlatform())) {
52+
return platformResolver.resolveAll(properties.getPlatform(), properties.getSchema());
53+
}
54+
return platformResolver.resolveAll(dataSource, properties.getSchema());
5455
}
5556

5657
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/main/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryProperties.java

Lines changed: 35 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,27 +17,58 @@
1717
package org.springframework.ai.model.chat.memory.jdbc.autoconfigure;
1818

1919
import org.springframework.boot.context.properties.ConfigurationProperties;
20+
import org.springframework.boot.sql.init.DatabaseInitializationMode;
2021

2122
/**
2223
* @author Jonathan Leijendekker
2324
* @author Thomas Vitale
25+
* @author Yanming Zhou
2426
* @since 1.0.0
2527
*/
2628
@ConfigurationProperties(JdbcChatMemoryProperties.CONFIG_PREFIX)
2729
public class JdbcChatMemoryProperties {
2830

2931
public static final String CONFIG_PREFIX = "spring.ai.chat.memory.repository.jdbc";
3032

33+
private static final String DEFAULT_SCHEMA_LOCATION = "classpath:org/springframework/ai/chat/memory/jdbc/schema-@@platform@@.sql";
34+
35+
/**
36+
* Path to the SQL file to use to initialize the database schema.
37+
*/
38+
private String schema = DEFAULT_SCHEMA_LOCATION;
39+
3140
/**
32-
* Whether to initialize the schema on startup.
41+
* Platform to use in initialization scripts if the @@platform@@ placeholder is used.
42+
* Auto-detected by default.
3343
*/
34-
private boolean initializeSchema = true;
44+
private String platform;
45+
46+
/**
47+
* Database schema initialization mode.
48+
*/
49+
private DatabaseInitializationMode initializeSchema = DatabaseInitializationMode.EMBEDDED;
50+
51+
public String getPlatform() {
52+
return platform;
53+
}
54+
55+
public void setPlatform(String platform) {
56+
this.platform = platform;
57+
}
58+
59+
public String getSchema() {
60+
return schema;
61+
}
62+
63+
public void setSchema(String schema) {
64+
this.schema = schema;
65+
}
3566

36-
public boolean isInitializeSchema() {
67+
public DatabaseInitializationMode getInitializeSchema() {
3768
return this.initializeSchema;
3869
}
3970

40-
public void setInitializeSchema(boolean initializeSchema) {
71+
public void setInitializeSchema(DatabaseInitializationMode initializeSchema) {
4172
this.initializeSchema = initializeSchema;
4273
}
4374

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryDataSourceScriptDatabaseInitializerPostgresqlTests.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333

3434
/**
3535
* @author Jonathan Leijendekker
36+
* @author Yanming Zhou
3637
*/
3738
@Testcontainers
3839
class JdbcChatMemoryDataSourceScriptDatabaseInitializerPostgresqlTests {
@@ -57,7 +58,8 @@ class JdbcChatMemoryDataSourceScriptDatabaseInitializerPostgresqlTests {
5758
void getSettings_shouldHaveSchemaLocations() {
5859
this.contextRunner.run(context -> {
5960
var dataSource = context.getBean(DataSource.class);
60-
var settings = JdbcChatMemoryDataSourceScriptDatabaseInitializer.getSettings(dataSource);
61+
var properties = context.getBean(JdbcChatMemoryProperties.class);
62+
var settings = JdbcChatMemoryDataSourceScriptDatabaseInitializer.getSettings(dataSource, properties);
6163

6264
assertThat(settings.getSchemaLocations())
6365
.containsOnly("classpath:org/springframework/ai/chat/memory/jdbc/schema-postgresql.sql");

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryPostgresqlAutoConfigurationIT.java

Lines changed: 43 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
* @author Jonathan Leijendekker
3939
* @author Thomas Vitale
4040
* @author Linar Abzaltdinov
41+
* @author Yanming Zhou
4142
*/
4243
class JdbcChatMemoryPostgresqlAutoConfigurationIT {
4344

@@ -48,84 +49,85 @@ class JdbcChatMemoryPostgresqlAutoConfigurationIT {
4849

4950
@Test
5051
void jdbcChatMemoryScriptDatabaseInitializer_shouldBeLoaded() {
51-
this.contextRunner.withPropertyValues("spring.ai.chat.memory.jdbc.initialize-schema=true")
52-
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
53-
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=true")
52+
this.contextRunner.withPropertyValues(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema=always")
5453
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
5554
}
5655

5756
@Test
5857
void jdbcChatMemoryScriptDatabaseInitializer_shouldNotBeLoaded() {
59-
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=false")
58+
this.contextRunner.withPropertyValues(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema=never")
6059
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isFalse());
6160
}
6261

6362
@Test
6463
void initializeSchemaEnabledWithProperty() {
65-
this.contextRunner.withPropertyValues("spring.ai.chat.memory.repository.jdbc.initialize-schema=true")
64+
this.contextRunner.withPropertyValues(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema=always")
6665
.run(context -> assertThat(context.containsBean("jdbcChatMemoryScriptDatabaseInitializer")).isTrue());
6766
}
6867

6968
@Test
7069
void useAutoConfiguredJdbcChatMemoryRepository() {
71-
this.contextRunner.run(context -> {
72-
var chatMemoryRepository = context.getBean(JdbcChatMemoryRepository.class);
73-
var conversationId = UUID.randomUUID().toString();
74-
var userMessage = new UserMessage("Message from the user");
70+
this.contextRunner.withPropertyValues(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema=always")
71+
.run(context -> {
72+
var chatMemoryRepository = context.getBean(JdbcChatMemoryRepository.class);
73+
var conversationId = UUID.randomUUID().toString();
74+
var userMessage = new UserMessage("Message from the user");
7575

76-
chatMemoryRepository.saveAll(conversationId, List.of(userMessage));
76+
chatMemoryRepository.saveAll(conversationId, List.of(userMessage));
7777

78-
assertThat(chatMemoryRepository.findByConversationId(conversationId)).hasSize(1);
79-
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEqualTo(List.of(userMessage));
78+
assertThat(chatMemoryRepository.findByConversationId(conversationId)).hasSize(1);
79+
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEqualTo(List.of(userMessage));
8080

81-
chatMemoryRepository.deleteByConversationId(conversationId);
81+
chatMemoryRepository.deleteByConversationId(conversationId);
8282

83-
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEmpty();
83+
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEmpty();
8484

85-
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
86-
new AssistantMessage("Message from the assistant 1"));
85+
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
86+
new AssistantMessage("Message from the assistant 1"));
8787

88-
chatMemoryRepository.saveAll(conversationId, multipleMessages);
88+
chatMemoryRepository.saveAll(conversationId, multipleMessages);
8989

90-
assertThat(chatMemoryRepository.findByConversationId(conversationId)).hasSize(multipleMessages.size());
91-
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEqualTo(multipleMessages);
92-
});
90+
assertThat(chatMemoryRepository.findByConversationId(conversationId)).hasSize(multipleMessages.size());
91+
assertThat(chatMemoryRepository.findByConversationId(conversationId)).isEqualTo(multipleMessages);
92+
});
9393
}
9494

9595
@Test
9696
void useAutoConfiguredChatMemoryWithJdbc() {
97-
this.contextRunner.withConfiguration(AutoConfigurations.of(ChatMemoryAutoConfiguration.class)).run(context -> {
98-
assertThat(context).hasSingleBean(ChatMemory.class);
99-
assertThat(context).hasSingleBean(JdbcChatMemoryRepository.class);
97+
this.contextRunner.withPropertyValues(JdbcChatMemoryProperties.CONFIG_PREFIX + ".initialize-schema=always")
98+
.withConfiguration(AutoConfigurations.of(ChatMemoryAutoConfiguration.class))
99+
.run(context -> {
100+
assertThat(context).hasSingleBean(ChatMemory.class);
101+
assertThat(context).hasSingleBean(JdbcChatMemoryRepository.class);
100102

101-
var chatMemory = context.getBean(ChatMemory.class);
102-
var conversationId = UUID.randomUUID().toString();
103-
var userMessage = new UserMessage("Message from the user");
103+
var chatMemory = context.getBean(ChatMemory.class);
104+
var conversationId = UUID.randomUUID().toString();
105+
var userMessage = new UserMessage("Message from the user");
104106

105-
chatMemory.add(conversationId, userMessage);
107+
chatMemory.add(conversationId, userMessage);
106108

107-
assertThat(chatMemory.get(conversationId)).hasSize(1);
108-
assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage));
109+
assertThat(chatMemory.get(conversationId)).hasSize(1);
110+
assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage));
109111

110-
var assistantMessage = new AssistantMessage("Message from the assistant");
112+
var assistantMessage = new AssistantMessage("Message from the assistant");
111113

112-
chatMemory.add(conversationId, List.of(assistantMessage));
114+
chatMemory.add(conversationId, List.of(assistantMessage));
113115

114-
assertThat(chatMemory.get(conversationId)).hasSize(2);
115-
assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage, assistantMessage));
116+
assertThat(chatMemory.get(conversationId)).hasSize(2);
117+
assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage, assistantMessage));
116118

117-
chatMemory.clear(conversationId);
119+
chatMemory.clear(conversationId);
118120

119-
assertThat(chatMemory.get(conversationId)).isEmpty();
121+
assertThat(chatMemory.get(conversationId)).isEmpty();
120122

121-
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
122-
new AssistantMessage("Message from the assistant 1"));
123+
var multipleMessages = List.<Message>of(new UserMessage("Message from the user 1"),
124+
new AssistantMessage("Message from the assistant 1"));
123125

124-
chatMemory.add(conversationId, multipleMessages);
126+
chatMemory.add(conversationId, multipleMessages);
125127

126-
assertThat(chatMemory.get(conversationId)).hasSize(multipleMessages.size());
127-
assertThat(chatMemory.get(conversationId)).isEqualTo(multipleMessages);
128-
});
128+
assertThat(chatMemory.get(conversationId)).hasSize(multipleMessages.size());
129+
assertThat(chatMemory.get(conversationId)).isEqualTo(multipleMessages);
130+
});
129131
}
130132

131133
}

‎auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryPropertiesTests.java

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -17,27 +17,29 @@
1717
package org.springframework.ai.model.chat.memory.jdbc.autoconfigure;
1818

1919
import org.junit.jupiter.api.Test;
20+
import org.springframework.boot.sql.init.DatabaseInitializationMode;
2021

2122
import static org.assertj.core.api.Assertions.assertThat;
2223

2324
/**
2425
* @author Jonathan Leijendekker
26+
* @author Yanming Zhou
2527
*/
2628
class JdbcChatMemoryPropertiesTests {
2729

2830
@Test
2931
void defaultValues() {
3032
var props = new JdbcChatMemoryProperties();
3133

32-
assertThat(props.isInitializeSchema()).isTrue();
34+
assertThat(props.getInitializeSchema()).isEqualTo(DatabaseInitializationMode.EMBEDDED);
3335
}
3436

3537
@Test
3638
void customValues() {
3739
var props = new JdbcChatMemoryProperties();
38-
props.setInitializeSchema(false);
40+
props.setInitializeSchema(DatabaseInitializationMode.ALWAYS);
3941

40-
assertThat(props.isInitializeSchema()).isFalse();
42+
assertThat(props.getInitializeSchema()).isEqualTo(DatabaseInitializationMode.ALWAYS);
4143
}
4244

4345
}

0 commit comments

Comments
 (0)