diff --git a/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidator.java b/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidator.java index e42d1f1ffd..57aa83e620 100644 --- a/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidator.java +++ b/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidator.java @@ -109,24 +109,32 @@ void validateTableSchema(@Nullable String schemaName, String tableName, String i // Check each column against expected fields List availableColumns = new ArrayList<>(); + String rawIdColumnName = null; for (Map column : columns) { String columnName = (String) column.get("COLUMN_NAME"); Assert.state(columnName != null, "COLUMN_NAME result should not be null"); - columnName = validateAndEnquoteIdentifier(columnName, false); - availableColumns.add(columnName); + String quotedColumnName = validateAndEnquoteIdentifier(columnName, false); + availableColumns.add(quotedColumnName); + if (quotedColumnName.equals(idFieldName)) { + rawIdColumnName = columnName; + } } - // TODO ensure id is a primary key for batch update - expectedColumns.removeAll(availableColumns); - if (expectedColumns.isEmpty()) { - logger.info("MariaDB VectorStore schema validation successful"); - } - else { + if (!expectedColumns.isEmpty()) { throw new IllegalStateException("Missing fields " + expectedColumns); } + Assert.state(rawIdColumnName != null, "id column must be present after passing schema validation"); + if (!isSoleUniqueOrPrimaryKeyColumn(schemaName, tableName, rawIdColumnName)) { + throw new IllegalStateException(String + .format("Column '%s' must be the sole column of a PRIMARY KEY or UNIQUE constraint on table '%s' " + + "so that 'INSERT ... ON DUPLICATE KEY UPDATE' can detect duplicates during batch updates", + idFieldName, tableName)); + } + + logger.info("MariaDB VectorStore schema validation successful"); } catch (DataAccessException | IllegalStateException e) { if (logger.isErrorEnabled()) { @@ -152,6 +160,26 @@ VECTOR INDEX (%s) } } + /** + * Checks whether the given column is, on its own, guaranteed to be unique i.e. it is + * the sole column of either the table's PRIMARY KEY or a UNIQUE constraint. + *

+ * A column that is merely part of a multi-column PRIMARY KEY or UNIQUE constraint is + * intentionally excluded, since the column's value alone does not guarantee + * uniqueness in that case, and 'INSERT ... ON DUPLICATE KEY UPDATE' would not detect + * a "duplicate" row that only matches on that column. + */ + private boolean isSoleUniqueOrPrimaryKeyColumn(@Nullable String schemaName, String tableName, + String rawColumnName) { + String query = "SELECT s.INDEX_NAME FROM INFORMATION_SCHEMA.STATISTICS s " + + "WHERE s.TABLE_SCHEMA = ? AND s.TABLE_NAME = ? AND s.NON_UNIQUE = 0 " + "GROUP BY s.INDEX_NAME " + + "HAVING COUNT(*) = 1 AND SUM(s.COLUMN_NAME = ?) = 1"; + + List> matches = this.jdbcTemplate.queryForList(query, schemaName, tableName, + rawColumnName); + return !matches.isEmpty(); + } + /** * Escaped identifier according to MariaDB requirement. * @param identifier identifier diff --git a/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidatorIT.java b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidatorIT.java new file mode 100644 index 0000000000..d29147b5ec --- /dev/null +++ b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBSchemaValidatorIT.java @@ -0,0 +1,174 @@ +/* + * Copyright 2023-present 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.vectorstore.mariadb; + +import java.util.function.Consumer; + +import javax.sql.DataSource; + +import com.zaxxer.hikari.HikariDataSource; +import org.junit.jupiter.api.Test; +import org.testcontainers.containers.MariaDBContainer; +import org.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; + +import org.springframework.boot.SpringBootConfiguration; +import org.springframework.boot.autoconfigure.EnableAutoConfiguration; +import org.springframework.boot.jdbc.autoconfigure.DataSourceAutoConfiguration; +import org.springframework.boot.jdbc.autoconfigure.DataSourceProperties; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThatIllegalStateException; +import static org.assertj.core.api.Assertions.assertThatNoException; + +/** + * Integration tests for {@link MariaDBSchemaValidator}, run against a real MariaDB + * instance so that the SQL used to inspect {@code INFORMATION_SCHEMA} is verified against + * actual server behavior rather than a mocked {@link JdbcTemplate}. + * + * @author dev-xong + */ +@Testcontainers +class MariaDBSchemaValidatorIT { + + private static final String SCHEMA_NAME = "testdb"; + + @Container + @SuppressWarnings("resource") + static MariaDBContainer mariadbContainer = new MariaDBContainer<>(MariaDBImage.DEFAULT_IMAGE) + .withUsername("mariadb") + .withPassword("mariadbpwd") + .withDatabaseName(SCHEMA_NAME); + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withUserConfiguration(TestApplication.class); + + @Test + void validatesSuccessfullyWhenIdIsSolePrimaryKey() { + withSchemaValidator(""" + CREATE TABLE vector_store ( + id VARCHAR(36) NOT NULL PRIMARY KEY, + content TEXT, + metadata JSON, + embedding VECTOR(1024) NOT NULL, + VECTOR INDEX (embedding) + ) ENGINE=InnoDB""", schemaValidator -> assertThatNoException().isThrownBy(() -> schemaValidator + .validateTableSchema(SCHEMA_NAME, "vector_store", "id", "content", "metadata", "embedding", 1024))); + } + + @Test + void validatesSuccessfullyWhenIdIsSoleUniqueColumn() { + withSchemaValidator(""" + CREATE TABLE vector_store ( + id VARCHAR(36) NOT NULL UNIQUE, + content TEXT, + metadata JSON, + embedding VECTOR(1024) NOT NULL, + VECTOR INDEX (embedding) + ) ENGINE=InnoDB""", schemaValidator -> assertThatNoException().isThrownBy(() -> schemaValidator + .validateTableSchema(SCHEMA_NAME, "vector_store", "id", "content", "metadata", "embedding", 1024))); + } + + @Test + void rejectsWhenIdHasNoUniqueConstraint() { + withSchemaValidator(""" + CREATE TABLE vector_store ( + id VARCHAR(36) NOT NULL, + content TEXT, + metadata JSON, + embedding VECTOR(1024) NOT NULL, + VECTOR INDEX (embedding) + ) ENGINE=InnoDB""", + schemaValidator -> assertThatIllegalStateException() + .isThrownBy(() -> schemaValidator.validateTableSchema(SCHEMA_NAME, "vector_store", "id", "content", + "metadata", "embedding", 1024)) + .withMessageContaining("sole column")); + } + + @Test + void rejectsWhenIdIsPartOfCompositePrimaryKey() { + withSchemaValidator(""" + CREATE TABLE vector_store ( + id VARCHAR(36) NOT NULL, + tenant_id VARCHAR(36) NOT NULL, + content TEXT, + metadata JSON, + embedding VECTOR(1024) NOT NULL, + PRIMARY KEY (id, tenant_id), + VECTOR INDEX (embedding) + ) ENGINE=InnoDB""", + schemaValidator -> assertThatIllegalStateException() + .isThrownBy(() -> schemaValidator.validateTableSchema(SCHEMA_NAME, "vector_store", "id", "content", + "metadata", "embedding", 1024)) + .withMessageContaining("sole column")); + } + + @Test + void rejectsWhenIdIsPartOfCompositeUniqueConstraint() { + withSchemaValidator(""" + CREATE TABLE vector_store ( + id VARCHAR(36) NOT NULL, + tenant_id VARCHAR(36) NOT NULL, + content TEXT, + metadata JSON, + embedding VECTOR(1024) NOT NULL, + UNIQUE KEY ux_id_tenant (id, tenant_id), + VECTOR INDEX (embedding) + ) ENGINE=InnoDB""", + schemaValidator -> assertThatIllegalStateException() + .isThrownBy(() -> schemaValidator.validateTableSchema(SCHEMA_NAME, "vector_store", "id", "content", + "metadata", "embedding", 1024)) + .withMessageContaining("sole column")); + } + + private void withSchemaValidator(String createTableSql, Consumer assertions) { + this.contextRunner.run(context -> { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + jdbcTemplate.execute("DROP TABLE IF EXISTS vector_store"); + jdbcTemplate.execute(createTableSql); + assertions.accept(new MariaDBSchemaValidator(jdbcTemplate)); + }); + } + + @SpringBootConfiguration + @EnableAutoConfiguration(exclude = DataSourceAutoConfiguration.class) + public static class TestApplication { + + @Bean + public JdbcTemplate jdbcTemplate(DataSource dataSource) { + return new JdbcTemplate(dataSource); + } + + @Bean + public DataSourceProperties dataSourceProperties() { + DataSourceProperties properties = new DataSourceProperties(); + properties.setUrl(mariadbContainer.getJdbcUrl()); + properties.setUsername(mariadbContainer.getUsername()); + properties.setPassword(mariadbContainer.getPassword()); + return properties; + } + + @Bean + public HikariDataSource dataSource(DataSourceProperties dataSourceProperties) { + return dataSourceProperties.initializeDataSourceBuilder().type(HikariDataSource.class).build(); + } + + } + +}