optimize milvus connector

Xiong-Da · Xiong-Da · commit 50df8ec1f4fb · 2026-04-11T17:54:01.000+08:00
diff --git a/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/calcite/MilvusCalciteTableFactory.java b/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/calcite/MilvusCalciteTableFactory.java
@@ -8,7 +8,6 @@
 import org.apache.calcite.schema.Table;
 
 import java.util.ArrayList;
-import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
diff --git a/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/flink/MilvusDynamicTableFactory.java b/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/flink/MilvusDynamicTableFactory.java
@@ -28,7 +28,6 @@ private MilvusConfig getMilvusConfig(Context context) {
         Map<String, String> options = context.getCatalogTable().getOptions();
         ResolvedSchema tableSchema = context.getCatalogTable().getResolvedSchema();
         MilvusConfig milvusConfig = MilvusOptions.getMilvusConfig(options);
-        milvusConfig.database = context.getObjectIdentifier().getDatabaseName();
         milvusConfig.fieldSchemas = FlinkSchemaUtils.getFieldSchemas(tableSchema);
         milvusConfig.primaryKey = FlinkSchemaUtils.getPrimaryKey(tableSchema);
         milvusConfig.primaryKeyIndex = HiveTableUtils.getTablePrimaryKeyIndex(
diff --git a/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/handler/MilvusHandler.java b/sqlrec-connectors/sqlrec-connector-milvus/src/main/java/com/sqlrec/connectors/milvus/handler/MilvusHandler.java
@@ -20,13 +20,16 @@
 import io.milvus.v2.service.vector.response.SearchResp;
 import org.apache.calcite.DataContext;
 import org.apache.calcite.rex.RexNode;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.time.Duration;
 import java.util.*;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.stream.Collectors;
 
 public class MilvusHandler {
+    private static final Logger logger = LoggerFactory.getLogger(MilvusHandler.class);
     private static Map<String, MilvusClientV2Pool> clientPools = new ConcurrentHashMap<>();
     private static Gson gson = new GsonBuilder()
             .setObjectToNumberStrategy(ToNumberPolicy.LONG_OR_DOUBLE)
@@ -236,20 +239,27 @@ private Object[] toRowWithScore(Map<String, Object> entity, List<FieldSchema> fi
     }
 
     public static void returnClient(MilvusClientV2 client, MilvusConfig milvusConfig) {
-        String key = milvusConfig.url + milvusConfig.token;
+        String key = getClientPoolKey(milvusConfig);
         if (clientPools.containsKey(key)) {
             clientPools.get(key).returnClient(key, client);
+        } else {
+            logger.warn("client pool {} is not found", key);
+            client.close();
         }
     }
 
     public static MilvusClientV2 getClient(MilvusConfig milvusConfig) {
-        String key = milvusConfig.url + milvusConfig.token;
+        String key = getClientPoolKey(milvusConfig);
         if (!clientPools.containsKey(key)) {
             openClientPool(key, milvusConfig);
         }
         return clientPools.get(key).getClient(key);
     }
 
+    private static String getClientPoolKey(MilvusConfig milvusConfig) {
+        return milvusConfig.url + "|" + milvusConfig.token;
+    }
+
     private static synchronized void openClientPool(String key, MilvusConfig milvusConfig) {
         if (clientPools.containsKey(key)) {
             return;
