nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/requesters/utils.py

@@ -18,23 +18,14 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
-from enum import Enum
-from typing import Any, Optional, TypeVar, Union, overload
+import json
+from enum import Enum, auto
+from typing import Any, Optional, Sequence, TypeVar, Union, overload
 
 from fastapi import HTTPException
+from google.protobuf.json_format import MessageToDict
 from grpc import StatusCode as GrpcStatusCode
 from grpc.aio import AioRpcError
-from nucliadb_protos.nodereader_pb2 import (
-    ParagraphSearchRequest,
-    ParagraphSearchResponse,
-    RelationSearchRequest,
-    RelationSearchResponse,
-    SearchRequest,
-    SearchResponse,
-    SuggestRequest,
-    SuggestResponse,
-)
-from nucliadb_protos.writer_pb2 import ShardObject as PBShardObject
 
 from nucliadb.common.cluster import manager as cluster_manager
 from nucliadb.common.cluster.base import AbstractIndexNode
@@ -42,64 +33,51 @@ from nucliadb.common.cluster.exceptions import ShardsNotFound
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.search import logger
 from nucliadb.search.search.shards import (
-    query_paragraph_shard,
     query_shard,
-    relations_shard,
     suggest_shard,
 )
 from nucliadb.search.settings import settings
+from nucliadb_protos.nodereader_pb2 import (
+    SearchRequest,
+    SearchResponse,
+    SuggestRequest,
+    SuggestResponse,
+)
+from nucliadb_protos.writer_pb2 import ShardObject as PBShardObject
 from nucliadb_telemetry import errors
 from nucliadb_utils import const
 from nucliadb_utils.utilities import has_feature
 
 
 class Method(Enum):
-    SEARCH = 1
-    PARAGRAPH = 2
-    SUGGEST = 3
-    RELATIONS = 4
+    SEARCH = auto()
+    SUGGEST = auto()
 
 
 METHODS = {
     Method.SEARCH: query_shard,
-    Method.PARAGRAPH: query_paragraph_shard,
     Method.SUGGEST: suggest_shard,
-    Method.RELATIONS: relations_shard,
 }
 
-REQUEST_TYPE = Union[
-    SuggestRequest, ParagraphSearchRequest, SearchRequest, RelationSearchRequest
-]
+REQUEST_TYPE = Union[SuggestRequest, SearchRequest]
 
 T = TypeVar(
     "T",
     SuggestResponse,
-    ParagraphSearchResponse,
     SearchResponse,
-    RelationSearchResponse,
 )
 
 
-@overload  # type: ignore
-async def node_query(
-    kbid: str,
-    method: Method,
-    pb_query: SuggestRequest,
-    target_shard_replicas: Optional[list[str]] = None,
-    use_read_replica_nodes: bool = True,
-) -> tuple[list[SuggestResponse], bool, list[tuple[AbstractIndexNode, str]]]:
-    ...
-
-
 @overload
 async def node_query(
     kbid: str,
     method: Method,
-    pb_query: ParagraphSearchRequest,
+    pb_query: SuggestRequest,
     target_shard_replicas: Optional[list[str]] = None,
     use_read_replica_nodes: bool = True,
-) -> tuple[list[ParagraphSearchResponse], bool, list[tuple[AbstractIndexNode, str]]]:
-    ...
+    timeout: Optional[float] = None,
+    retry_on_primary: bool = True,
+) -> tuple[list[SuggestResponse], bool, list[tuple[AbstractIndexNode, str]]]: ...
 
 
 @overload
@@ -109,19 +87,9 @@ async def node_query(
     pb_query: SearchRequest,
     target_shard_replicas: Optional[list[str]] = None,
     use_read_replica_nodes: bool = True,
-) -> tuple[list[SearchResponse], bool, list[tuple[AbstractIndexNode, str]]]:
-    ...
-
-
-@overload
-async def node_query(
-    kbid: str,
-    method: Method,
-    pb_query: RelationSearchRequest,
-    target_shard_replicas: Optional[list[str]] = None,
-    use_read_replica_nodes: bool = True,
-) -> tuple[list[RelationSearchResponse], bool, list[tuple[AbstractIndexNode, str]]]:
-    ...
+    timeout: Optional[float] = None,
+    retry_on_primary: bool = True,
+) -> tuple[list[SearchResponse], bool, list[tuple[AbstractIndexNode, str]]]: ...
 
 
 async def node_query(
@@ -130,7 +98,10 @@ async def node_query(
     pb_query: REQUEST_TYPE,
     target_shard_replicas: Optional[list[str]] = None,
     use_read_replica_nodes: bool = True,
-) -> tuple[list[T], bool, list[tuple[AbstractIndexNode, str]]]:
+    timeout: Optional[float] = None,
+    retry_on_primary: bool = True,
+) -> tuple[Sequence[Union[T, BaseException]], bool, list[tuple[AbstractIndexNode, str]]]:
+    timeout = timeout or settings.search_timeout
     use_read_replica_nodes = use_read_replica_nodes and has_feature(
         const.Features.READ_REPLICA_SEARCHES, context={"kbid": kbid}
     )
@@ -152,6 +123,7 @@ async def node_query(
         try:
             node, shard_id = cluster_manager.choose_node(
                 shard_obj,
+                use_nidx=has_feature(const.Features.NIDX_READS, context={"kbid": kbid}),
                 use_read_replica_nodes=use_read_replica_nodes,
                 target_shard_replicas=target_shard_replicas,
             )
@@ -173,9 +145,9 @@ async def node_query(
         )
 
     try:
-        results = await asyncio.wait_for(  # type: ignore
-            asyncio.gather(*ops, return_exceptions=True),  # type: ignore
-            timeout=settings.search_timeout,
+        results: list[Union[T, BaseException]] = await asyncio.wait_for(
+            asyncio.gather(*ops, return_exceptions=True),
+            timeout=timeout,
         )
     except asyncio.TimeoutError as exc:  # pragma: no cover
         logger.warning(
@@ -186,10 +158,20 @@ async def node_query(
 
     error = validate_node_query_results(results or [])
     if error is not None:
+        query_dict = MessageToDict(pb_query)
+        query_dict.pop("vector", None)
+        logger.error(
+            "Error while querying nodes",
+            extra={
+                "kbid": kbid,
+                "query": json.dumps(query_dict),
+            },
+        )
         if (
             error.status_code >= 500
             and use_read_replica_nodes
             and any([node.is_read_replica() for node, _ in queried_nodes])
+            and retry_on_primary
         ):
             # We had an error querying a secondary node, instead of raising an
             # error directly, retry query to primaries and hope it works
@@ -220,9 +202,7 @@ def validate_node_query_results(results: list[Any]) -> Optional[HTTPException]:
     Handling of exception is responsibility of caller.
     """
     if results is None or len(results) == 0:
-        return HTTPException(
-            status_code=500, detail=f"Error while executing shard queries. No results."
-        )
+        return HTTPException(status_code=500, detail=f"Error while executing shard queries. No results.")
 
     for result in results:
         if isinstance(result, Exception):
@@ -251,9 +231,7 @@ def validate_node_query_results(results: list[Any]) -> Optional[HTTPException]:
     return None
 
 
-def debug_nodes_info(
-    nodes: list[tuple[AbstractIndexNode, str]]
-) -> list[dict[str, str]]:
+def debug_nodes_info(nodes: list[tuple[AbstractIndexNode, str]]) -> list[dict[str, str]]:
     details: list[dict[str, str]] = []
     for node, shard_id in nodes:
         info = {

nucliadb/search/search/cache.py

@@ -18,42 +18,77 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import asyncio
+import contextlib
+import logging
 from contextvars import ContextVar
 from typing import Optional
 
-from lru import LRU  # type: ignore
+from lru import LRU
 
-from nucliadb.common.maindb.driver import Transaction
+from nucliadb.common.ids import FieldId
+from nucliadb.common.maindb.utils import get_driver
+from nucliadb.ingest.fields.base import Field
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
-from nucliadb.middleware.transaction import get_read_only_transaction
 from nucliadb.search import SERVICE_NAME
+from nucliadb_protos.utils_pb2 import ExtractedText
 from nucliadb_telemetry import metrics
 from nucliadb_utils.utilities import get_storage
 
-rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar(
-    "rcache", default=None
-)
+logger = logging.getLogger(__name__)
+
+rcache: ContextVar[Optional[dict[str, ResourceORM]]] = ContextVar("rcache", default=None)
+etcache: ContextVar[Optional["ExtractedTextCache"]] = ContextVar("etcache", default=None)
 
 
 RESOURCE_LOCKS: dict[str, asyncio.Lock] = LRU(1000)  # type: ignore
 RESOURCE_CACHE_OPS = metrics.Counter("nucliadb_resource_cache_ops", labels={"type": ""})
+EXTRACTED_CACHE_OPS = metrics.Counter("nucliadb_extracted_text_cache_ops", labels={"type": ""})
+
+
+def set_extracted_text_cache() -> None:
+    value = ExtractedTextCache()
+    etcache.set(value)
+
+
+def get_extracted_text_cache() -> Optional["ExtractedTextCache"]:
+    return etcache.get()
+
+
+def clear_extracted_text_cache() -> None:
+    value = etcache.get()
+    if value is not None:
+        value.clear()
+        etcache.set(None)
+
 
+def set_resource_cache() -> None:
+    value: dict[str, ResourceORM] = {}
+    rcache.set(value)
 
-def get_resource_cache(clear: bool = False) -> dict[str, ResourceORM]:
-    value: Optional[dict[str, ResourceORM]] = rcache.get()
-    if value is None or clear:
-        value = {}
-        rcache.set(value)
-    return value
 
+def get_resource_cache() -> Optional[dict[str, ResourceORM]]:
+    return rcache.get()
 
-async def get_resource_from_cache(
-    kbid: str, uuid: str, txn: Optional[Transaction] = None
-) -> Optional[ResourceORM]:
+
+def clear_resource_cache() -> None:
+    value = rcache.get()
+    if value is not None:
+        value.clear()
+        rcache.set(None)
+
+
+async def get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
+    """
+    Will try to get the resource from the cache, if it's not there it will fetch it from the ORM and cache it.
+    """
     orm_resource: Optional[ResourceORM] = None
 
     resource_cache = get_resource_cache()
+    if resource_cache is None:
+        RESOURCE_CACHE_OPS.inc({"type": "miss"})
+        logger.warning("Resource cache not set")
+        return await _orm_get_resource(kbid, uuid)
 
     if uuid not in RESOURCE_LOCKS:
         RESOURCE_LOCKS[uuid] = asyncio.Lock()
@@ -61,11 +96,7 @@ async def get_resource_from_cache(
     async with RESOURCE_LOCKS[uuid]:
         if uuid not in resource_cache:
             RESOURCE_CACHE_OPS.inc({"type": "miss"})
-            if txn is None:
-                txn = await get_read_only_transaction()
-            storage = await get_storage(service_name=SERVICE_NAME)
-            kb = KnowledgeBoxORM(txn, storage, kbid)
-            orm_resource = await kb.get(uuid)
+            orm_resource = await _orm_get_resource(kbid, uuid)
         else:
             RESOURCE_CACHE_OPS.inc({"type": "hit"})
 
@@ -75,3 +106,101 @@ async def get_resource_from_cache(
             orm_resource = resource_cache.get(uuid)
 
     return orm_resource
+
+
+async def _orm_get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
+    async with get_driver().transaction(read_only=True) as txn:
+        storage = await get_storage(service_name=SERVICE_NAME)
+        kb = KnowledgeBoxORM(txn, storage, kbid)
+        return await kb.get(uuid)
+
+
+class ExtractedTextCache:
+    """
+    Used to cache extracted text from a resource in memory during the process
+    of search results hydration.
+
+    This is needed to avoid fetching the same extracted text multiple times,
+    as matching text blocks are processed in parallel and the extracted text is
+    fetched for each field where the text block is found.
+    """
+
+    def __init__(self):
+        self.locks = {}
+        self.values = {}
+
+    def get_value(self, key: str) -> Optional[ExtractedText]:
+        return self.values.get(key)
+
+    def get_lock(self, key: str) -> asyncio.Lock:
+        return self.locks.setdefault(key, asyncio.Lock())
+
+    def set_value(self, key: str, value: ExtractedText) -> None:
+        self.values[key] = value
+
+    def clear(self):
+        self.values.clear()
+        self.locks.clear()
+
+
+async def get_field_extracted_text(field: Field) -> Optional[ExtractedText]:
+    cache = get_extracted_text_cache()
+    if cache is None:
+        logger.warning("Extracted text cache not set")
+        EXTRACTED_CACHE_OPS.inc({"type": "miss"})
+        return await field.get_extracted_text()
+
+    key = f"{field.kbid}/{field.uuid}/{field.id}"
+    extracted_text = cache.get_value(key)
+    if extracted_text is not None:
+        EXTRACTED_CACHE_OPS.inc({"type": "hit"})
+        return extracted_text
+
+    async with cache.get_lock(key):
+        # Check again in case another task already fetched it
+        extracted_text = cache.get_value(key)
+        if extracted_text is not None:
+            EXTRACTED_CACHE_OPS.inc({"type": "hit"})
+            return extracted_text
+
+        EXTRACTED_CACHE_OPS.inc({"type": "miss"})
+        extracted_text = await field.get_extracted_text()
+        if extracted_text is not None:
+            # Only cache if we actually have extracted text
+            cache.set_value(key, extracted_text)
+        return extracted_text
+
+
+async def get_extracted_text_from_field_id(kbid: str, field: FieldId) -> Optional[ExtractedText]:
+    rid = field.rid
+    orm_resource = await get_resource(kbid, rid)
+    if orm_resource is None:
+        return None
+    field_obj = await orm_resource.get_field(
+        key=field.key,
+        type=field.pb_type,
+        load=False,
+    )
+    return await get_field_extracted_text(field_obj)
+
+
+@contextlib.contextmanager
+def request_caches():
+    """
+    This context manager sets the caches for extracted text and resources for a request.
+
+    It should used at the beginning of a request handler to avoid fetching the same
+    resources and extracted text multiple times.
+
+    Makes sure to clean the caches at the end of the context manager.
+    >>> with request_caches():
+    ...     resource = await get_resource(kbid, uuid)
+    ...     extracted_text = await get_extracted_text_from_field_id(kbid, rid, field_id)
+    """
+    set_resource_cache()
+    set_extracted_text_cache()
+    try:
+        yield
+    finally:
+        clear_resource_cache()
+        clear_extracted_text_cache()