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

nucliadb/writer/app.py

@@ -18,74 +18,38 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-import functools
+import importlib.metadata
 
-import pkg_resources
 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.requests import ClientDisconnect
 from starlette.responses import HTMLResponse
 
-from nucliadb.common.context.fastapi import get_app_context, set_app_context
 from nucliadb.writer import API_PREFIX
 from nucliadb.writer.api.v1.router import api as api_v1
-from nucliadb.writer.lifecycle import finalize, initialize
+from nucliadb.writer.lifecycle import lifespan
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
+from nucliadb_telemetry.fastapi.utils import (
+    client_disconnect_handler,
+    global_exception_handler,
+)
 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
 
 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())]
-)
-
-
-errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
+middleware.extend([Middleware(AuthenticationMiddleware, backend=NucliaCloudAuthenticationBackend())])
 
-on_startup = [initialize]
-on_shutdown = [finalize]
-
-
-async def global_exception_handler(request: Request, exc: Exception):
-    errors.capture_exception(exc)
-    return JSONResponse(
-        status_code=500,
-        content={"detail": "Something went wrong, please contact your administrator"},
-    )
-
-
-async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
-    return JSONResponse(
-        status_code=200,
-        content={"detail": "Client disconnected while an operation was in course"},
-    )
 
+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,
@@ -115,18 +79,4 @@ def create_application() -> FastAPI:
     # Use raw starlette routes to avoid unnecessary overhead
     application.add_route("/", homepage)
 
-    set_app_context(application)
-    maybe_configure_back_pressure(application)
     return application
-
-
-def maybe_configure_back_pressure(application: FastAPI):
-    from nucliadb.writer.back_pressure import start_materializer, stop_materializer
-    from nucliadb.writer.settings import back_pressure_settings
-    from nucliadb_utils.settings import is_onprem_nucliadb
-
-    if back_pressure_settings.enabled and not is_onprem_nucliadb():
-        context = get_app_context(application)
-        start_materializer_with_context = functools.partial(start_materializer, context)
-        application.add_event_handler("startup", start_materializer_with_context)
-        application.add_event_handler("shutdown", stop_materializer)

nucliadb/writer/back_pressure.py

@@ -28,7 +28,6 @@ from typing import Optional
 from async_lru import alru_cache
 from cachetools import TTLCache
 from fastapi import HTTPException, Request
-from nucliadb_protos.writer_pb2 import ShardObject
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.manager import get_index_nodes
@@ -37,11 +36,11 @@ from nucliadb.common.context.fastapi import get_app_context
 from nucliadb.common.http_clients.processing import ProcessingHTTPClient
 from nucliadb.writer import logger
 from nucliadb.writer.settings import back_pressure_settings as settings
+from nucliadb_protos.writer_pb2 import ShardObject
 from nucliadb_telemetry import metrics
 from nucliadb_utils import const
 from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.settings import is_onprem_nucliadb
-from nucliadb_utils.utilities import has_feature
 
 __all__ = ["maybe_back_pressure"]
 
@@ -49,7 +48,7 @@ __all__ = ["maybe_back_pressure"]
 back_pressure_observer = metrics.Observer("nucliadb_back_pressure", labels={"type": ""})
 
 
-rate_limited_requests_counter = metrics.Counter(
+RATE_LIMITED_REQUESTS_COUNTER = metrics.Counter(
     "nucliadb_rate_limited_requests", labels={"type": "", "cached": ""}
 )
 
@@ -113,8 +112,15 @@ def cached_back_pressure(kbid: str, resource_uuid: Optional[str] = None):
     if data is not None:
         try_after = data.try_after
         back_pressure_type = data.type
-        rate_limited_requests_counter.inc(
-            {"type": back_pressure_type, "cached": "true"}
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "true"})
+        logger.info(
+            "Back pressure applied from cache",
+            extra={
+                "type": back_pressure_type,
+                "try_after": try_after,
+                "kbid": kbid,
+                "resource_uuid": resource_uuid,
+            },
         )
         raise HTTPException(
             status_code=429,
@@ -129,9 +135,7 @@ def cached_back_pressure(kbid: str, resource_uuid: Optional[str] = None):
     except BackPressureException as exc:
         try_after = exc.data.try_after
         back_pressure_type = exc.data.type
-        rate_limited_requests_counter.inc(
-            {"type": back_pressure_type, "cached": "false"}
-        )
+        RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "false"})
         _cache.set(cache_key, exc.data)
         raise HTTPException(
             status_code=429,
@@ -216,6 +220,11 @@ class Materializer:
                 )
                 return 0
 
+            if pending > 0:
+                logger.info(
+                    f"Processing returned {pending} pending messages for KB",
+                    extra={"kbid": kbid},
+                )
             self.processing_pending_cache[kbid] = pending
             return pending
 
@@ -235,9 +244,7 @@ class Materializer:
                 for node in get_index_nodes():
                     try:
                         with back_pressure_observer({"type": "get_indexing_pending"}):
-                            self.indexing_pending[
-                                node.id
-                            ] = await get_nats_consumer_pending_messages(
+                            self.indexing_pending[node.id] = await get_nats_consumer_pending_messages(
                                 self.nats_manager,
                                 stream=const.Streams.INDEX.name,
                                 consumer=const.Streams.INDEX.group.format(node=node.id),
@@ -321,22 +328,17 @@ def get_materializer() -> Materializer:
     return MATERIALIZER
 
 
-async def maybe_back_pressure(
-    request: Request, kbid: str, resource_uuid: Optional[str] = None
-) -> None:
+async def maybe_back_pressure(request: Request, kbid: str, resource_uuid: Optional[str] = None) -> None:
     """
     This function does system checks to see if we need to put back pressure on writes.
     In that case, a HTTP 429 will be raised with the estimated time to try again.
     """
-    if has_feature(const.Features.BACK_PRESSURE, context={"kbid": kbid}):
-        if not is_back_pressure_enabled() or is_onprem_nucliadb():
-            return
-        await back_pressure_checks(request, kbid, resource_uuid)
+    if not is_back_pressure_enabled() or is_onprem_nucliadb():
+        return
+    await back_pressure_checks(request, kbid, resource_uuid)
 
 
-async def back_pressure_checks(
-    request: Request, kbid: str, resource_uuid: Optional[str] = None
-):
+async def back_pressure_checks(request: Request, kbid: str, resource_uuid: Optional[str] = None):
     """
     Will raise a 429 if back pressure is needed:
     - If the processing engine is behind.
@@ -347,9 +349,7 @@ async def back_pressure_checks(
     materializer = get_materializer()
     with cached_back_pressure(kbid, resource_uuid):
         check_ingest_behind(materializer.get_ingest_pending())
-        await check_indexing_behind(
-            context, kbid, resource_uuid, materializer.get_indexing_pending()
-        )
+        await check_indexing_behind(context, kbid, resource_uuid, materializer.get_indexing_pending())
         await check_processing_behind(materializer, kbid)
 
 
@@ -366,9 +366,19 @@ async def check_processing_behind(materializer: Materializer, kbid: str):
     kb_pending = await materializer.get_processing_pending(kbid)
     if kb_pending > max_pending:
         try_after = estimate_try_after(
-            rate=settings.processing_rate, pending=kb_pending
+            rate=settings.processing_rate,
+            pending=kb_pending,
+            max_wait=settings.max_wait_time,
         )
         data = BackPressureData(type="processing", try_after=try_after)
+        logger.info(
+            "Processing back pressure applied",
+            extra={
+                "kbid": kbid,
+                "try_after": try_after,
+                "pending": kb_pending,
+            },
+        )
         raise BackPressureException(data)
 
 
@@ -394,9 +404,7 @@ async def check_indexing_behind(
 
     # Get nodes that are involved in the indexing of the request
     if resource_uuid is not None:
-        nodes_to_check = await get_nodes_for_resource_shard(
-            context, kbid, resource_uuid
-        )
+        nodes_to_check = await get_nodes_for_resource_shard(context, kbid, resource_uuid)
     else:
         nodes_to_check = await get_nodes_for_kb_active_shards(context, kbid)
 
@@ -418,9 +426,20 @@ async def check_indexing_behind(
 
     if highest_pending > max_pending:
         try_after = estimate_try_after(
-            rate=settings.indexing_rate, pending=highest_pending
+            rate=settings.indexing_rate,
+            pending=highest_pending,
+            max_wait=settings.max_wait_time,
         )
         data = BackPressureData(type="indexing", try_after=try_after)
+        logger.info(
+            "Indexing back pressure applied",
+            extra={
+                "kbid": kbid,
+                "resource_uuid": resource_uuid,
+                "try_after": try_after,
+                "pending": highest_pending,
+            },
+        )
         raise BackPressureException(data)
 
 
@@ -432,24 +451,28 @@ def check_ingest_behind(ingest_pending: int):
 
     if ingest_pending > max_pending:
         try_after = estimate_try_after(
-            rate=settings.ingest_rate, pending=ingest_pending
+            rate=settings.ingest_rate,
+            pending=ingest_pending,
+            max_wait=settings.max_wait_time,
         )
         data = BackPressureData(type="ingest", try_after=try_after)
+        logger.info(
+            "Ingest back pressure applied",
+            extra={"try_after": try_after, "pending": ingest_pending},
+        )
         raise BackPressureException(data)
 
 
-def estimate_try_after(rate: float, pending: int) -> datetime:
+def estimate_try_after(rate: float, pending: int, max_wait: int) -> datetime:
     """
     This function estimates the time to try again based on the rate and the number of pending messages.
     """
-    delta_seconds = pending / rate
+    delta_seconds = min(pending / rate, max_wait)
     return datetime.utcnow() + timedelta(seconds=delta_seconds)
 
 
 @alru_cache(maxsize=1024, ttl=60 * 15)
-async def get_nodes_for_kb_active_shards(
-    context: ApplicationContext, kbid: str
-) -> list[str]:
+async def get_nodes_for_kb_active_shards(context: ApplicationContext, kbid: str) -> list[str]:
     with back_pressure_observer({"type": "get_kb_active_shard"}):
         active_shard = await get_kb_active_shard(context, kbid)
     if active_shard is None:
@@ -480,20 +503,16 @@ async def get_nats_consumer_pending_messages(
     return consumer_info.num_pending
 
 
-async def get_kb_active_shard(
-    context: ApplicationContext, kbid: str
-) -> Optional[ShardObject]:
-    async with context.kv_driver.transaction() as txn:
+async def get_kb_active_shard(context: ApplicationContext, kbid: str) -> Optional[ShardObject]:
+    async with context.kv_driver.transaction(read_only=True) as txn:
         return await context.shard_manager.get_current_active_shard(txn, kbid)
 
 
 async def get_resource_shard(
     context: ApplicationContext, kbid: str, resource_uuid: str
 ) -> Optional[ShardObject]:
-    async with datamanagers.with_transaction(read_only=True) as txn:
-        shard_id = await datamanagers.resources.get_resource_shard_id(
-            txn, kbid=kbid, rid=resource_uuid
-        )
+    async with datamanagers.with_ro_transaction() as txn:
+        shard_id = await datamanagers.resources.get_resource_shard_id(txn, kbid=kbid, rid=resource_uuid)
         if shard_id is None:
             # Resource does not exist
             logger.debug(

nucliadb/writer/exceptions.py

@@ -19,10 +19,6 @@
 #
 
 
-class ConflictError(Exception):
-    pass
-
-
 class ResourceNotFound(Exception):
     pass
 

nucliadb/writer/lifecycle.py

@@ -17,13 +17,20 @@
 # 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 nucliadb.ingest.processing import start_processing_engine
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from nucliadb.common.context.fastapi import inject_app_context
+from nucliadb.ingest.processing import start_processing_engine, stop_processing_engine
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.writer import SERVICE_NAME
+from nucliadb.writer.back_pressure import start_materializer, stop_materializer
+from nucliadb.writer.settings import back_pressure_settings
 from nucliadb.writer.tus import finalize as storage_finalize
 from nucliadb.writer.tus import initialize as storage_initialize
-from nucliadb.writer.utilities import get_processing
 from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
+from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.utilities import (
     finalize_utilities,
     start_partitioning_utility,
@@ -32,29 +39,28 @@ from nucliadb_utils.utilities import (
 )
 
 
-async def initialize():
-    await setup_telemetry(SERVICE_NAME)
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    back_pressure_enabled = back_pressure_settings.enabled and not is_onprem_nucliadb()
 
+    await setup_telemetry(SERVICE_NAME)
     await start_ingest(SERVICE_NAME)
-
     await start_processing_engine()
-
     start_partitioning_utility()
-
     await start_transaction_utility(SERVICE_NAME)
     await storage_initialize()
 
+    # Inject application context into the fastapi app's state
+    async with inject_app_context(app) as context:
+        if back_pressure_enabled:
+            await start_materializer(context)
+        yield
 
-async def finalize():
+    if back_pressure_enabled:
+        await stop_materializer()
     await stop_transaction_utility()
-
     await stop_ingest()
-    processing = get_processing()
-    if processing is not None:
-        await processing.finalize()
-
+    await stop_processing_engine()
     await storage_finalize()
-
     await clean_telemetry(SERVICE_NAME)
-
     await finalize_utilities()

nucliadb/writer/resource/audit.py

@@ -19,9 +19,10 @@
 #
 from datetime import datetime
 
-from nucliadb_protos.writer_pb2 import Audit
 from starlette.requests import Request
 
+from nucliadb_protos.writer_pb2 import Audit
+
 
 def parse_audit(audit: Audit, request: Request):
     audit.user = request.headers.get("X-NUCLIADB-USER", "")

nucliadb/writer/resource/basic.py

@@ -18,8 +18,31 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from datetime import datetime
+from typing import Optional
 
 from fastapi import HTTPException
+
+from nucliadb.common.models_utils import to_proto
+from nucliadb.common.models_utils.from_proto import (
+    RelationNodeTypeMap,
+    RelationTypeMap,
+)
+from nucliadb.ingest.orm.utils import set_title
+from nucliadb.ingest.processing import PushPayload
+from nucliadb_models.content_types import GENERIC_MIME_TYPE
+from nucliadb_models.file import FileField
+from nucliadb_models.link import LinkField
+from nucliadb_models.metadata import (
+    ParagraphAnnotation,
+    QuestionAnswerAnnotation,
+)
+from nucliadb_models.text import TEXT_FORMAT_TO_MIMETYPE, PushTextFormat, Text
+from nucliadb_models.writer import (
+    ComingResourcePayload,
+    CreateResourcePayload,
+    UpdateResourcePayload,
+)
+from nucliadb_protos.knowledgebox_pb2 import KnowledgeBoxConfig
 from nucliadb_protos.resources_pb2 import (
     Answers,
     Basic,
@@ -30,38 +53,19 @@ from nucliadb_protos.resources_pb2 import (
     Metadata,
     PageSelections,
     Paragraph,
+    TokenSplit,
+    UserFieldMetadata,
+    VisualSelection,
 )
 from nucliadb_protos.resources_pb2 import ParagraphAnnotation as PBParagraphAnnotation
 from nucliadb_protos.resources_pb2 import (
     QuestionAnswerAnnotation as PBQuestionAnswerAnnotation,
 )
-from nucliadb_protos.resources_pb2 import TokenSplit, UserFieldMetadata, VisualSelection
 from nucliadb_protos.utils_pb2 import Relation, RelationNode
 from nucliadb_protos.writer_pb2 import BrokerMessage
 
-from nucliadb.ingest.orm.utils import set_title
-from nucliadb.ingest.processing import ProcessingInfo, PushPayload
-from nucliadb_models.common import FIELD_TYPES_MAP_REVERSE
-from nucliadb_models.file import FileField
-from nucliadb_models.link import LinkField
-from nucliadb_models.metadata import (
-    ParagraphAnnotation,
-    QuestionAnswerAnnotation,
-    RelationNodeTypeMap,
-    RelationTypeMap,
-)
-from nucliadb_models.text import TEXT_FORMAT_TO_MIMETYPE, PushTextFormat, Text
-from nucliadb_models.writer import (
-    GENERIC_MIME_TYPE,
-    ComingResourcePayload,
-    CreateResourcePayload,
-    UpdateResourcePayload,
-)
-
 
-def parse_basic_modify(
-    bm: BrokerMessage, item: ComingResourcePayload, toprocess: PushPayload
-):
+def parse_basic_modify(bm: BrokerMessage, item: ComingResourcePayload, toprocess: PushPayload):
     bm.basic.modified.FromDatetime(datetime.now())
     if item.title:
         set_title(bm, toprocess, item.title)
@@ -83,14 +87,9 @@ def parse_basic_modify(
         bm.basic.metadata.useful = True
         bm.basic.metadata.status = Metadata.Status.PENDING
 
-        toprocess.genericfield["summary"] = Text(
-            body=item.summary, format=PushTextFormat.PLAIN
-        )
+        toprocess.genericfield["summary"] = Text(body=item.summary, format=PushTextFormat.PLAIN)
     if item.thumbnail:
         bm.basic.thumbnail = item.thumbnail
-    if item.layout:
-        bm.basic.layout = item.layout
-
     if item.metadata is not None:
         bm.basic.metadata.metadata.update(item.metadata.metadata)
         if item.metadata.language:
@@ -147,9 +146,8 @@ def parse_basic_modify(
                 userfieldmetadata.question_answers.append(qa_annotation_pb)
 
             userfieldmetadata.field.field = fieldmetadata.field.field
-            userfieldmetadata.field.field_type = FIELD_TYPES_MAP_REVERSE[  # type: ignore
-                fieldmetadata.field.field_type.value
-            ]
+
+            userfieldmetadata.field.field_type = to_proto.field_type(fieldmetadata.field.field_type)
 
             bm.basic.fieldmetadata.append(userfieldmetadata)
 
@@ -167,9 +165,7 @@ def parse_basic_modify(
             ]
         )
 
-        relation_node_resource = RelationNode(
-            value=bm.uuid, ntype=RelationNode.NodeType.RESOURCE
-        )
+        relation_node_resource = RelationNode(value=bm.uuid, ntype=RelationNode.NodeType.RESOURCE)
         relations = []
         for relation in item.usermetadata.relations:
             if relation.from_ is None:
@@ -205,8 +201,16 @@ def parse_basic_modify(
         unique_groups = list(set(item.security.access_groups))
         bm.security.access_groups.extend(unique_groups)
 
+    if item.hidden is not None:
+        bm.basic.hidden = item.hidden
 
-def parse_basic(bm: BrokerMessage, item: CreateResourcePayload, toprocess: PushPayload):
+
+def parse_basic_creation(
+    bm: BrokerMessage,
+    item: CreateResourcePayload,
+    toprocess: PushPayload,
+    kb_config: Optional[KnowledgeBoxConfig],
+):
     bm.basic.created.FromDatetime(datetime.now())
 
     if item.title is None:
@@ -215,6 +219,10 @@ def parse_basic(bm: BrokerMessage, item: CreateResourcePayload, toprocess: PushP
 
     parse_basic_modify(bm, item, toprocess)
 
+    if item.hidden is None:
+        if kb_config and kb_config.hidden_resources_hide_on_creation:
+            bm.basic.hidden = True
+
 
 def set_status(basic: Basic, item: CreateResourcePayload):
     basic.metadata.status = Metadata.Status.PENDING
@@ -224,34 +232,14 @@ def set_status_modify(basic: Basic, item: UpdateResourcePayload):
     basic.metadata.status = Metadata.Status.PENDING
 
 
-def set_processing_info(bm: BrokerMessage, processing_info: ProcessingInfo):
-    """
-    Processing V2 does not have this awkward processing info data field and storage
-    but keeping for b/w compatibility.
-
-    Once V1 is removed, this code can be removed because status checking will be done
-    in a separate API that is not part of NucliaDB.
-    """
-    if processing_info.seqid is not None:
-        bm.basic.last_seqid = processing_info.seqid
-    if processing_info.account_seq is not None:
-        bm.basic.last_account_seq = processing_info.account_seq
-    if processing_info.queue is not None:
-        bm.basic.queue = bm.basic.QueueType.Value(processing_info.queue.name)
-
-
 def validate_classifications(paragraph: ParagraphAnnotation):
     classifications = paragraph.classifications
     if len(classifications) == 0:
-        raise HTTPException(
-            status_code=422, detail="ensure classifications has at least 1 items"
-        )
+        raise HTTPException(status_code=422, detail="ensure classifications has at least 1 items")
 
-    unique_classifications = {tuple(cf.dict().values()) for cf in classifications}
+    unique_classifications = {tuple(cf.model_dump().values()) for cf in classifications}
     if len(unique_classifications) != len(classifications):
-        raise HTTPException(
-            status_code=422, detail="Paragraph classifications need to be unique"
-        )
+        raise HTTPException(status_code=422, detail="Paragraph classifications need to be unique")
 
 
 def compute_title(item: CreateResourcePayload, rid: str) -> str:
@@ -289,9 +277,7 @@ def build_question_answer_annotation_pb(
     pb.cancelled_by_user = qa_annotation.cancelled_by_user
     pb.question_answer.question.text = qa_annotation.question_answer.question.text
     if qa_annotation.question_answer.question.language is not None:
-        pb.question_answer.question.language = (
-            qa_annotation.question_answer.question.language
-        )
+        pb.question_answer.question.language = qa_annotation.question_answer.question.language
     pb.question_answer.question.ids_paragraphs.extend(
         qa_annotation.question_answer.question.ids_paragraphs
     )