nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/api/v1/suggest.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
-from time import time
 from typing import Optional, Union
 
 from fastapi import Header, Request, Response
@@ -28,9 +27,11 @@ from nucliadb.models.responses import HTTPClientError
 from nucliadb.search.api.v1.router import KB_PREFIX, api
 from nucliadb.search.api.v1.utils import fastapi_query
 from nucliadb.search.requesters.utils import Method, node_query
+from nucliadb.search.search import cache
 from nucliadb.search.search.exceptions import InvalidQueryError
 from nucliadb.search.search.merge import merge_suggest_results
 from nucliadb.search.search.query import suggest_query_to_pb
+from nucliadb.search.search.utils import filter_hidden_resources
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.search import (
@@ -40,8 +41,8 @@ from nucliadb_models.search import (
     SearchParamDefaults,
     SuggestOptions,
 )
+from nucliadb_models.utils import DateTime
 from nucliadb_utils.authentication import requires
-from nucliadb_utils.utilities import get_audit
 
 
 @api.get(
@@ -63,21 +64,15 @@ async def suggest_knowledgebox(
     fields: list[str] = fastapi_query(SearchParamDefaults.fields),
     filters: list[str] = fastapi_query(SearchParamDefaults.filters),
     faceted: list[str] = fastapi_query(SearchParamDefaults.faceted),
-    range_creation_start: Optional[datetime] = fastapi_query(
-        SearchParamDefaults.range_creation_start
-    ),
-    range_creation_end: Optional[datetime] = fastapi_query(
-        SearchParamDefaults.range_creation_end
-    ),
-    range_modification_start: Optional[datetime] = fastapi_query(
+    range_creation_start: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_start),
+    range_creation_end: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_end),
+    range_modification_start: Optional[DateTime] = fastapi_query(
         SearchParamDefaults.range_modification_start
     ),
-    range_modification_end: Optional[datetime] = fastapi_query(
+    range_modification_end: Optional[DateTime] = fastapi_query(
         SearchParamDefaults.range_modification_end
     ),
-    features: list[SuggestOptions] = fastapi_query(
-        SearchParamDefaults.suggest_features
-    ),
+    features: list[SuggestOptions] = fastapi_query(SearchParamDefaults.suggest_features),
     show: list[ResourceProperties] = fastapi_query(SearchParamDefaults.show),
     field_type_filter: list[FieldTypeName] = fastapi_query(
         SearchParamDefaults.field_type_filter, alias="field_type"
@@ -87,6 +82,7 @@ async def suggest_knowledgebox(
     x_forwarded_for: str = Header(""),
     debug: bool = fastapi_query(SearchParamDefaults.debug),
     highlight: bool = fastapi_query(SearchParamDefaults.highlight),
+    show_hidden: bool = fastapi_query(SearchParamDefaults.show_hidden),
 ) -> Union[KnowledgeboxSuggestResults, HTTPClientError]:
     try:
         return await suggest(
@@ -108,6 +104,7 @@ async def suggest_knowledgebox(
             x_forwarded_for,
             debug,
             highlight,
+            show_hidden,
         )
     except InvalidQueryError as exc:
         return HTTPClientError(status_code=412, detail=str(exc))
@@ -132,49 +129,37 @@ async def suggest(
     x_forwarded_for: str,
     debug: bool,
     highlight: bool,
+    show_hidden: bool,
 ) -> KnowledgeboxSuggestResults:
-    # We need the nodes/shards that are connected to the KB
-    audit = get_audit()
-    start_time = time()
-
-    # We need to query all nodes
-    pb_query = suggest_query_to_pb(
-        features,
-        query,
-        fields,
-        filters,
-        faceted,
-        range_creation_start,
-        range_creation_end,
-        range_modification_start,
-        range_modification_end,
-    )
-    results, incomplete_results, queried_nodes = await node_query(
-        kbid, Method.SUGGEST, pb_query
-    )
-
-    # We need to merge
-    search_results = await merge_suggest_results(
-        results,
-        kbid=kbid,
-        show=show,
-        field_type_filter=field_type_filter,
-        highlight=highlight,
-    )
+    with cache.request_caches():
+        hidden = await filter_hidden_resources(kbid, show_hidden)
+        pb_query = suggest_query_to_pb(
+            features,
+            query,
+            fields,
+            filters,
+            faceted,
+            range_creation_start,
+            range_creation_end,
+            range_modification_start,
+            range_modification_end,
+            hidden,
+        )
+        results, incomplete_results, queried_nodes = await node_query(kbid, Method.SUGGEST, pb_query)
 
-    response.status_code = 206 if incomplete_results else 200
+        # We need to merge
+        search_results = await merge_suggest_results(
+            results,
+            kbid=kbid,
+            show=show,
+            field_type_filter=field_type_filter,
+            highlight=highlight,
+        )
 
-    queried_shards = [shard_id for _, shard_id in queried_nodes]
-    if debug and queried_shards:
-        search_results.shards = queried_shards
+        response.status_code = 206 if incomplete_results else 200
 
-    if audit is not None:
-        await audit.suggest(
-            kbid,
-            x_nucliadb_user,
-            x_ndb_client.to_proto(),
-            x_forwarded_for,
-            time() - start_time,
-        )
+        queried_shards = [shard_id for _, shard_id in queried_nodes]
+        if debug and queried_shards:
+            search_results.shards = queried_shards
 
-    return search_results
+        return search_results

nucliadb/search/app.py

@@ -17,72 +17,52 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
-import pkg_resources
+import importlib.metadata
+
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
-from starlette.middleware.cors import CORSMiddleware
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import HTMLResponse
 
 from nucliadb.common.cluster import manager
+from nucliadb.ingest.settings import DriverConfig
 from nucliadb.middleware import ProcessTimeHeaderMiddleware
-from nucliadb.middleware.transaction import ReadOnlyTransactionMiddleware
 from nucliadb.search import API_PREFIX
 from nucliadb.search.api.v1.router import api as api_v1
-from nucliadb.search.lifecycle import finalize, initialize
+from nucliadb.search.lifecycle import lifespan
 from nucliadb.search.settings import settings
 from nucliadb_telemetry import errors
 from nucliadb_telemetry.fastapi.utils import (
     client_disconnect_handler,
     global_exception_handler,
 )
-from nucliadb_utils import const
+from nucliadb_utils.audit.stream import AuditMiddleware
 from nucliadb_utils.authentication import NucliaCloudAuthenticationBackend
 from nucliadb_utils.fastapi.openapi import extend_openapi
 from nucliadb_utils.fastapi.versioning import VersionedFastAPI
-from nucliadb_utils.settings import http_settings, running_settings
-from nucliadb_utils.utilities import has_feature
+from nucliadb_utils.settings import running_settings
+from nucliadb_utils.utilities import get_audit
 
 middleware = []
-
-if has_feature(const.Features.CORS_MIDDLEWARE, default=False):
-    middleware.append(
-        Middleware(
-            CORSMiddleware,
-            allow_origins=http_settings.cors_origins,
-            allow_methods=["*"],
-            # Authorization will be exluded from * in the future, (CORS non-wildcard request-header).
-            # Browsers already showing deprecation notices, so it needs to be specified explicitly
-            allow_headers=["*", "Authorization"],
-        )
-    )
-
 middleware.extend(
     [
-        Middleware(
-            AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()
-        ),
-        Middleware(ReadOnlyTransactionMiddleware),
+        Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend()),
+        Middleware(AuditMiddleware, audit_utility_getter=get_audit),
     ]
 )
 
 if running_settings.debug:
     middleware.append(Middleware(ProcessTimeHeaderMiddleware))
 
-errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
-
-
-on_startup = [initialize]
-on_shutdown = [finalize]
+errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
 
 fastapi_settings = dict(
     debug=running_settings.debug,
     middleware=middleware,
-    on_startup=on_startup,
-    on_shutdown=on_shutdown,
+    lifespan=lifespan,
     exception_handlers={
         Exception: global_exception_handler,
         ClientDisconnect: client_disconnect_handler,
@@ -127,7 +107,7 @@ async def node_members(request: Request) -> JSONResponse:
 
 
 async def alive(request: Request) -> JSONResponse:
-    if len(manager.get_index_nodes()) == 0 and settings.driver != "local":
+    if len(manager.get_index_nodes()) == 0 and settings.driver != DriverConfig.LOCAL:
         return JSONResponse({"status": "error"}, status_code=503)
     else:
         return JSONResponse({"status": "ok"})

nucliadb/search/lifecycle.py

@@ -17,8 +17,13 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
-from nucliadb.common.maindb.utils import setup_driver  # type: ignore
+from nucliadb.common.maindb.utils import setup_driver
+from nucliadb.common.nidx import start_nidx_utility
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.search import SERVICE_NAME
 from nucliadb.search.predict import start_predict_engine
@@ -33,7 +38,8 @@ from nucliadb_utils.utilities import (
 )
 
 
-async def initialize() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
 
     await start_ingest(SERVICE_NAME)
@@ -41,11 +47,12 @@ async def initialize() -> None:
 
     await setup_driver()
     await setup_cluster()
+    await start_nidx_utility()
 
     await start_audit_utility(SERVICE_NAME)
 
+    yield
 
-async def finalize() -> None:
     await stop_ingest()
     if get_utility(Utility.PARTITION):
         clean_utility(Utility.PARTITION)