fix: get query for MSSQL Server

Signed-off-by: Xavier Chopin <[email protected]>

Author: xchopin

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/pom.xml

@@ -66,6 +66,12 @@
 			<artifactId>postgresql</artifactId>
 			<scope>test</scope>
 		</dependency>
+
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>mssqlserver</artifactId>
+			<scope>test</scope>
+		</dependency>
 	</dependencies>
 
 </project>

memory/spring-ai-model-chat-memory-jdbc/pom.xml

@@ -69,6 +69,12 @@
 			<optional>true</optional>
 		</dependency>
 
+		<dependency>
+			<groupId>com.microsoft.sqlserver</groupId>
+			<artifactId>mssql-jdbc</artifactId>
+			<optional>true</optional>
+		</dependency>
+
 		<!-- TESTING -->
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
@@ -82,6 +88,12 @@
 			<scope>test</scope>
 		</dependency>
 
+		<dependency>
+			<groupId>org.testcontainers</groupId>
+			<artifactId>mssqlserver</artifactId>
+			<scope>test</scope>
+		</dependency>
+
 		<dependency>
 			<groupId>org.testcontainers</groupId>
 			<artifactId>postgresql</artifactId>

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

@@ -16,27 +16,27 @@
 
 package org.springframework.ai.chat.memory.jdbc;
 
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
-import java.util.List;
-
 import org.springframework.ai.chat.memory.ChatMemory;
-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.SystemMessage;
-import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.messages.*;
+import org.springframework.boot.jdbc.DatabaseDriver;
 import org.springframework.jdbc.core.BatchPreparedStatementSetter;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.core.RowMapper;
+import org.springframework.util.Assert;
+
+import java.sql.Connection;
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.List;
 
 /**
  * An implementation of {@link ChatMemory} for JDBC. Creating an instance of
  * JdbcChatMemory example:
  * <code>JdbcChatMemory.create(JdbcChatMemoryConfig.builder().jdbcTemplate(jdbcTemplate).build());</code>
  *
  * @author Jonathan Leijendekker
+ * @author Xavier Chopin
  * @since 1.0.0
  */
 public class JdbcChatMemory implements ChatMemory {
@@ -45,14 +45,33 @@ public class JdbcChatMemory implements ChatMemory {
 			INSERT INTO ai_chat_memory (conversation_id, content, type) VALUES (?, ?, ?)""";
 
 	private static final String QUERY_GET = """
-			SELECT content, type FROM ai_chat_memory WHERE conversation_id = ? ORDER BY "timestamp" DESC LIMIT ?""";
+		SELECT content, type \
+		FROM ai_chat_memory \
+		WHERE conversation_id = ? \
+		ORDER BY "timestamp" DESC  \
+		LIMIT ?
+	""";
+
+	private static final String MSSQL_QUERY_GET = """
+    	SELECT content, type \
+    	  FROM ( \
+    	    SELECT TOP (?) content, type, [timestamp] \
+    	    FROM ai_chat_memory \
+    	    WHERE conversation_id = ? \
+    	    ORDER BY [timestamp] DESC \
+    	  ) AS recent \
+    	 ORDER BY [timestamp] ASC \
+	""";
 
 	private static final String QUERY_CLEAR = "DELETE FROM ai_chat_memory WHERE conversation_id = ?";
 
 	private final JdbcTemplate jdbcTemplate;
 
+	private final DatabaseDriver driver;
+
 	public JdbcChatMemory(JdbcChatMemoryConfig config) {
 		this.jdbcTemplate = config.getJdbcTemplate();
+		this.driver = this.detectDialect(this.jdbcTemplate);
 	}
 
 	public static JdbcChatMemory create(JdbcChatMemoryConfig config) {
@@ -66,16 +85,19 @@ public void add(String conversationId, List<Message> messages) {
 
 	@Override
 	public List<Message> get(String conversationId, int lastN) {
-		return this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId, lastN);
+		return switch (driver) {
+			case SQLSERVER -> this.jdbcTemplate.query(MSSQL_QUERY_GET, new MessageRowMapper(), lastN, conversationId);
+			default -> this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId, lastN);
+		};
 	}
 
 	@Override
 	public void clear(String conversationId) {
 		this.jdbcTemplate.update(QUERY_CLEAR, conversationId);
 	}
 
-	private record AddBatchPreparedStatement(String conversationId,
-			List<Message> messages) implements BatchPreparedStatementSetter {
+	private record AddBatchPreparedStatement(String conversationId, List<Message> messages)
+			implements BatchPreparedStatementSetter {
 		@Override
 		public void setValues(PreparedStatement ps, int i) throws SQLException {
 			var message = this.messages.get(i);
@@ -108,4 +130,15 @@ public Message mapRow(ResultSet rs, int i) throws SQLException {
 
 	}
 
+	private DatabaseDriver detectDialect(JdbcTemplate jdbcTemplate) {
+		try {
+			Assert.notNull(jdbcTemplate.getDataSource(), "jdbcTemplate.dataSource must not be null");
+			try (Connection conn = jdbcTemplate.getDataSource().getConnection()) {
+				String url = conn.getMetaData().getURL();
+				return DatabaseDriver.fromJdbcUrl(url);
+			}
+		} catch (SQLException ex) {
+			throw new IllegalStateException("Impossible to detect dialect", ex);
+		}
+	}
 }

memory/spring-ai-model-chat-memory-jdbc/src/main/java/org/springframework/ai/chat/memory/jdbc/aot/hint/JdbcChatMemoryRuntimeHints.java

@@ -36,6 +36,7 @@ public void registerHints(RuntimeHints hints, ClassLoader classLoader) {
 
 		hints.resources()
 			.registerPattern("org/springframework/ai/chat/memory/jdbc/schema-mariadb.sql")
+			.registerPattern("org/springframework/ai/chat/memory/jdbc/schema-mssql.sql")
 			.registerPattern("org/springframework/ai/chat/memory/jdbc/schema-postgresql.sql");
 	}
 

memory/spring-ai-model-chat-memory-jdbc/src/main/resources/org/springframework/ai/chat/memory/jdbc/schema-mssql.sql

@@ -0,0 +1,9 @@
+CREATE TABLE ai_chat_memory (
+    conversation_id VARCHAR(36) NOT NULL,
+    content NVARCHAR(MAX) NOT NULL,
+    type VARCHAR(10) NOT NULL,
+    [timestamp] DATETIME2 NOT NULL DEFAULT SYSDATETIME(),
+    CONSTRAINT type_check CHECK (type IN ('USER', 'ASSISTANT', 'SYSTEM', 'TOOL'))
+);
+
+CREATE INDEX ai_chat_memory_conversation_id_timestamp_idx ON ai_chat_memory(conversation_id, [timestamp] DESC);