From 2ea5022a8d38e666cb92266ea0f8b041766657d9 Mon Sep 17 00:00:00 2001 From: jiajingda Date: Tue, 7 Apr 2026 09:35:57 +0800 Subject: [PATCH] Add health indicator for PgVector vector store Adds a Spring Boot Actuator HealthIndicator for the PgVector vector store that verifies the vector PostgreSQL extension is installed and reachable. Reports the extension version as part of the health details when UP. The health indicator can be disabled via: management.health.pgvector.enabled=false Closes gh-1611 Signed-off-by: jiajingda --- .../pom.xml | 6 + .../PgVectorStoreHealthAutoConfiguration.java | 50 +++++++ .../PgVectorStoreHealthIndicator.java | 55 +++++++ ...ot.autoconfigure.AutoConfiguration.imports | 1 + ...ctorStoreHealthAutoConfigurationTests.java | 139 ++++++++++++++++++ 5 files changed, 251 insertions(+) create mode 100644 auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfiguration.java create mode 100644 auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthIndicator.java create mode 100644 auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/test/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfigurationTests.java diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/pom.xml b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/pom.xml index 9985d74f0f..2ac4f57fc9 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/pom.xml +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/pom.xml @@ -96,5 +96,11 @@ ${project.parent.version} test + + + org.springframework.boot + spring-boot-health + true + diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfiguration.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfiguration.java new file mode 100644 index 0000000000..f64fc6506a --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfiguration.java @@ -0,0 +1,50 @@ +/* + * 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.pgvector.autoconfigure; + +import javax.sql.DataSource; + +import org.springframework.ai.vectorstore.pgvector.PgVectorStore; +import org.springframework.boot.autoconfigure.AutoConfiguration; +import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; +import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; +import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean; +import org.springframework.boot.health.autoconfigure.contributor.ConditionalOnEnabledHealthIndicator; +import org.springframework.boot.health.contributor.HealthIndicator; +import org.springframework.context.annotation.Bean; +import org.springframework.jdbc.core.JdbcTemplate; + +/** + * {@link AutoConfiguration Auto-configuration} for {@link PgVectorStore} health + * indicator. + * + * @author jiajingda + * @since 1.1.0 + */ +@AutoConfiguration(after = PgVectorStoreAutoConfiguration.class) +@ConditionalOnClass({ PgVectorStore.class, DataSource.class, JdbcTemplate.class, HealthIndicator.class }) +@ConditionalOnBean(PgVectorStore.class) +@ConditionalOnEnabledHealthIndicator("pgvector") +public class PgVectorStoreHealthAutoConfiguration { + + @Bean + @ConditionalOnMissingBean(name = "pgvectorHealthIndicator") + public HealthIndicator pgvectorHealthIndicator(JdbcTemplate jdbcTemplate) { + return new PgVectorStoreHealthIndicator(jdbcTemplate); + } + +} diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthIndicator.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthIndicator.java new file mode 100644 index 0000000000..b8917c80f3 --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthIndicator.java @@ -0,0 +1,55 @@ +/* + * 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.pgvector.autoconfigure; + +import org.springframework.boot.health.contributor.AbstractHealthIndicator; +import org.springframework.boot.health.contributor.Health; +import org.springframework.jdbc.core.JdbcTemplate; +import org.springframework.util.Assert; + +/** + * {@link AbstractHealthIndicator Health indicator} for PgVector vector store. + *

+ * Verifies that the {@code vector} PostgreSQL extension is installed and reachable. + * Reports the extension version as part of the health details when {@code UP}. + * + * @author jiajingda + * @since 1.1.0 + */ +public class PgVectorStoreHealthIndicator extends AbstractHealthIndicator { + + private static final String VECTOR_EXTENSION_QUERY = "SELECT extversion FROM pg_extension WHERE extname = 'vector'"; + + private final JdbcTemplate jdbcTemplate; + + public PgVectorStoreHealthIndicator(JdbcTemplate jdbcTemplate) { + super("PgVector health check failed"); + Assert.notNull(jdbcTemplate, "JdbcTemplate must not be null"); + this.jdbcTemplate = jdbcTemplate; + } + + @Override + protected void doHealthCheck(Health.Builder builder) throws Exception { + String version = this.jdbcTemplate.queryForObject(VECTOR_EXTENSION_QUERY, String.class); + if (version == null || version.isBlank()) { + builder.down().withDetail("reason", "PgVector extension not installed"); + return; + } + builder.up().withDetail("vectorExtensionVersion", version).withDetail("database", "postgresql"); + } + +} diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports index 1d6ddc67ec..3b28264c7f 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/main/resources/META-INF/spring/org.springframework.boot.autoconfigure.AutoConfiguration.imports @@ -14,3 +14,4 @@ # limitations under the License. # org.springframework.ai.vectorstore.pgvector.autoconfigure.PgVectorStoreAutoConfiguration +org.springframework.ai.vectorstore.pgvector.autoconfigure.PgVectorStoreHealthAutoConfiguration \ No newline at end of file diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/test/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfigurationTests.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/test/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfigurationTests.java new file mode 100644 index 0000000000..7d81f825d1 --- /dev/null +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-pgvector/src/test/java/org/springframework/ai/vectorstore/pgvector/autoconfigure/PgVectorStoreHealthAutoConfigurationTests.java @@ -0,0 +1,139 @@ +/* + * 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.pgvector.autoconfigure; + +import org.junit.jupiter.api.Test; + +import org.springframework.ai.vectorstore.pgvector.PgVectorStore; +import org.springframework.boot.autoconfigure.AutoConfigurations; +import org.springframework.boot.health.contributor.Health; +import org.springframework.boot.health.contributor.HealthIndicator; +import org.springframework.boot.health.contributor.Status; +import org.springframework.boot.test.context.runner.ApplicationContextRunner; +import org.springframework.context.annotation.Bean; +import org.springframework.context.annotation.Configuration; +import org.springframework.jdbc.core.JdbcTemplate; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isA; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Unit tests for {@link PgVectorStoreHealthAutoConfiguration}. + * + * @author jiajingda + */ +class PgVectorStoreHealthAutoConfigurationTests { + + private final ApplicationContextRunner contextRunner = new ApplicationContextRunner() + .withConfiguration(AutoConfigurations.of(PgVectorStoreHealthAutoConfiguration.class)) + .withUserConfiguration(MockBeansConfiguration.class); + + @Test + void healthIndicatorRegisteredWhenPgVectorStorePresent() { + this.contextRunner.run(context -> assertThat(context).hasSingleBean(PgVectorStoreHealthIndicator.class)); + } + + @Test + void healthIndicatorReportsUpWhenVectorExtensionInstalled() { + this.contextRunner.run(context -> { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + when(jdbcTemplate.queryForObject(isA(String.class), eq(String.class))).thenReturn("0.5.1"); + + HealthIndicator indicator = context.getBean(HealthIndicator.class); + Health health = indicator.health(); + + assertThat(health.getStatus()).isEqualTo(Status.UP); + assertThat(health.getDetails()).containsEntry("vectorExtensionVersion", "0.5.1"); + assertThat(health.getDetails()).containsEntry("database", "postgresql"); + }); + } + + @Test + void healthIndicatorReportsDownWhenVectorExtensionMissing() { + this.contextRunner.run(context -> { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + when(jdbcTemplate.queryForObject(isA(String.class), eq(String.class))).thenReturn(null); + + HealthIndicator indicator = context.getBean(HealthIndicator.class); + Health health = indicator.health(); + + assertThat(health.getStatus()).isEqualTo(Status.DOWN); + assertThat(health.getDetails()).containsEntry("reason", "PgVector extension not installed"); + }); + } + + @Test + void healthIndicatorReportsDownOnQueryException() { + this.contextRunner.run(context -> { + JdbcTemplate jdbcTemplate = context.getBean(JdbcTemplate.class); + when(jdbcTemplate.queryForObject(isA(String.class), eq(String.class))) + .thenThrow(new RuntimeException("Connection refused")); + + HealthIndicator indicator = context.getBean(HealthIndicator.class); + Health health = indicator.health(); + + assertThat(health.getStatus()).isEqualTo(Status.DOWN); + }); + } + + @Test + void healthIndicatorNotRegisteredWhenDisabled() { + this.contextRunner.withPropertyValues("management.health.pgvector.enabled=false").run(context -> { + assertThat(context).doesNotHaveBean(PgVectorStoreHealthIndicator.class); + assertThat(context).doesNotHaveBean("pgvectorHealthIndicator"); + }); + } + + @Test + void healthIndicatorBacksOffWhenCustomBeanProvided() { + this.contextRunner.withUserConfiguration(CustomHealthIndicatorConfiguration.class).run(context -> { + assertThat(context).hasSingleBean(HealthIndicator.class); + assertThat(context).doesNotHaveBean(PgVectorStoreHealthIndicator.class); + HealthIndicator indicator = context.getBean(HealthIndicator.class); + assertThat(indicator.health().getStatus()).isEqualTo(Status.UNKNOWN); + }); + } + + @Configuration(proxyBeanMethods = false) + static class MockBeansConfiguration { + + @Bean + PgVectorStore pgVectorStore() { + return mock(PgVectorStore.class); + } + + @Bean + JdbcTemplate jdbcTemplate() { + return mock(JdbcTemplate.class); + } + + } + + @Configuration(proxyBeanMethods = false) + static class CustomHealthIndicatorConfiguration { + + @Bean(name = "pgvectorHealthIndicator") + HealthIndicator pgvectorHealthIndicator() { + return () -> Health.unknown().build(); + } + + } + +}