Fix: use Infinity connPool.get_conn() in metadata filter push-down (#18438)

api/db/services/doc_metadata_service.py

@@ -24,16 +24,15 @@ import json
 import logging
 import re
 from copy import deepcopy
-from typing import Any, Dict, List, Optional
+from typing import Any
 
-from api.db.db_models import DB, Document
+from api.db.db_models import DB, Document, Knowledgebase
 from common import settings
-from common.metadata_utils import dedupe_list
-from api.db.db_models import Knowledgebase
 from common.doc_store.doc_store_base import OrderByExpr
+from common.metadata_utils import dedupe_list
 
 
-def _es_response_total(response: Any) -> Optional[int]:
+def _es_response_total(response: Any) -> int | None:
     """Extract the exact total hit count from an ES search response.
 
     Returns ``None`` when the field is missing or in an unexpected shape
@@ -80,7 +79,7 @@ class DocMetadataService:
         return f"ragflow_doc_meta_{tenant_id}"
 
     @staticmethod
-    def _extract_metadata(flat_meta: Dict) -> Dict:
+    def _extract_metadata(flat_meta: dict) -> dict:
         """
         Extract metadata from ES/Infinity document format.
 
@@ -113,7 +112,7 @@ class DocMetadataService:
         return {}
 
     @staticmethod
-    def _extract_doc_id(doc: Dict, hit: Dict = None) -> str:
+    def _extract_doc_id(doc: dict, hit: dict = None) -> str:
         """
         Extract document ID from various formats.
 
@@ -188,7 +187,7 @@ class DocMetadataService:
                     yield doc_id, doc
 
     @classmethod
-    def _search_metadata(cls, kb_id: str, condition: Dict = None):
+    def _search_metadata(cls, kb_id: str, condition: dict = None):
         """
         Common search logic for metadata queries.
         Uses pagination internally to retrieve data from the index.
@@ -302,7 +301,7 @@ class DocMetadataService:
         return all_results
 
     @classmethod
-    def _split_combined_values(cls, meta_fields: Dict) -> Dict:
+    def _split_combined_values(cls, meta_fields: dict) -> dict:
         """
         Post-process metadata to split combined values by common delimiters.
 
@@ -350,7 +349,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def insert_document_metadata(cls, doc_id: str, meta_fields: Dict) -> bool:
+    def insert_document_metadata(cls, doc_id: str, meta_fields: dict) -> bool:
         """
         Insert document metadata into ES/Infinity.
 
@@ -439,7 +438,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def update_document_metadata(cls, doc_id: str, meta_fields: Dict) -> bool:
+    def update_document_metadata(cls, doc_id: str, meta_fields: dict) -> bool:
         """
         Update document metadata in ES/Infinity.
 
@@ -709,7 +708,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def get_document_metadata(cls, doc_id: str) -> Dict:
+    def get_document_metadata(cls, doc_id: str) -> dict:
         """
         Get document metadata from ES/Infinity.
 
@@ -749,7 +748,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def get_flatted_meta_by_kbs(cls, kb_ids: List[str]) -> Dict:
+    def get_flatted_meta_by_kbs(cls, kb_ids: list[str]) -> dict:
         """
         Get flattened metadata for documents in knowledge bases.
 
@@ -846,11 +845,11 @@ class DocMetadataService:
     @classmethod
     def filter_doc_ids_by_meta_pushdown(
         cls,
-        kb_ids: List[str],
-        filters: List[Dict],
+        kb_ids: list[str],
+        filters: list[dict],
         logic: str = "and",
         limit: int = 10000,
-    ) -> Optional[List[str]]:
+    ) -> list[str] | None:
         """Run a metadata filter directly against ES or Infinity, returning matching doc IDs.
 
         Returns ``None`` to signal "push-down not viable, use the in-memory
@@ -893,11 +892,11 @@ class DocMetadataService:
     def _filter_doc_ids_by_metadata_es(
         cls,
         index_name: str,
-        kb_ids: List[str],
-        filters: List[Dict],
+        kb_ids: list[str],
+        filters: list[dict],
         logic: str,
         limit: int,
-    ) -> Optional[List[str]]:
+    ) -> list[str] | None:
         """ES push-down path for metadata filtering."""
         from common.metadata_es_filter import (
             UnsupportedMetaFilter,
@@ -937,7 +936,7 @@ class DocMetadataService:
 
         doc_ids = extract_doc_ids(response if isinstance(response, dict) else dict(response))
         seen: set[str] = set()
-        unique: List[str] = []
+        unique: list[str] = []
         for did in doc_ids:
             if did in seen:
                 continue
@@ -967,11 +966,11 @@ class DocMetadataService:
     def _filter_doc_ids_by_metadata_gaussdb(
         cls,
         index_name: str,
-        kb_ids: List[str],
-        filters: List[Dict],
+        kb_ids: list[str],
+        filters: list[dict],
         logic: str,
         limit: int,
-    ) -> Optional[List[str]]:
+    ) -> list[str] | None:
         """GaussDB push-down path for metadata filtering."""
         from common.metadata_gaussdb_filter import (
             UnsupportedGaussDBMetaFilter,
@@ -997,7 +996,7 @@ class DocMetadataService:
             return None
 
         seen: set[str] = set()
-        unique: List[str] = []
+        unique: list[str] = []
         for did in doc_ids:
             if did in seen:
                 continue
@@ -1015,10 +1014,10 @@ class DocMetadataService:
     def _filter_doc_ids_by_metadata_infinity(
         cls,
         index_name: str,
-        kb_ids: List[str],
-        filters: List[Dict],
+        kb_ids: list[str],
+        filters: list[dict],
         logic: str,
-    ) -> Optional[List[str]]:
+    ) -> list[str] | None:
         """Infinity push-down path for metadata filtering."""
         from common.metadata_infinity_filter import (
             build_infinity_filter,
@@ -1036,7 +1035,7 @@ class DocMetadataService:
             where_clause = f"{kb_filter} AND {sql_filter}"
             logging.debug(f"Infinity metadata filter: {where_clause}")
 
-            inf_conn = settings.docStoreConn.acquire_conn()
+            inf_conn = settings.docStoreConn.connPool.get_conn()
             try:
                 db_instance = inf_conn.get_database(settings.docStoreConn.dbName)
                 table_instance = db_instance.get_table(index_name)
@@ -1051,7 +1050,7 @@ class DocMetadataService:
             return None
 
     @classmethod
-    def get_metadata_keys_by_kbs(cls, kb_ids: List[str]) -> List[str]:
+    def get_metadata_keys_by_kbs(cls, kb_ids: list[str]) -> list[str]:
         """
         Get unique metadata field names across multiple knowledge bases.
 
@@ -1081,7 +1080,7 @@ class DocMetadataService:
             return []
 
     @classmethod
-    def get_metadata_for_documents(cls, doc_ids: Optional[List[str]], kb_id: str) -> Dict[str, Dict]:
+    def get_metadata_for_documents(cls, doc_ids: list[str] | None, kb_id: str) -> dict[str, dict]:
         """
         Get metadata fields for specific documents.
         Returns a mapping of doc_id -> meta_fields
@@ -1120,7 +1119,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def get_metadata_summary(cls, kb_id: str, doc_ids=None) -> Dict:
+    def get_metadata_summary(cls, kb_id: str, doc_ids=None) -> dict:
         """
         Get metadata summary for documents in a knowledge base.
 
@@ -1213,7 +1212,7 @@ class DocMetadataService:
 
     @classmethod
     @DB.connection_context()
-    def batch_update_metadata(cls, kb_id: str, doc_ids: List[str], updates=None, deletes=None) -> int:
+    def batch_update_metadata(cls, kb_id: str, doc_ids: list[str], updates=None, deletes=None) -> int:
         """
         Batch update metadata for documents in a knowledge base.