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

nucliadb/common/cluster/rollover.py

@@ -21,7 +21,6 @@ import argparse
 import asyncio
 import logging
 from datetime import datetime
-from typing import Optional
 
 from nidx_protos.nodewriter_pb2 import (
     NewShardRequest,
@@ -34,10 +33,12 @@ from nucliadb.common.external_index_providers.base import ExternalIndexManager
 from nucliadb.common.external_index_providers.manager import (
     get_external_index_manager,
 )
+from nucliadb.common.maindb.utils import get_driver
 from nucliadb.common.nidx import get_nidx_api_client
 from nucliadb.common.vector_index_config import nucliadb_index_config_to_nidx
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.migrator.settings import settings
-from nucliadb_protos import utils_pb2, writer_pb2
+from nucliadb_protos import writer_pb2
 from nucliadb_telemetry import errors
 
 from .utils import (
@@ -45,6 +46,7 @@ from .utils import (
     get_resource,
     get_rollover_resource_index_message,
     index_resource_to_shard,
+    wait_for_nidx,
 )
 
 logger = logging.getLogger(__name__)
@@ -59,7 +61,7 @@ class UnexpectedRolloverError(Exception):
 async def create_rollover_index(
     app_context: ApplicationContext,
     kbid: str,
-    external: Optional[ExternalIndexManager] = None,
+    external: ExternalIndexManager | None = None,
 ) -> None:
     """
     Creates a new index for a knowledgebox in the index node cluster (and to the external index provider if configured).
@@ -148,7 +150,6 @@ async def create_rollover_shards(
 
             req = NewShardRequest(
                 kbid=kbid,
-                release_channel=utils_pb2.ReleaseChannel.STABLE,
                 vectorsets_configs=vectorsets,
             )
 
@@ -172,7 +173,7 @@ async def create_rollover_shards(
         return kb_shards
 
 
-def _get_shard(shards: writer_pb2.Shards, shard_id: str) -> Optional[writer_pb2.ShardObject]:
+def _get_shard(shards: writer_pb2.Shards, shard_id: str) -> writer_pb2.ShardObject | None:
     for shard in shards.shards:
         if shard_id == shard.shard:
             return shard
@@ -220,7 +221,7 @@ def _to_ts(dt: datetime) -> int:
 
 
 async def index_to_rollover_index(
-    app_context: ApplicationContext, kbid: str, external: Optional[ExternalIndexManager] = None
+    app_context: ApplicationContext, kbid: str, external: ExternalIndexManager | None = None
 ) -> None:
     """
     Indexes all data in a kb in rollover indexes. This happens before the cutover.
@@ -254,6 +255,7 @@ async def index_to_rollover_index(
             for rid in resource_ids
         ]
         await asyncio.gather(*batch)
+        await wait_for_indexing_to_catch_up(app_context)
 
     async with datamanagers.with_transaction() as txn:
         state.resources_indexed = True
@@ -262,12 +264,28 @@ async def index_to_rollover_index(
         await txn.commit()
 
 
+async def wait_for_indexing_to_catch_up(app_context: ApplicationContext):
+    try:
+        app_context.nats_manager
+    except AssertionError:
+        logger.warning("Nats manager not initialized. Cannot wait for indexing to catch up")
+        return
+    max_pending = 1000
+    while True:
+        try:
+            await wait_for_nidx(app_context.nats_manager, max_wait_seconds=60, max_pending=max_pending)
+            return
+        except asyncio.TimeoutError:
+            logger.warning(f"Nidx is behind more than {max_pending} messages. Throttling rollover.")
+            await asyncio.sleep(30)
+
+
 async def _index_resource_to_rollover_index(
     app_context: ApplicationContext,
     rollover_shards: writer_pb2.Shards,
     kbid: str,
     resource_id: str,
-    external: Optional[ExternalIndexManager] = None,
+    external: ExternalIndexManager | None = None,
 ) -> None:
     async with resource_index_semaphore:
         async with datamanagers.with_transaction() as txn:
@@ -321,7 +339,7 @@ async def _index_resource_to_rollover_index(
 
 
 async def cutover_index(
-    app_context: ApplicationContext, kbid: str, external: Optional[ExternalIndexManager] = None
+    app_context: ApplicationContext, kbid: str, external: ExternalIndexManager | None = None
 ) -> None:
     """
     Swaps our the current active index for a knowledgebox.
@@ -415,9 +433,18 @@ async def cutover_shards(app_context: ApplicationContext, kbid: str) -> None:
 
         await txn.commit()
 
+    # For KBs with pre-warm enabled, we must configure the new shards. There may
+    # be some small delay between this call and the shards being actually
+    # prewarmed, but rollovers are quite unusual and we prefer this rather than
+    # prewarming old and new shards at the same time
+    kb_config = await datamanagers.atomic.kb.get_config(kbid=kbid)
+    if kb_config is not None and kb_config.prewarm_enabled:
+        driver = get_driver()
+        await KnowledgeBox.configure_shards(driver, kbid, prewarm=True)
+
 
 async def validate_indexed_data(
-    app_context: ApplicationContext, kbid: str, external: Optional[ExternalIndexManager] = None
+    app_context: ApplicationContext, kbid: str, external: ExternalIndexManager | None = None
 ) -> list[str]:
     """
     Goes through all the resources in a knowledgebox and validates it

nucliadb/common/cluster/settings.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import enum
-from typing import Optional
 
 from pydantic import Field
 from pydantic_settings import BaseSettings
@@ -42,7 +41,7 @@ class Settings(BaseSettings):
         description="Maximum number of paragraphs to target per shard",
     )
     max_resource_paragraphs: int = Field(
-        default=50_000,
+        default=300_000,
         title="Max paragraphs per resource",
         description="Maximum number of paragraphs allowed on a single resource",
     )
@@ -52,13 +51,9 @@ class Settings(BaseSettings):
         description="Maximum number of entity labels (/e/) per field that are indexed (excess is not indexed)",
     )
 
-    nidx_api_address: Optional[str] = Field(default=None, description="NIDX gRPC API address")
-    nidx_searcher_address: Optional[str] = Field(
-        default=None, description="NIDX gRPC searcher API address"
-    )
-    nidx_indexer_address: Optional[str] = Field(
-        default=None, description="NIDX gRPC indexer API address"
-    )
+    nidx_api_address: str | None = Field(default=None, description="NIDX gRPC API address")
+    nidx_searcher_address: str | None = Field(default=None, description="NIDX gRPC searcher API address")
+    nidx_indexer_address: str | None = Field(default=None, description="NIDX gRPC indexer API address")
 
 
 settings = Settings()

nucliadb/common/cluster/utils.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Optional, Union
+from typing import TYPE_CHECKING
 
 import backoff
 from nidx_protos import nodereader_pb2
@@ -32,6 +32,7 @@ from nucliadb.common.cluster.settings import settings
 from nucliadb.ingest.orm import index_message
 from nucliadb.ingest.orm.resource import Resource
 from nucliadb_protos import writer_pb2
+from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -45,13 +46,13 @@ logger = logging.getLogger(__name__)
 _lock = asyncio.Lock()
 
 
-async def setup_cluster() -> Union[KBShardManager, StandaloneKBShardManager]:
+async def setup_cluster() -> KBShardManager | StandaloneKBShardManager:
     async with _lock:
         if get_utility(Utility.SHARD_MANAGER) is not None:
             # already setup
             return get_utility(Utility.SHARD_MANAGER)
 
-        mng: Union[KBShardManager, StandaloneKBShardManager]
+        mng: KBShardManager | StandaloneKBShardManager
         if settings.standalone_mode:
             mng = StandaloneKBShardManager()
         else:
@@ -69,17 +70,17 @@ def get_shard_manager() -> KBShardManager:
     return get_utility(Utility.SHARD_MANAGER)  # type: ignore
 
 
-async def get_resource(kbid: str, resource_id: str) -> Optional[Resource]:
+async def get_resource(kbid: str, resource_id: str) -> Resource | None:
     async with datamanagers.with_ro_transaction() as txn:
-        return await datamanagers.resources.get_resource(txn, kbid=kbid, rid=resource_id)
+        return await Resource.get(txn, kbid=kbid, rid=resource_id)
 
 
 @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
 async def get_rollover_resource_index_message(
     kbid: str, resource_id: str
-) -> Optional[nodereader_pb2.Resource]:
+) -> nodereader_pb2.Resource | None:
     async with datamanagers.with_ro_transaction() as txn:
-        resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=resource_id)
+        resource = await Resource.get(txn, kbid=kbid, rid=resource_id)
         if resource is None:
             logger.warning(
                 "Resource not found while indexing, skipping",
@@ -97,7 +98,7 @@ async def index_resource_to_shard(
     kbid: str,
     resource_id: str,
     shard: writer_pb2.ShardObject,
-    resource_index_message: Optional[nodereader_pb2.Resource] = None,
+    resource_index_message: nodereader_pb2.Resource | None = None,
 ) -> None:
     logger.info("Indexing resource", extra={"kbid": kbid, "resource_id": resource_id})
     sm = app_context.shard_manager
@@ -125,3 +126,28 @@ async def delete_resource_from_shard(
     partition = partitioning.generate_partition(kbid, resource_id)
 
     await sm.delete_resource(shard, resource_id, 0, str(partition), kbid)
+
+
+async def get_nats_consumer_pending_messages(
+    nats_manager: NatsConnectionManager, *, stream: str, consumer: str
+) -> int:
+    # get raw js client
+    js = nats_manager.js
+    consumer_info = await js.consumer_info(stream, consumer)
+    return consumer_info.num_pending
+
+
+async def wait_for_nidx(
+    nats_manager: NatsConnectionManager,
+    max_pending: int,
+    poll_interval_seconds: int = 5,
+    max_wait_seconds: int = 60,
+):
+    async with asyncio.timeout(max_wait_seconds):  # type: ignore
+        while True:
+            pending = await get_nats_consumer_pending_messages(
+                nats_manager, stream="nidx", consumer="nidx"
+            )
+            if pending < max_pending:
+                return
+            await asyncio.sleep(poll_interval_seconds)

nucliadb/common/context/__init__.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-from typing import Optional
 
 from nucliadb.common.cluster.manager import KBShardManager
 from nucliadb.common.cluster.settings import in_standalone_mode
@@ -58,13 +57,13 @@ class ApplicationContext:
         self.service_name = service_name
         self._initialized: bool = False
         self._lock = asyncio.Lock()
-        self._kv_driver: Optional[Driver] = None
-        self._blob_storage: Optional[Storage] = None
-        self._shard_manager: Optional[KBShardManager] = None
-        self._partitioning: Optional[PartitionUtility] = None
-        self._nats_manager: Optional[NatsConnectionManager] = None
-        self._transaction: Optional[TransactionUtility] = None
-        self._nidx: Optional[NidxUtility] = None
+        self._kv_driver: Driver | None = None
+        self._blob_storage: Storage | None = None
+        self._shard_manager: KBShardManager | None = None
+        self._partitioning: PartitionUtility | None = None
+        self._nats_manager: NatsConnectionManager | None = None
+        self._transaction: TransactionUtility | None = None
+        self._nidx: NidxUtility | None = None
         self.enabled_kv_driver = kv_driver
         self.enabled_blob_storage = blob_storage
         self.enabled_shard_manager = shard_manager

nucliadb/common/context/fastapi.py

@@ -19,7 +19,6 @@
 #
 
 from contextlib import asynccontextmanager
-from typing import Optional
 
 from fastapi import FastAPI
 from starlette.routing import Mount
@@ -28,7 +27,7 @@ from nucliadb.common.context import ApplicationContext
 
 
 @asynccontextmanager
-async def inject_app_context(app: FastAPI, context: Optional[ApplicationContext] = None):
+async def inject_app_context(app: FastAPI, context: ApplicationContext | None = None):
     if context is None:
         context = ApplicationContext()
 

nucliadb/common/datamanagers/__init__.py

@@ -31,7 +31,6 @@
 from . import (
     atomic,
     cluster,
-    entities,
     exceptions,
     fields,
     kb,
@@ -47,7 +46,6 @@ from .utils import with_ro_transaction, with_rw_transaction, with_transaction
 __all__ = (
     "atomic",
     "cluster",
-    "entities",
     "exceptions",
     "fields",
     "kb",
@@ -57,7 +55,7 @@ __all__ = (
     "search_configurations",
     "synonyms",
     "vectorsets",
-    "with_transaction",
-    "with_rw_transaction",
     "with_ro_transaction",
+    "with_rw_transaction",
+    "with_transaction",
 )

nucliadb/common/datamanagers/atomic.py

@@ -35,13 +35,15 @@ it's transaction
 
 """
 
+from collections.abc import Awaitable, Callable
 from functools import wraps
-from typing import Awaitable, Callable, TypeVar
+from typing import Concatenate, TypeVar
 
-from typing_extensions import Concatenate, ParamSpec
+from typing_extensions import ParamSpec
 
 from nucliadb.common.maindb.driver import Transaction
 
+from . import cluster as cluster_dm
 from . import kb as kb_dm
 from . import labels as labels_dm
 from . import resources as resources_dm
@@ -73,6 +75,10 @@ def rw_txn_wrap(fun: Callable[Concatenate[Transaction, P], Awaitable[T]]) -> Cal
     return wrapper
 
 
+class cluster:
+    get_kb_shards = ro_txn_wrap(cluster_dm.get_kb_shards)
+
+
 class kb:
     exists_kb = ro_txn_wrap(kb_dm.exists_kb)
     get_config = ro_txn_wrap(kb_dm.get_config)
@@ -83,6 +89,7 @@ class resources:
     get_resource_uuid_from_slug = ro_txn_wrap(resources_dm.get_resource_uuid_from_slug)
     resource_exists = ro_txn_wrap(resources_dm.resource_exists)
     slug_exists = ro_txn_wrap(resources_dm.slug_exists)
+    get_all_field_ids = ro_txn_wrap(resources_dm.get_all_field_ids)
 
 
 class labelset:

nucliadb/common/datamanagers/cluster.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import Optional
 
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb_protos import writer_pb2
@@ -33,7 +32,7 @@ KB_SHARDS = "/kbs/{kbid}/shards"
 
 async def get_kb_shards(
     txn: Transaction, *, kbid: str, for_update: bool = False
-) -> Optional[writer_pb2.Shards]:
+) -> writer_pb2.Shards | None:
     key = KB_SHARDS.format(kbid=kbid)
     return await get_kv_pb(txn, key, writer_pb2.Shards, for_update=for_update)
 

nucliadb/common/datamanagers/fields.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
 
 from google.protobuf.message import Message
 
@@ -34,7 +33,7 @@ KB_RESOURCE_FIELD_STATUS = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/status"
 
 async def get_raw(
     txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
-) -> Optional[bytes]:
+) -> bytes | None:
     key = KB_RESOURCE_FIELD.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
     return await txn.get(key)
 
@@ -62,7 +61,7 @@ async def delete(txn: Transaction, *, kbid: str, rid: str, field_type: str, fiel
 
 async def get_error(
     txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
-) -> Optional[writer_pb2.Error]:
+) -> writer_pb2.Error | None:
     key = KB_RESOURCE_FIELD_ERROR.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
     return await get_kv_pb(txn, key, writer_pb2.Error)
 
@@ -85,7 +84,7 @@ async def set_error(
 
 async def get_status(
     txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
-) -> Optional[writer_pb2.FieldStatus]:
+) -> writer_pb2.FieldStatus | None:
     key = KB_RESOURCE_FIELD_STATUS.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
     return await get_kv_pb(txn, key, writer_pb2.FieldStatus)
 

nucliadb/common/datamanagers/kb.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import AsyncIterator, Optional
+from collections.abc import AsyncIterator
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.maindb.driver import Transaction
@@ -47,7 +47,7 @@ async def exists_kb(txn: Transaction, *, kbid: str) -> bool:
     return await get_config(txn, kbid=kbid, for_update=False) is not None
 
 
-async def get_kb_uuid(txn: Transaction, *, slug: str) -> Optional[str]:
+async def get_kb_uuid(txn: Transaction, *, slug: str) -> str | None:
     uuid = await txn.get(KB_SLUGS.format(slug=slug), for_update=False)
     if uuid is not None:
         return uuid.decode()
@@ -67,7 +67,7 @@ async def delete_kb_slug(txn: Transaction, *, slug: str):
 
 async def get_config(
     txn: Transaction, *, kbid: str, for_update: bool = False
-) -> Optional[knowledgebox_pb2.KnowledgeBoxConfig]:
+) -> knowledgebox_pb2.KnowledgeBoxConfig | None:
     key = KB_UUID.format(kbid=kbid)
     payload = await txn.get(key, for_update=for_update)
     if payload is None:
@@ -105,8 +105,8 @@ async def get_matryoshka_vector_dimension(
     txn: Transaction,
     *,
     kbid: str,
-    vectorset_id: Optional[str] = None,
-) -> Optional[int]:
+    vectorset_id: str | None = None,
+) -> int | None:
     """Return vector dimension for matryoshka models"""
     from . import vectorsets
 
@@ -145,7 +145,7 @@ async def get_matryoshka_vector_dimension(
 
 async def get_external_index_provider_metadata(
     txn: Transaction, *, kbid: str
-) -> Optional[knowledgebox_pb2.StoredExternalIndexProviderMetadata]:
+) -> knowledgebox_pb2.StoredExternalIndexProviderMetadata | None:
     kb_config = await get_config(txn, kbid=kbid)
     if kb_config is None:
         return None

nucliadb/common/datamanagers/labels.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import Optional
 
 import orjson
 
@@ -50,7 +49,7 @@ async def get_labels(txn: Transaction, *, kbid: str) -> kb_pb2.Labels:
     return labels
 
 
-async def _get_labelset_ids(txn: Transaction, *, kbid: str) -> Optional[list[str]]:
+async def _get_labelset_ids(txn: Transaction, *, kbid: str) -> list[str] | None:
     key = KB_LABELSET_IDS.format(kbid=kbid)
     data = await txn.get(key, for_update=True)
     if not data:
@@ -84,7 +83,7 @@ async def _set_labelset_ids(txn: Transaction, *, kbid: str, labelsets: list[str]
     await txn.set(key, data)
 
 
-async def get_labelset(txn: Transaction, *, kbid: str, labelset_id: str) -> Optional[kb_pb2.LabelSet]:
+async def get_labelset(txn: Transaction, *, kbid: str, labelset_id: str) -> kb_pb2.LabelSet | None:
     labelset_key = KB_LABELSET.format(kbid=kbid, id=labelset_id)
     payload = await txn.get(labelset_key)
     if payload:

nucliadb/common/datamanagers/resources.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 TYPE_CHECKING, AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 import backoff
 
@@ -28,14 +28,9 @@ from nucliadb.common.maindb.exceptions import ConflictError, NotFoundError
 # These should be refactored
 from nucliadb.ingest.settings import settings as ingest_settings
 from nucliadb_protos import resources_pb2
-from nucliadb_utils.utilities import get_storage
 
 from .utils import with_ro_transaction
 
-if TYPE_CHECKING:
-    from nucliadb.ingest.orm.resource import Resource as ResourceORM
-
-
 KB_RESOURCE_BASIC = "/kbs/{kbid}/r/{uuid}"
 KB_RESOURCE_BASIC_FS = "/kbs/{kbid}/r/{uuid}/basic"  # Only used on FS driver
 KB_RESOURCE_ORIGIN = "/kbs/{kbid}/r/{uuid}/origin"
@@ -61,7 +56,7 @@ async def resource_exists(txn: Transaction, *, kbid: str, rid: str) -> bool:
 # id and slug
 
 
-async def get_resource_uuid_from_slug(txn: Transaction, *, kbid: str, slug: str) -> Optional[str]:
+async def get_resource_uuid_from_slug(txn: Transaction, *, kbid: str, slug: str) -> str | None:
     encoded_uuid = await txn.get(KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug, for_update=False))
     if not encoded_uuid:
         return None
@@ -70,7 +65,7 @@ async def get_resource_uuid_from_slug(txn: Transaction, *, kbid: str, slug: str)
 
 async def slug_exists(txn: Transaction, *, kbid: str, slug: str) -> bool:
     key = KB_RESOURCE_SLUG.format(kbid=kbid, slug=slug)
-    encoded_slug: Optional[bytes] = await txn.get(key)
+    encoded_slug: bytes | None = await txn.get(key)
     return encoded_slug not in (None, b"")
 
 
@@ -102,7 +97,7 @@ async def modify_slug(txn: Transaction, *, kbid: str, rid: str, new_slug: str) -
 @backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
 async def get_resource_shard_id(
     txn: Transaction, *, kbid: str, rid: str, for_update: bool = False
-) -> Optional[str]:
+) -> str | None:
     key = KB_RESOURCE_SHARD.format(kbid=kbid, uuid=rid)
     shard = await txn.get(key, for_update=for_update)
     if shard is not None:
@@ -118,7 +113,7 @@ async def set_resource_shard_id(txn: Transaction, *, kbid: str, rid: str, shard:
 # Basic
 
 
-async def get_basic(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Basic]:
+async def get_basic(txn: Transaction, *, kbid: str, rid: str) -> resources_pb2.Basic | None:
     raw = await get_basic_raw(txn, kbid=kbid, rid=rid)
     if raw is None:
         return None
@@ -127,7 +122,7 @@ async def get_basic(txn: Transaction, *, kbid: str, rid: str) -> Optional[resour
     return basic
 
 
-async def get_basic_raw(txn: Transaction, *, kbid: str, rid: str) -> Optional[bytes]:
+async def get_basic_raw(txn: Transaction, *, kbid: str, rid: str) -> bytes | None:
     if ingest_settings.driver == "local":
         raw_basic = await txn.get(KB_RESOURCE_BASIC_FS.format(kbid=kbid, uuid=rid))
     else:
@@ -151,7 +146,7 @@ async def set_basic(txn: Transaction, *, kbid: str, rid: str, basic: resources_p
 # Origin
 
 
-async def get_origin(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Origin]:
+async def get_origin(txn: Transaction, *, kbid: str, rid: str) -> resources_pb2.Origin | None:
     key = KB_RESOURCE_ORIGIN.format(kbid=kbid, uuid=rid)
     return await get_kv_pb(txn, key, resources_pb2.Origin)
 
@@ -164,7 +159,7 @@ async def set_origin(txn: Transaction, *, kbid: str, rid: str, origin: resources
 # Extra
 
 
-async def get_extra(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Extra]:
+async def get_extra(txn: Transaction, *, kbid: str, rid: str) -> resources_pb2.Extra | None:
     key = KB_RESOURCE_EXTRA.format(kbid=kbid, uuid=rid)
     return await get_kv_pb(txn, key, resources_pb2.Extra)
 
@@ -177,7 +172,7 @@ async def set_extra(txn: Transaction, *, kbid: str, rid: str, extra: resources_p
 # Security
 
 
-async def get_security(txn: Transaction, *, kbid: str, rid: str) -> Optional[resources_pb2.Security]:
+async def get_security(txn: Transaction, *, kbid: str, rid: str) -> resources_pb2.Security | None:
     key = KB_RESOURCE_SECURITY.format(kbid=kbid, uuid=rid)
     return await get_kv_pb(txn, key, resources_pb2.Security)
 
@@ -265,7 +260,7 @@ async def set_number_of_resources(txn: Transaction, kbid: str, value: int) -> No
 
 async def get_all_field_ids(
     txn: Transaction, *, kbid: str, rid: str, for_update: bool = False
-) -> Optional[resources_pb2.AllFieldIDs]:
+) -> resources_pb2.AllFieldIDs | None:
     key = KB_RESOURCE_ALL_FIELDS.format(kbid=kbid, uuid=rid)
     return await get_kv_pb(txn, key, resources_pb2.AllFieldIDs, for_update=for_update)
 
@@ -285,21 +280,3 @@ async def has_field(txn: Transaction, *, kbid: str, rid: str, field_id: resource
         if field_id == resource_field_id:
             return True
     return False
-
-
-# ORM mix (this functions shouldn't belong here)
-
-
-@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
-async def get_resource(txn: Transaction, *, kbid: str, rid: str) -> Optional["ResourceORM"]:
-    """
-    Not ideal to return Resource type here but refactoring would
-    require a lot of changes.
-
-    At least this isolated that dependency here.
-    """
-    # prevent circulat imports -- this is not ideal that we have the ORM mix here.
-    from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
-
-    kb_orm = KnowledgeBoxORM(txn, await get_storage(), kbid)
-    return await kb_orm.get(rid)

nucliadb/common/datamanagers/rollover.py

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import AsyncGenerator, Optional
+from collections.abc import AsyncGenerator
 
 import orjson
 from pydantic import BaseModel
@@ -56,7 +56,7 @@ class RolloverStateNotFoundError(Exception):
     ...
 
 
-async def get_kb_rollover_shards(txn: Transaction, *, kbid: str) -> Optional[writer_pb2.Shards]:
+async def get_kb_rollover_shards(txn: Transaction, *, kbid: str) -> writer_pb2.Shards | None:
     key = KB_ROLLOVER_SHARDS.format(kbid=kbid)
     return await get_kv_pb(txn, key, writer_pb2.Shards)
 
@@ -90,7 +90,7 @@ async def add_batch_to_index(txn: Transaction, *, kbid: str, batch: list[str]) -
         await txn.set(key, b"")
 
 
-async def get_to_index(txn: Transaction, *, kbid: str, count: int) -> Optional[list[str]]:
+async def get_to_index(txn: Transaction, *, kbid: str, count: int) -> list[str] | None:
     key = KB_ROLLOVER_RESOURCES_TO_INDEX.format(kbid=kbid, resource="")
     found = [key async for key in txn.keys(key, count=count)]
     if found:
@@ -118,9 +118,7 @@ async def add_indexed(
     await txn.set(indexed, orjson.dumps(data))
 
 
-async def get_indexed_data(
-    txn: Transaction, *, kbid: str, resource_id: str
-) -> Optional[tuple[str, int]]:
+async def get_indexed_data(txn: Transaction, *, kbid: str, resource_id: str) -> tuple[str, int] | None:
     key = KB_ROLLOVER_RESOURCES_INDEXED.format(kbid=kbid, resource=resource_id)
     val = await txn.get(key)
     if val is not None:
@@ -213,7 +211,7 @@ async def update_kb_rollover_external_index_metadata(
 
 async def get_kb_rollover_external_index_metadata(
     txn: Transaction, *, kbid: str
-) -> Optional[kb_pb2.StoredExternalIndexProviderMetadata]:
+) -> kb_pb2.StoredExternalIndexProviderMetadata | None:
     key = KB_ROLLOVER_EXTERNAL_INDEX_METADATA.format(kbid=kbid)
     val = await txn.get(key)
     if not val:

nucliadb/common/datamanagers/search_configurations.py

@@ -18,7 +18,6 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import logging
-from typing import Optional
 
 from pydantic import TypeAdapter
 
@@ -31,7 +30,7 @@ KB_SEARCH_CONFIGURATION_PREFIX = "/kbs/{kbid}/search_configuration"
 KB_SEARCH_CONFIGURATION = "/kbs/{kbid}/search_configuration/{name}"
 
 
-async def get(txn: Transaction, *, kbid: str, name: str) -> Optional[SearchConfiguration]:
+async def get(txn: Transaction, *, kbid: str, name: str) -> SearchConfiguration | None:
     key = KB_SEARCH_CONFIGURATION.format(kbid=kbid, name=name)
     data = await txn.get(key, for_update=True)
     if not data: