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

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,
@@ -67,6 +66,11 @@ class DriverSettings(BaseSettings):
     )
 
 
+class CatalogConfig(Enum):
+    UNSET = "unset"
+    PG = "pg"
+
+
 # For use during migration from pull v1 to pull v2
 class ProcessingPullMode(Enum):
     OFF = "off"
@@ -75,26 +79,43 @@ class ProcessingPullMode(Enum):
 
 
 class Settings(DriverSettings):
-    grpc_port: int = 8030
-
-    partitions: list[str] = ["1"]
+    # Catalog settings
+    catalog: CatalogConfig = Field(default=CatalogConfig.PG, description="Catalog backend")
 
+    # Pull worker settings
     pull_time_error_backoff: int = 30
     pull_api_timeout: int = 60
-    disable_pull_worker: bool = False
+    disable_pull_worker: bool = Field(
+        default=False, description="Set to true to disable the pull worker task"
+    )
 
-    # ingest consumer sts replica settings
-    replica_number: int = -1
-    total_replicas: int = 1  # number of ingest processor replicas in the cluster
-    nuclia_partitions: int = 50
+    # Ingest consumer sts replica settings
+    replica_number: int = Field(
+        default=-1,
+        description="The replica number of this ingest statefulset instance. Leave to -1 to auto-assign based on hostname.",
+    )
+    total_replicas: int = Field(default=1, description="Number of ingest statefulset replicas deployed")
+    nuclia_partitions: int = Field(
+        default=50, description="Total number of partitions of the nats stream."
+    )
+    partitions: list[str] = Field(
+        default=["1"],
+        description="List of partitions assigned to this ingest statefulset instance. This is automatically assigned based on the replica number and total replicas.",
+    )
+    max_concurrent_ingest_processing: int = Field(
+        default=5,
+        description="Controls the number of concurrent messages from different partitions that can be processed at the same time by ingest statefulset consumers.",
+    )
 
-    max_receive_message_length: int = 500  # In MB
+    # Grpc server settings
+    grpc_port: int = 8030
+    max_receive_message_length: int = Field(
+        default=500, description="Maximum receive grpc message length in MB."
+    )
 
     # Search query timeouts
     relation_search_timeout: float = 10.0
     relation_types_timeout: float = 10.0
 
-    max_concurrent_ingest_processing: int = 5
-
 
 settings = Settings()

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
@@ -43,8 +43,8 @@ logger = logging.getLogger(SERVICE_NAME)
 WHITELISTED_HEADERS = {
     "x-nucliadb-user",
     "x-nucliadb-roles",
-    "x-stf-roles",
-    "x-stf-user",
+    "x-nucliadb-account",
+    "x-nucliadb-account-type",
     "x-forwarded-for",
     "x-forwarded-host",
     "x-forwarded-proto",
@@ -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)
 
 
@@ -204,14 +204,14 @@ async def learning_config_proxy(
     request: Request,
     method: str,
     url: str,
-    extra_headers: Optional[dict[str, str]] = None,
-) -> Union[Response, StreamingResponse]:
+    headers: dict[str, str] = {},
+) -> Response | StreamingResponse:
     return await proxy(
         service=LearningService.CONFIG,
         request=request,
         method=method,
         url=url,
-        extra_headers=extra_headers,
+        headers=headers,
     )
 
 
@@ -244,24 +244,21 @@ async def proxy(
     request: Request,
     method: str,
     url: str,
-    extra_headers: Optional[dict[str, str]] = None,
-) -> Union[Response, StreamingResponse]:
+    headers: dict[str, str] = {},
+) -> Response | StreamingResponse:
     """
     Proxy the request to a learning API.
 
-    service: LearningService. The learning service to proxy the request to.
-    request: Request. The incoming request.
-    method: str. The HTTP method to use.
-    url: str. The URL to proxy the request to.
-    extra_headers: Optional[dict[str, str]]. Extra headers to include in the proxied request.
+    service: The learning service to proxy the request to.
+    request: The incoming request.
+    method: The HTTP method to use.
+    url: The URL to proxy the request to.
+    headers: Extra headers to include in the proxied request.
 
     Returns: Response. The response from the learning API. If the response is chunked, a StreamingResponse is returned.
     """
-
-    proxied_headers = extra_headers or {}
-    proxied_headers.update(
-        {k.lower(): v for k, v in request.headers.items() if is_white_listed_header(k)}
-    )
+    proxied_headers = {k.lower(): v for k, v in request.headers.items() if is_white_listed_header(k)}
+    proxied_headers.update(**headers)
 
     async with service_client(
         base_url=get_base_url(service=service),
@@ -429,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: ...
@@ -442,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:
@@ -486,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
@@ -37,8 +38,11 @@ MIGRATION_COUNT = metrics.Gauge("nucliadb_migration", labels={"type": "", "versi
 
 PENDING_RESOURCE_COUNT = metrics.Gauge("nucliadb_pending_resources_count")
 
+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.
     """
@@ -47,9 +51,11 @@ async def iter_kbids(context: ApplicationContext) -> AsyncGenerator[str, None]:
             yield kbid
 
 
-async def update_migration_metrics(context: ApplicationContext):
+async def update_kb_metrics(context: ApplicationContext):
     """
-    Report the global migration version and the number of KBs per migration version.
+    Report metrics at the kb level:
+      - total number of KBs
+      - the global migration version and the number of KBs per migration version.
     """
     # Clear previoulsy set values so that we report only the current state
     MIGRATION_COUNT.gauge.clear()
@@ -60,12 +66,16 @@ async def update_migration_metrics(context: ApplicationContext):
         MIGRATION_COUNT.set(1, labels=dict(type="global", version=str(global_info.current_version)))
 
     version_count: dict[str, int] = {}
+    n_kbs = 0
     async for kbid in iter_kbids(context):
         kb_info = await mdm.get_kb_info(kbid)
         if kb_info is not None:
             current_version = str(kb_info.current_version)
             version_count.setdefault(current_version, 0)
             version_count[current_version] += 1
+            n_kbs += 1
+
+    KBS_COUNT.set(n_kbs)
 
     for version, count in version_count.items():
         MIGRATION_COUNT.set(count, labels=dict(type="kb", version=version))
@@ -85,9 +95,13 @@ 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]
+        RESOURCES_COUNT.set(count)
+
 
 async def run_exporter_task(context: ApplicationContext, exporter_task: Callable, interval: int):
     """
@@ -109,7 +123,7 @@ async def run_exporter(context: ApplicationContext):
     # Schedule exporter tasks
     tasks = []
     for export_task, interval in [
-        (update_migration_metrics, 60 * 3),
+        (update_kb_metrics, 60 * 3),
         (update_resource_metrics, 60 * 5),
     ]:
         tasks.append(asyncio.create_task(run_exporter_task(context, export_task, interval=interval)))

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:
@@ -77,13 +76,7 @@ class MigrationsDataManager:
         return KnowledgeBoxInfo(current_version=kb_config.migration_version)
 
     async def update_kb_info(self, *, kbid: str, current_version: int) -> None:
-        async with self.driver.rw_transaction() as txn:
-            kb_config = await datamanagers.kb.get_config(txn, kbid=kbid, for_update=True)
-            if kb_config is None:
-                raise Exception(f"KB {kbid} does not exist")
-            kb_config.migration_version = current_version
-            await KnowledgeBoxORM.update(txn, kbid, config=kb_config)
-            await txn.commit()
+        await KnowledgeBoxORM.update(self.driver, kbid, migration_version=current_version)
 
     async def get_global_info(self) -> GlobalInfo:
         async with self.driver.ro_transaction() as txn:
@@ -97,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.",