nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/ingest/service/writer.py

@@ -18,12 +18,11 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import uuid
-from typing import AsyncIterator
+from collections.abc import AsyncIterator
 
 from nucliadb.backups import tasks as backup_tasks
 from nucliadb.backups import utils as backup_utils
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.exceptions import AlreadyExists, EntitiesGroupNotFound
 from nucliadb.common.cluster.utils import get_shard_manager
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.external_index_providers.exceptions import ExternalIndexCreationError
@@ -49,7 +48,6 @@ from nucliadb_protos.knowledgebox_pb2 import (
 )
 from nucliadb_protos.writer_pb2 import (
     BrokerMessage,
-    DelEntitiesRequest,
     GetEntitiesGroupRequest,
     GetEntitiesGroupResponse,
     GetEntitiesRequest,
@@ -58,12 +56,7 @@ from nucliadb_protos.writer_pb2 import (
     IndexStatus,
     ListEntitiesGroupsRequest,
     ListEntitiesGroupsResponse,
-    NewEntitiesGroupRequest,
-    NewEntitiesGroupResponse,
     OpStatusWriter,
-    SetEntitiesRequest,
-    UpdateEntitiesGroupRequest,
-    UpdateEntitiesGroupResponse,
     WriterStatusRequest,
     WriterStatusResponse,
 )
@@ -224,29 +217,6 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
             logger.info(f"Processed {message.uuid}")
         return response
 
-    async def NewEntitiesGroup(  # type: ignore
-        self, request: NewEntitiesGroupRequest, context=None
-    ) -> NewEntitiesGroupResponse:
-        response = NewEntitiesGroupResponse()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = NewEntitiesGroupResponse.Status.KB_NOT_FOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.create_entities_group(request.group, request.entities)
-            except AlreadyExists:
-                response.status = NewEntitiesGroupResponse.Status.ALREADY_EXISTS
-                return response
-
-            await txn.commit()
-            response.status = NewEntitiesGroupResponse.Status.OK
-            return response
-
     async def GetEntities(  # type: ignore
         self, request: GetEntitiesRequest, context=None
     ) -> GetEntitiesResponse:
@@ -320,81 +290,6 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
             return response
 
-    async def SetEntities(self, request: SetEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
-        response = OpStatusWriter()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = OpStatusWriter.Status.NOTFOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.set_entities_group(request.group, request.entities)
-            except Exception as e:
-                errors.capture_exception(e)
-                logger.error("Error in ingest gRPC servicer", exc_info=True)
-                response.status = OpStatusWriter.Status.ERROR
-            else:
-                response.status = OpStatusWriter.Status.OK
-                await txn.commit()
-            return response
-
-    async def UpdateEntitiesGroup(  # type: ignore
-        self, request: UpdateEntitiesGroupRequest, context=None
-    ) -> UpdateEntitiesGroupResponse:
-        response = UpdateEntitiesGroupResponse()
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = UpdateEntitiesGroupResponse.Status.KB_NOT_FOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.set_entities_group_metadata(
-                    request.group,
-                    title=request.title,
-                    color=request.color,
-                )
-                updates = {**request.add, **request.update}
-                await entities_manager.update_entities(request.group, updates)
-                await entities_manager.delete_entities(request.group, request.delete)  # type: ignore
-            except EntitiesGroupNotFound:
-                response.status = UpdateEntitiesGroupResponse.Status.ENTITIES_GROUP_NOT_FOUND
-                return response
-
-            await txn.commit()
-            response.status = UpdateEntitiesGroupResponse.Status.OK
-            return response
-
-    async def DelEntities(self, request: DelEntitiesRequest, context=None) -> OpStatusWriter:  # type: ignore
-        response = OpStatusWriter()
-
-        async with self.driver.ro_transaction() as ro_txn:
-            kbobj = await self.proc.get_kb_obj(ro_txn, request.kb)
-            if kbobj is None:
-                response.status = OpStatusWriter.Status.NOTFOUND
-                return response
-
-        async with self.driver.rw_transaction() as txn:
-            kbobj.txn = txn
-            entities_manager = EntitiesManager(kbobj, txn)
-            try:
-                await entities_manager.delete_entities_group(request.group)
-            except Exception as e:
-                errors.capture_exception(e)
-                logger.error("Error in ingest gRPC servicer", exc_info=True)
-                response.status = OpStatusWriter.Status.ERROR
-            else:
-                await txn.commit()
-                response.status = OpStatusWriter.Status.OK
-            return response
-
     async def Status(  # type: ignore
         self, request: WriterStatusRequest, context=None
     ) -> WriterStatusResponse:
@@ -413,8 +308,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
 
     async def Index(self, request: IndexResource, context=None) -> IndexStatus:  # type: ignore
         async with self.driver.ro_transaction() as txn:
-            kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-            resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
+            resobj = ResourceORM(txn, self.storage, request.kbid, request.rid)
             bm = await generate_broker_message(resobj)
             transaction = get_transaction_utility()
             partitioning = get_partitioning()
@@ -428,7 +322,7 @@ class WriterServicer(writer_pb2_grpc.WriterServicer):
         try:
             async with self.driver.rw_transaction() as txn:
                 kbobj = KnowledgeBoxORM(txn, self.storage, request.kbid)
-                resobj = ResourceORM(txn, self.storage, kbobj, request.rid)
+                resobj = ResourceORM(txn, self.storage, request.kbid, request.rid)
                 resobj.disable_vectors = not request.reindex_vectors
                 index_message = await get_resource_index_message(resobj, reindex=True)
                 shard = await self.proc.get_or_assign_resource_shard(txn, kbobj, request.rid)

nucliadb/ingest/settings.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from enum import Enum
-from typing import Optional
 
 from pydantic import Field
 from pydantic_settings import BaseSettings
@@ -41,13 +40,13 @@ class DriverConfig(Enum):
 
 class DriverSettings(BaseSettings):
     driver: DriverConfig = Field(default=DriverConfig.PG, description="K/V storage driver")
-    driver_local_url: Optional[str] = Field(
+    driver_local_url: str | None = Field(
         default=None,
         description="Local path to store data on file system. Example: /nucliadb/data/main",
     )
-    driver_pg_url: Optional[str] = Field(
+    driver_pg_url: str | None = Field(
         default=None,
-        description="PostgreSQL DSN. The connection string to the PG server. Example: postgres://username:password@postgres:5432/nucliadb.",  # noqa
+        description="PostgreSQL DSN. The connection string to the PG server. Example: postgres://username:password@postgres:5432/nucliadb.",
     )
     driver_pg_connection_pool_min_size: int = Field(
         default=10,

nucliadb/ingest/utils.py

@@ -17,7 +17,6 @@
 # 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 Optional
 
 from nucliadb.common.maindb.utils import setup_driver
 from nucliadb_protos.writer_pb2_grpc import WriterStub
@@ -26,7 +25,7 @@ from nucliadb_utils.settings import nucliadb_settings
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 
-async def start_ingest(service_name: Optional[str] = None):
+async def start_ingest(service_name: str | None = None):
     await setup_driver()
 
     actual_service = get_utility(Utility.INGEST)

nucliadb/learning_proxy.py

@@ -24,7 +24,7 @@ import os
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
 from enum import Enum, IntEnum
-from typing import Any, Optional, Union
+from typing import Any
 
 import backoff
 import httpx
@@ -89,12 +89,12 @@ class LearningConfiguration(BaseModel):
     # aka similarity function
     semantic_vector_similarity: str
     # aka vector_dimension
-    semantic_vector_size: Optional[int] = None
+    semantic_vector_size: int | None = None
     # aka min_score
-    semantic_threshold: Optional[float] = None
+    semantic_threshold: float | None = None
     # List of possible subdivisions of the matryoshka embeddings (if the model
     # supports it)
-    semantic_matryoshka_dimensions: Optional[list[int]] = Field(
+    semantic_matryoshka_dimensions: list[int] | None = Field(
         default=None, alias="semantic_matryoshka_dims"
     )
 
@@ -154,7 +154,7 @@ class LearningConfiguration(BaseModel):
 
 
 class ProxiedLearningConfigError(Exception):
-    def __init__(self, status_code: int, content: Union[str, dict[str, Any]]):
+    def __init__(self, status_code: int, content: str | dict[str, Any]):
         self.status_code = status_code
         self.content = content
 
@@ -176,7 +176,7 @@ def raise_for_status(response: httpx.Response) -> None:
 
 async def get_configuration(
     kbid: str,
-) -> Optional[LearningConfiguration]:
+) -> LearningConfiguration | None:
     return await learning_config_service().get_configuration(kbid)
 
 
@@ -205,7 +205,7 @@ async def learning_config_proxy(
     method: str,
     url: str,
     headers: dict[str, str] = {},
-) -> Union[Response, StreamingResponse]:
+) -> Response | StreamingResponse:
     return await proxy(
         service=LearningService.CONFIG,
         request=request,
@@ -245,7 +245,7 @@ async def proxy(
     method: str,
     url: str,
     headers: dict[str, str] = {},
-) -> Union[Response, StreamingResponse]:
+) -> Response | StreamingResponse:
     """
     Proxy the request to a learning API.
 
@@ -426,7 +426,7 @@ class DummyClient(httpx.AsyncClient):
 
 class LearningConfigService(ABC):
     @abstractmethod
-    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]: ...
+    async def get_configuration(self, kbid: str) -> LearningConfiguration | None: ...
 
     @abstractmethod
     async def set_configuration(self, kbid: str, config: dict[str, Any]) -> LearningConfiguration: ...
@@ -439,7 +439,7 @@ class LearningConfigService(ABC):
 
 
 class ProxiedLearningConfig(LearningConfigService):
-    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]:
+    async def get_configuration(self, kbid: str) -> LearningConfiguration | None:
         async with self._client() as client:
             resp = await client.get(f"config/{kbid}")
             try:
@@ -483,7 +483,7 @@ class InMemoryLearningConfig(LearningConfigService):
     def __init__(self):
         self.in_memory_configs = {}
 
-    async def get_configuration(self, kbid: str) -> Optional[LearningConfiguration]:
+    async def get_configuration(self, kbid: str) -> LearningConfiguration | None:
         return _IN_MEMORY_CONFIGS.get(kbid, None)
 
     async def set_configuration(self, kbid: str, config: dict[str, Any]) -> LearningConfiguration:

nucliadb/metrics_exporter.py

@@ -20,7 +20,8 @@
 from __future__ import annotations
 
 import asyncio
-from typing import AsyncGenerator, Callable, Tuple, cast
+from collections.abc import AsyncGenerator, Callable
+from typing import cast
 
 from nucliadb import logger
 from nucliadb.common import datamanagers
@@ -41,7 +42,7 @@ KBS_COUNT = metrics.Gauge("nucliadb_kbs_count")
 RESOURCES_COUNT = metrics.Gauge("nucliadb_resources_count")
 
 
-async def iter_kbids(context: ApplicationContext) -> AsyncGenerator[str, None]:
+async def iter_kbids(context: ApplicationContext) -> AsyncGenerator[str]:
     """
     Return a list of all KB ids.
     """
@@ -94,11 +95,11 @@ async def update_resource_metrics(context: ApplicationContext):
             "WHERE labels @> '{/n/s/PENDING}' "
             "AND COALESCE(modified_at, created_at) BETWEEN NOW() - INTERVAL '1 month' AND NOW() - INTERVAL '6 hours'"
         )
-        count = cast(Tuple[int], await cur.fetchone())[0]
+        count = cast(tuple[int], await cur.fetchone())[0]
         PENDING_RESOURCE_COUNT.set(count)
 
         await cur.execute("SELECT COUNT(*) FROM catalog")
-        count = cast(Tuple[int], await cur.fetchone())[0]
+        count = cast(tuple[int], await cur.fetchone())[0]
         RESOURCES_COUNT.set(count)
 
 

nucliadb/middleware/__init__.py

@@ -17,7 +17,10 @@
 # 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 logging
 import time
+from collections import deque
+from typing import ClassVar
 
 from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
 from starlette.requests import Request
@@ -27,6 +30,9 @@ PROCESS_TIME_HEADER = "X-PROCESS-TIME"
 ACCESS_CONTROL_EXPOSE_HEADER = "Access-Control-Expose-Headers"
 
 
+logger = logging.getLogger("nucliadb.middleware")
+
+
 class ProcessTimeHeaderMiddleware(BaseHTTPMiddleware):
     def capture_process_time(self, response, duration: float):
         response.headers[PROCESS_TIME_HEADER] = str(duration)
@@ -44,9 +50,84 @@ class ProcessTimeHeaderMiddleware(BaseHTTPMiddleware):
         start = time.perf_counter()
         try:
             response = await call_next(request)
+            return response
         finally:
             if response is not None:
                 duration = time.perf_counter() - start
                 self.capture_process_time(response, duration)
                 self.expose_process_time_header(response)
-                return response
+
+
+class ClientErrorPayloadLoggerMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware that logs the payload of client error responses (HTTP 412 and 422).
+    This helps supporting clients by providing more context about the errors they
+    encounter which otherwise we don't have much visibility on.
+
+    There is a limit of logs per IP to avoid flooding the logs in case of
+    misbehaving clients.
+    """
+
+    log_counters: ClassVar[dict[str, "HourlyLogCounter"]] = {}
+    max_logs: int = 200
+
+    def get_request_host(self, request: Request) -> str:
+        return request.client.host if request.client else "unknown"
+
+    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
+        response = await call_next(request)
+
+        host = self.get_request_host(request)
+        counter = self.log_counters.setdefault(host, HourlyLogCounter())
+        if response.status_code in (412, 422) and counter.get_count() < self.max_logs:
+            counter.log_event()
+
+            response_body = b""
+            chunk: bytes
+            async for chunk in response.body_iterator:  # type: ignore
+                response_body += chunk
+
+            logger.info(
+                f"Client payload validation error",
+                extra={
+                    "request_method": request.method,
+                    "request_path": request.url.path,
+                    "response_status_code": response.status_code,
+                    "response_payload": response_body.decode("utf-8", errors="replace"),
+                },
+            )
+            # Recreate the response body iterator since it has been consumed
+            response = Response(
+                content=response_body,
+                status_code=response.status_code,
+                headers=dict(response.headers),
+                media_type=response.media_type,
+                background=response.background,
+            )
+        return response
+
+
+class EventCounter:
+    def __init__(self, window_seconds: int = 3600):
+        self.window_seconds = window_seconds
+        self.events: deque[float] = deque()
+
+    def log_event(self):
+        current_time = time.time()
+        # Remove events older than the window
+        while self.events and self.events[0] < current_time - self.window_seconds:
+            self.events.popleft()
+        # Add current event
+        self.events.append(current_time)
+
+    def get_count(self) -> int:
+        current_time = time.time()
+        # Remove old events and return count
+        while self.events and self.events[0] < current_time - self.window_seconds:
+            self.events.popleft()
+        return len(self.events)
+
+
+class HourlyLogCounter(EventCounter):
+    def __init__(self):
+        super().__init__(window_seconds=3600)

nucliadb/migrator/datamanager.py

@@ -17,7 +17,6 @@
 # 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 Optional, Union
 
 from nucliadb.common import datamanagers
 from nucliadb.common.maindb.driver import Driver
@@ -69,7 +68,7 @@ class MigrationsDataManager:
             await txn.delete(MIGRATIONS_KEY.format(kbid=kbid))
             await txn.commit()
 
-    async def get_kb_info(self, kbid: str) -> Optional[KnowledgeBoxInfo]:
+    async def get_kb_info(self, kbid: str) -> KnowledgeBoxInfo | None:
         async with self.driver.ro_transaction() as txn:
             kb_config = await datamanagers.kb.get_config(txn, kbid=kbid)
             if kb_config is None:
@@ -91,8 +90,8 @@ class MigrationsDataManager:
     async def update_global_info(
         self,
         *,
-        current_version: Union[int, _Unset] = _UNSET,
-        target_version: Union[int, None, _Unset] = _UNSET,
+        current_version: int | _Unset = _UNSET,
+        target_version: int | None | _Unset = _UNSET,
     ) -> None:
         async with self.driver.rw_transaction() as txn:
             raw_pb = await txn.get(MIGRATION_INFO_KEY, for_update=True)

nucliadb/migrator/migrator.py

@@ -19,7 +19,6 @@
 #
 import asyncio
 import logging
-from typing import Optional
 
 from nucliadb.common import locking
 from nucliadb.common.cluster.rollover import rollover_kb_index
@@ -233,7 +232,7 @@ async def run_pg_schema_migrations(driver: PGDriver):
                 await tx.commit()
 
 
-async def run(context: ExecutionContext, target_version: Optional[int] = None) -> None:
+async def run(context: ExecutionContext, target_version: int | None = None) -> None:
     # Run schema migrations first, since they create the `resources` table needed for the lock below
     # Schema migrations use their own locking system
     if isinstance(context.kv_driver, PGDriver):

nucliadb/migrator/models.py

@@ -19,7 +19,6 @@
 #
 import types
 from dataclasses import dataclass
-from typing import Optional
 
 
 @dataclass
@@ -36,4 +35,4 @@ class KnowledgeBoxInfo:
 @dataclass
 class GlobalInfo:
     current_version: int
-    target_version: Optional[int] = None
+    target_version: int | None = None

nucliadb/migrator/settings.py

@@ -17,14 +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 typing import Optional
 
 import pydantic
 import pydantic_settings
 
 
 class Settings(pydantic_settings.BaseSettings):
-    redis_url: Optional[str] = None
+    redis_url: str | None = None
     max_concurrent_migrations: int = pydantic.Field(
         default=5,
         description="Maximum number of concurrent KB migrations allowed.",