wangyunpeng пре 6 дана
родитељ
комит
3c25563147

+ 2 - 12
core/src/main/java/com/tzld/videoVector/config/db/PgVectorDBConfig.java

@@ -1,13 +1,10 @@
 package com.tzld.videoVector.config.db;
 
-import com.tzld.videoVector.config.mybatis.HnswEfSearchInterceptor;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.ibatis.session.SqlSessionFactory;
 import org.mybatis.spring.SqlSessionFactoryBean;
 import org.springframework.beans.factory.annotation.Qualifier;
-import org.springframework.beans.factory.annotation.Value;
 import org.springframework.boot.context.properties.ConfigurationProperties;
-import org.springframework.boot.jdbc.DataSourceBuilder;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
@@ -21,18 +18,11 @@ import javax.sql.DataSource;
 @EnableTransactionManagement
 public class PgVectorDBConfig {
 
-    @Value("${spring.datasource.pg-vector.connection-init-sql:}")
-    private String connectionInitSql;
-
     // 1. 配置 pgVector 数据源
     @Bean(name = "pgVectorDataSource")
     @ConfigurationProperties(prefix = "spring.datasource.pg-vector")
     public DataSource pgVectorDataSource() {
-        HikariDataSource dataSource = (HikariDataSource) DataSourceBuilder.create().build();
-        if (connectionInitSql != null && !connectionInitSql.isEmpty()) {
-            dataSource.setConnectionInitSql(connectionInitSql);
-        }
-        return dataSource;
+        return new HikariDataSource();
     }
 
     // 2. 配置 pgVector 专属 SqlSessionFactory(独立 Configuration,避免与 MySQL 数据源互相污染)
@@ -57,7 +47,7 @@ public class PgVectorDBConfig {
         org.apache.ibatis.session.Configuration configuration = new org.apache.ibatis.session.Configuration();
         configuration.setMapUnderscoreToCamelCase(true);
         configuration.setUseGeneratedKeys(true);
-        configuration.addInterceptor(new HnswEfSearchInterceptor());
+        configuration.addInterceptor(new com.tzld.videoVector.config.mybatis.HnswEfSearchInterceptor());
         sessionFactory.setConfiguration(configuration);
         sessionFactory.setTypeAliasesPackage("com.tzld.videoVector");
         return sessionFactory.getObject();

+ 25 - 8
core/src/main/java/com/tzld/videoVector/config/mybatis/HnswEfSearchInterceptor.java

@@ -1,30 +1,47 @@
 package com.tzld.videoVector.config.mybatis;
 
-import org.apache.ibatis.executor.statement.StatementHandler;
-import org.apache.ibatis.plugin.Interceptor;
-import org.apache.ibatis.plugin.Intercepts;
-import org.apache.ibatis.plugin.Invocation;
-import org.apache.ibatis.plugin.Signature;
+import org.apache.ibatis.executor.Executor;
+import org.apache.ibatis.mapping.MappedStatement;
+import org.apache.ibatis.plugin.*;
+import org.apache.ibatis.session.ResultHandler;
+import org.apache.ibatis.session.RowBounds;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.sql.Connection;
 import java.sql.Statement;
+import java.util.Properties;
 
 @Intercepts({
-    @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
+    @Signature(type = Executor.class, method = "query",
+        args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
 })
 public class HnswEfSearchInterceptor implements Interceptor {
 
+    private static final Logger log = LoggerFactory.getLogger(HnswEfSearchInterceptor.class);
+
+    private volatile boolean logged;
+
     @Override
     public Object intercept(Invocation invocation) throws Throwable {
-        Connection connection = (Connection) invocation.getArgs()[0];
+        Executor executor = (Executor) invocation.getTarget();
+        Connection connection = executor.getTransaction().getConnection();
         try (Statement stmt = connection.createStatement()) {
             stmt.execute("SET hnsw.ef_search = 200");
         }
+        if (!logged) {
+            logged = true;
+            log.info("HnswEfSearchInterceptor activated: hnsw.ef_search = 200");
+        }
         return invocation.proceed();
     }
 
     @Override
     public Object plugin(Object target) {
-        return org.apache.ibatis.plugin.Plugin.wrap(target, this);
+        return Plugin.wrap(target, this);
+    }
+
+    @Override
+    public void setProperties(Properties properties) {
     }
 }