nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/backups/tasks.py

@@ -17,7 +17,7 @@
 # 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 typing import Awaitable, Callable
+from collections.abc import Awaitable, Callable
 
 from nucliadb.backups.const import BackupsNatsConfig
 from nucliadb.backups.create import backup_kb_task
@@ -38,6 +38,7 @@ def creator_consumer() -> NatsTaskConsumer[CreateBackupRequest]:
         callback=backup_kb_task,
         msg_type=CreateBackupRequest,
         max_concurrent_messages=10,
+        max_retries=100,
     )
     return consumer
 
@@ -64,6 +65,7 @@ def restorer_consumer() -> NatsTaskConsumer[RestoreBackupRequest]:
         callback=restore_kb_task,
         msg_type=RestoreBackupRequest,
         max_concurrent_messages=10,
+        max_retries=100,
     )
     return consumer
 
@@ -90,6 +92,7 @@ def deleter_consumer() -> NatsTaskConsumer[DeleteBackupRequest]:
         callback=delete_backup_task,
         msg_type=DeleteBackupRequest,
         max_concurrent_messages=2,
+        max_retries=100,
     )
     return consumer
 

nucliadb/common/back_pressure/cache.py

@@ -21,7 +21,6 @@ import contextlib
 import logging
 import threading
 from datetime import datetime, timezone
-from typing import Optional
 
 from cachetools import TTLCache
 
@@ -47,7 +46,7 @@ class BackPressureCache:
         self._cache = TTLCache(maxsize=1024, ttl=5 * 60)
         self._lock = threading.Lock()
 
-    def get(self, key: str) -> Optional[BackPressureData]:
+    def get(self, key: str) -> BackPressureData | None:
         with self._lock:
             data = self._cache.get(key, None)
             if data is None:
@@ -72,7 +71,7 @@ def cached_back_pressure(cache_key: str):
     Context manager that handles the caching of the try again in time so that
     we don't recompute try again times if we have already applied back pressure.
     """
-    data: Optional[BackPressureData] = _cache.get(cache_key)
+    data: BackPressureData | None = _cache.get(cache_key)
     if data is not None:
         back_pressure_type = data.type
         RATE_LIMITED_REQUESTS_COUNTER.inc({"type": back_pressure_type, "cached": "true"})

nucliadb/common/back_pressure/materializer.py

@@ -20,7 +20,6 @@
 import asyncio
 import logging
 import threading
-from typing import Optional
 
 from cachetools import TTLCache
 from fastapi import HTTPException
@@ -118,12 +117,6 @@ class BackPressureMaterializer:
                     extra={"kbid": kbid},
                 )
                 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
 
@@ -184,7 +177,7 @@ class BackPressureMaterializer:
                 pending=pending,
                 max_wait=settings.max_wait_time,
             )
-            data = BackPressureData(type="indexing", try_after=try_after)
+            data = BackPressureData(type="indexing", try_after=try_after, pending=pending)
             raise BackPressureException(data)
 
     def check_ingest(self):
@@ -199,7 +192,7 @@ class BackPressureMaterializer:
                 pending=ingest_pending,
                 max_wait=settings.max_wait_time,
             )
-            data = BackPressureData(type="ingest", try_after=try_after)
+            data = BackPressureData(type="ingest", try_after=try_after, pending=ingest_pending)
             raise BackPressureException(data)
 
     async def check_processing(self, kbid: str):
@@ -215,11 +208,11 @@ class BackPressureMaterializer:
                 pending=kb_pending,
                 max_wait=settings.max_wait_time,
             )
-            data = BackPressureData(type="processing", try_after=try_after)
+            data = BackPressureData(type="processing", try_after=try_after, pending=kb_pending)
             raise BackPressureException(data)
 
 
-MATERIALIZER: Optional[BackPressureMaterializer] = None
+MATERIALIZER: BackPressureMaterializer | None = None
 materializer_lock = threading.Lock()
 
 
@@ -268,7 +261,7 @@ def get_materializer() -> BackPressureMaterializer:
     return MATERIALIZER
 
 
-async def maybe_back_pressure(kbid: str, resource_uuid: Optional[str] = None) -> None:
+async def maybe_back_pressure(kbid: str, resource_uuid: str | None = 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.
@@ -278,7 +271,7 @@ async def maybe_back_pressure(kbid: str, resource_uuid: Optional[str] = None) ->
     await back_pressure_checks(kbid, resource_uuid)
 
 
-async def back_pressure_checks(kbid: str, resource_uuid: Optional[str] = None):
+async def back_pressure_checks(kbid: str, resource_uuid: str | None = None):
     """
     Will raise a 429 if back pressure is needed:
     - If the processing engine is behind.
@@ -299,6 +292,7 @@ async def back_pressure_checks(kbid: str, resource_uuid: Optional[str] = None):
                 "resource_uuid": resource_uuid,
                 "try_after": exc.data.try_after,
                 "back_pressure_type": exc.data.type,
+                "pending": exc.data.pending,
             },
         )
         raise HTTPException(

nucliadb/common/back_pressure/settings.py

@@ -29,30 +29,30 @@ class BackPressureSettings(BaseSettings):
     )
     indexing_rate: float = Field(
         default=10,
-        description="Estimation of the indexing rate in messages per second. This is used to calculate the try again in time",  # noqa
+        description="Estimation of the indexing rate in messages per second. This is used to calculate the try again in time",
     )
     ingest_rate: float = Field(
         default=4,
-        description="Estimation of the ingest processed consumer rate in messages per second. This is used to calculate the try again in time",  # noqa
+        description="Estimation of the ingest processed consumer rate in messages per second. This is used to calculate the try again in time",
     )
     processing_rate: float = Field(
         default=1,
-        description="Estimation of the processing rate in messages per second. This is used to calculate the try again in time",  # noqa
+        description="Estimation of the processing rate in messages per second. This is used to calculate the try again in time",
     )
     max_indexing_pending: int = Field(
         default=1000,
-        description="Max number of messages pending to index in a node queue before rate limiting writes. Set to 0 to disable indexing back pressure checks",  # noqa
+        description="Max number of messages pending to index in a node queue before rate limiting writes. Set to 0 to disable indexing back pressure checks",
         alias="back_pressure_max_indexing_pending",
     )
     max_ingest_pending: int = Field(
         # Disabled by default
         default=0,
-        description="Max number of messages pending to be ingested by processed consumers before rate limiting writes. Set to 0 to disable ingest back pressure checks",  # noqa
+        description="Max number of messages pending to be ingested by processed consumers before rate limiting writes. Set to 0 to disable ingest back pressure checks",
         alias="back_pressure_max_ingest_pending",
     )
     max_processing_pending: int = Field(
         default=1000,
-        description="Max number of messages pending to process per Knowledge Box before rate limiting writes. Set to 0 to disable processing back pressure checks",  # noqa
+        description="Max number of messages pending to process per Knowledge Box before rate limiting writes. Set to 0 to disable processing back pressure checks",
         alias="back_pressure_max_processing_pending",
     )
     indexing_check_interval: int = Field(

nucliadb/common/back_pressure/utils.py

@@ -28,6 +28,7 @@ from nucliadb_utils.nats import NatsConnectionManager
 class BackPressureData:
     type: str
     try_after: datetime
+    pending: int = 0
 
 
 class BackPressureException(Exception):

nucliadb/common/cache.py

@@ -24,7 +24,7 @@ from abc import ABC, abstractmethod
 from contextvars import ContextVar
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Generic, Optional, TypeVar
+from typing import Generic, TypeVar
 
 import backoff
 from async_lru import _LRUCacheWrapper, alru_cache
@@ -66,9 +66,9 @@ class Cache(Generic[K, T], ABC):
 
     """
 
-    cache: _LRUCacheWrapper[Optional[T]]
+    cache: _LRUCacheWrapper[T | None]
 
-    async def get(self, *args: K.args, **kwargs: K.kwargs) -> Optional[T]:
+    async def get(self, *args: K.args, **kwargs: K.kwargs) -> T | None:
         result = await self.cache(*args)
         # Do not cache None
         if result is None:
@@ -88,7 +88,7 @@ class Cache(Generic[K, T], ABC):
 class ResourceCache(Cache[[str, str], ResourceORM]):
     def __init__(self, cache_size: int) -> None:
         @alru_cache(maxsize=cache_size)
-        async def _get_resource(kbid: str, rid: str) -> Optional[ResourceORM]:
+        async def _get_resource(kbid: str, rid: str) -> ResourceORM | None:
             storage = await get_storage()
             async with get_driver().ro_transaction() as txn:
                 kb = KnowledgeBoxORM(txn, storage, kbid)
@@ -115,7 +115,7 @@ class ExtractedTextCache(Cache[[str, FieldId], ExtractedText]):
     def __init__(self, cache_size: int) -> None:
         @alru_cache(maxsize=cache_size)
         @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
-        async def _get_extracted_text(kbid: str, field_id: FieldId) -> Optional[ExtractedText]:
+        async def _get_extracted_text(kbid: str, field_id: FieldId) -> ExtractedText | None:
             storage = await get_storage()
             try:
                 sf = storage.file_extracted(
@@ -144,18 +144,18 @@ class ExtractedTextCache(Cache[[str, FieldId], ExtractedText]):
 
 # Global caches (per asyncio task)
 
-rcache: ContextVar[Optional[ResourceCache]] = ContextVar("rcache", default=None)
-etcache: ContextVar[Optional[ExtractedTextCache]] = ContextVar("etcache", default=None)
+rcache: ContextVar[ResourceCache | None] = ContextVar("rcache", default=None)
+etcache: ContextVar[ExtractedTextCache | None] = ContextVar("etcache", default=None)
 
 
 # Cache management
 
 
-def get_resource_cache() -> Optional[ResourceCache]:
+def get_resource_cache() -> ResourceCache | None:
     return rcache.get()
 
 
-def get_extracted_text_cache() -> Optional[ExtractedTextCache]:
+def get_extracted_text_cache() -> ExtractedTextCache | None:
     return etcache.get()
 
 

nucliadb/common/catalog/__init__.py

@@ -0,0 +1,79 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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/>.
+#
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 nidx_protos.noderesources_pb2 import Resource as IndexMessage
+
+from nucliadb.common.catalog.dummy import DummyCatalog
+from nucliadb.common.catalog.interface import Catalog, CatalogQuery
+from nucliadb.common.catalog.pg import PGCatalog
+from nucliadb.common.catalog.utils import build_catalog_resource_data
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.ingest.settings import CatalogConfig, settings
+from nucliadb_models.search import CatalogFacetsRequest, Resources
+from nucliadb_utils.exceptions import ConfigurationError
+
+
+def get_catalog() -> Catalog:
+    if settings.catalog == CatalogConfig.UNSET:
+        return DummyCatalog()
+    elif settings.catalog == CatalogConfig.PG:
+        return PGCatalog()
+    else:
+        raise ConfigurationError(f"Unknown catalog configuration: {settings.catalog}")
+
+
+async def catalog_update(txn: Transaction, kbid: str, resource: Resource, index_message: IndexMessage):
+    catalog = get_catalog()
+    resource_data = build_catalog_resource_data(resource, index_message)
+    await catalog.update(txn, kbid, resource.uuid, resource_data)
+
+
+async def catalog_delete(txn: Transaction, kbid: str, rid: str):
+    catalog = get_catalog()
+    await catalog.delete(txn, kbid, rid)
+
+
+async def catalog_search(query: CatalogQuery) -> Resources:
+    catalog = get_catalog()
+    return await catalog.search(query)
+
+
+async def catalog_facets(kbid: str, request: CatalogFacetsRequest) -> dict[str, int]:
+    catalog = get_catalog()
+    return await catalog.facets(kbid, request)

nucliadb/common/catalog/dummy.py

@@ -0,0 +1,36 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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.common.catalog.interface import Catalog, CatalogQuery, CatalogResourceData
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb_models.search import CatalogFacetsRequest, Resources
+
+
+class DummyCatalog(Catalog):
+    async def update(self, txn: Transaction, kbid: str, rid: str, data: CatalogResourceData):
+        return
+
+    async def delete(self, txn: Transaction, kbid: str, rid: str):
+        return
+
+    async def search(self, query: CatalogQuery) -> Resources:
+        return Resources(results=[], min_score=0.0)
+
+    async def facets(self, kbid: str, request: CatalogFacetsRequest) -> dict[str, int]:
+        return {}

nucliadb/common/catalog/interface.py

@@ -0,0 +1,85 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# 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 __future__ import annotations
+
+import abc
+import datetime
+from dataclasses import dataclass
+from typing import Literal
+
+from pydantic import BaseModel, Field
+
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb_models import search as search_models
+from nucliadb_models.search import CatalogFacetsRequest, Resources
+
+
+class CatalogResourceData(BaseModel):
+    """
+    Data extracted from a resource to be indexed in the catalog
+    """
+
+    title: str = Field(description="Resource title")
+    created_at: datetime.datetime = Field(description="Resource creation date")
+    modified_at: datetime.datetime = Field(description="Resource last modification date")
+    labels: list[str] = Field(
+        description="Resource labels. This includes labels at the resource level and all classification labels of its fields"
+    )
+    slug: str = Field(description="Resource slug")
+
+
+@dataclass
+class CatalogExpression:
+    @dataclass
+    class Date:
+        field: Literal["created_at"] | Literal["modified_at"]
+        since: datetime.datetime | None
+        until: datetime.datetime | None
+
+    bool_and: list[CatalogExpression] | None = None
+    bool_or: list[CatalogExpression] | None = None
+    bool_not: CatalogExpression | None = None
+    date: Date | None = None
+    facet: str | None = None
+    resource_id: str | None = None
+
+
+class CatalogQuery(BaseModel):
+    kbid: str
+    query: search_models.CatalogQuery | None = Field(description="Full-text search query")
+    filters: CatalogExpression | None = Field(description="Filters to apply to the search")
+    sort: search_models.SortOptions = Field(description="Sorting option")
+    faceted: list[str] = Field(description="List of facets to compute during the search")
+    page_size: int = Field(description="Used for pagination. Maximum page size is 100")
+    page_number: int = Field(description="Used for pagination. First page is 0")
+
+
+class Catalog(abc.ABC, metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def update(self, txn: Transaction, kbid: str, rid: str, data: CatalogResourceData): ...
+
+    @abc.abstractmethod
+    async def delete(self, txn: Transaction, kbid: str, rid: str): ...
+
+    @abc.abstractmethod
+    async def search(self, query: CatalogQuery) -> Resources: ...
+
+    @abc.abstractmethod
+    async def facets(self, kbid: str, request: CatalogFacetsRequest) -> dict[str, int]: ...