nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/cluster/settings.py

@@ -18,6 +18,7 @@
 # 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
@@ -85,6 +86,11 @@ class Settings(BaseSettings):
     cluster_discovery_kubernetes_selector: str = "appType=node"
     cluster_discovery_manual_addresses: list[str] = []
 
+    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"
+    )
+
 
 settings = Settings()
 

nucliadb/common/cluster/standalone/grpc_node_binding.py

@@ -30,6 +30,7 @@ from nucliadb_protos.nodereader_pb2 import (
     DocumentItem,
     EdgeList,
     GetShardRequest,
+    IdCollection,
     ParagraphItem,
     ParagraphSearchRequest,
     ParagraphSearchResponse,
@@ -46,15 +47,13 @@ from nucliadb_protos.noderesources_pb2 import (
     EmptyResponse,
     Resource,
     ResourceID,
-)
-from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
-from nucliadb_protos.noderesources_pb2 import (
     ShardCreated,
     ShardId,
     ShardIds,
     VectorSetID,
     VectorSetList,
 )
+from nucliadb_protos.noderesources_pb2 import Shard as NodeResourcesShard
 from nucliadb_protos.nodewriter_pb2 import NewShardRequest, OpStatus
 
 from ..settings import settings
@@ -68,8 +67,7 @@ except ImportError:  # pragma: no cover
     IndexNodeException = Exception
 
 try:
-    from nucliadb_node_binding import NodeReader  # type: ignore
-    from nucliadb_node_binding import NodeWriter  # type: ignore
+    from nucliadb_node_binding import NodeReader, NodeWriter
 except ImportError:  # pragma: no cover
     NodeReader = None
     NodeWriter = None
@@ -80,15 +78,11 @@ class StandaloneReaderWrapper:
 
     def __init__(self):
         if NodeReader is None:
-            raise ImportError(
-                "NucliaDB index node bindings are not installed (reader not found)"
-            )
+            raise ImportError("NucliaDB index node bindings are not installed (reader not found)")
         self.reader = NodeReader()
         self.executor = ThreadPoolExecutor(settings.local_reader_threads)
 
-    async def Search(
-        self, request: SearchRequest, retry: bool = False
-    ) -> SearchResponse:
+    async def Search(self, request: SearchRequest, retry: bool = False) -> SearchResponse:
         try:
             loop = asyncio.get_running_loop()
             result = await loop.run_in_executor(
@@ -112,30 +106,6 @@ class StandaloneReaderWrapper:
             else:
                 raise
 
-    async def ParagraphSearch(
-        self, request: ParagraphSearchRequest
-    ) -> ParagraphSearchResponse:
-        loop = asyncio.get_running_loop()
-        result = await loop.run_in_executor(
-            self.executor, self.reader.paragraph_search, request.SerializeToString()
-        )
-        pb_bytes = bytes(result)
-        pb = ParagraphSearchResponse()
-        pb.ParseFromString(pb_bytes)
-        return pb
-
-    async def RelationSearch(
-        self, request: RelationSearchRequest
-    ) -> RelationSearchResponse:
-        loop = asyncio.get_running_loop()
-        result = await loop.run_in_executor(
-            self.executor, self.reader.relation_search, request.SerializeToString()
-        )
-        pb_bytes = bytes(result)
-        pb = RelationSearchResponse()
-        pb.ParseFromString(pb_bytes)
-        return pb
-
     async def GetShard(self, request: GetShardRequest) -> NodeResourcesShard:
         loop = asyncio.get_running_loop()
         result = await loop.run_in_executor(
@@ -200,9 +170,7 @@ class StandaloneReaderWrapper:
             raise exception
         await loop.run_in_executor(self.executor, t1.join)
 
-    async def Paragraphs(
-        self, stream_request: StreamRequest
-    ) -> AsyncIterator[ParagraphItem]:
+    async def Paragraphs(self, stream_request: StreamRequest) -> AsyncIterator[ParagraphItem]:
         loop = asyncio.get_running_loop()
         q: asyncio.Queue[ParagraphItem] = asyncio.Queue(1)
         exception = None
@@ -248,30 +216,15 @@ class StandaloneReaderWrapper:
         edge_list.ParseFromString(pb_bytes)
         return edge_list
 
-
-async def Search(self, request: SearchRequest, retry: bool = False) -> SearchResponse:
-    try:
+    async def VectorIds(self, request: VectorSetID) -> IdCollection:
         loop = asyncio.get_running_loop()
         result = await loop.run_in_executor(
-            self.executor, self.reader.search, request.SerializeToString()
+            self.executor, self.reader.vector_ids, request.SerializeToString()
         )
         pb_bytes = bytes(result)
-        pb = SearchResponse()
-        pb.ParseFromString(pb_bytes)
-        return pb
-    except IndexNodeException as exc:
-        if "IO error" not in str(exc):
-            # ignore any other error
-            raise
-
-        # try some mitigations...
-        logger.error(f"IndexNodeException in Search: {request}", exc_info=True)
-        if not retry:
-            # reinit?
-            self.reader = NodeReader()
-            return await self.Search(request, retry=True)
-        else:
-            raise
+        ids = IdCollection()
+        ids.ParseFromString(pb_bytes)
+        return ids
 
 
 class StandaloneWriterWrapper:
@@ -280,9 +233,7 @@ class StandaloneWriterWrapper:
     def __init__(self):
         os.makedirs(settings.data_path, exist_ok=True)
         if NodeWriter is None:
-            raise ImportError(
-                "NucliaDB index node bindings are not installed (writer not found)"
-            )
+            raise ImportError("NucliaDB index node bindings are not installed (writer not found)")
         self.writer = NodeWriter()
         self.executor = ThreadPoolExecutor(settings.local_writer_threads)
 
@@ -369,9 +320,7 @@ class StandaloneWriterWrapper:
 
     async def GC(self, request: ShardId) -> EmptyResponse:
         loop = asyncio.get_running_loop()
-        resp = await loop.run_in_executor(
-            self.executor, self.writer.gc, request.SerializeToString()
-        )
+        resp = await loop.run_in_executor(self.executor, self.writer.gc, request.SerializeToString())
         pb_bytes = bytes(resp)
         op_status = EmptyResponse()
         op_status.ParseFromString(pb_bytes)

nucliadb/common/cluster/standalone/index_node.py

@@ -20,10 +20,7 @@
 from typing import Any, Optional
 
 from nucliadb.common.cluster.base import AbstractIndexNode
-from nucliadb.common.cluster.grpc_node_dummy import (  # type: ignore
-    DummyReaderStub,
-    DummyWriterStub,
-)
+from nucliadb.common.cluster.grpc_node_dummy import DummyReaderStub, DummyWriterStub
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.common.cluster.standalone import grpc_node_binding
 from nucliadb_protos import standalone_pb2, standalone_pb2_grpc
@@ -79,7 +76,7 @@ class ProxyCallerWrapper:
         else:
             grpc_address = address
         self._channel = get_traced_grpc_channel(grpc_address, "standalone_proxy")
-        self._stub = standalone_pb2_grpc.StandaloneClusterServiceStub(self._channel)  # type: ignore
+        self._stub = standalone_pb2_grpc.StandaloneClusterServiceStub(self._channel)
 
     def __getattr__(self, name):
         async def call(request):
@@ -95,9 +92,7 @@ class ProxyCallerWrapper:
                 else:
                     raise NotImplementedError(f"Unknown type {self._type}")
             except KeyError:
-                raise NotImplementedError(
-                    f"Unknown method for type {self._type}: {name}"
-                )
+                raise NotImplementedError(f"Unknown method for type {self._type}: {name}")
             return_value = return_type()
             return_value.ParseFromString(resp.payload)
             return return_value
@@ -116,9 +111,7 @@ class ProxyStandaloneIndexNode(StandaloneIndexNode):
         available_disk: int,
         dummy: bool = False,
     ):
-        super().__init__(
-            id, address, shard_count, available_disk=available_disk, dummy=dummy
-        )
+        super().__init__(id, address, shard_count, available_disk=available_disk, dummy=dummy)
         if dummy:
             return
 

nucliadb/common/cluster/standalone/service.py

@@ -32,9 +32,7 @@ from nucliadb_protos import standalone_pb2, standalone_pb2_grpc
 from nucliadb_utils.grpc import get_traced_grpc_server
 
 
-class StandaloneClusterServiceServicer(
-    standalone_pb2_grpc.StandaloneClusterServiceServicer
-):
+class StandaloneClusterServiceServicer(standalone_pb2_grpc.StandaloneClusterServiceServicer):
     @backoff.on_exception(backoff.expo, (AioRpcError,), max_time=60)
     async def NodeAction(  # type: ignore
         self, request: standalone_pb2.NodeActionRequest, context
@@ -61,9 +59,7 @@ class StandaloneClusterServiceServicer(
         self, request: standalone_pb2.NodeInfoRequest, context
     ) -> standalone_pb2.NodeInfoResponse:
         index_node = get_self()
-        index_node.shard_count = len(
-            os.listdir(os.path.join(cluster_settings.data_path, "shards"))
-        )
+        index_node.shard_count = len(os.listdir(os.path.join(cluster_settings.data_path, "shards")))
         total_disk, _, available_disk = shutil.disk_usage(cluster_settings.data_path)
         return standalone_pb2.NodeInfoResponse(
             id=index_node.id,

nucliadb/common/cluster/standalone/utils.py

@@ -56,9 +56,7 @@ def get_self() -> StandaloneIndexNode:
     make another grpc request since this node can service it directly.
     """
     if not is_index_node():
-        raise Exception(
-            "This node is not an Index Node. You should not reach this code path."
-        )
+        raise Exception("This node is not an Index Node. You should not reach this code path.")
     global _SELF_INDEX_NODE
     node_id = get_standalone_node_id()
     if _SELF_INDEX_NODE is None or node_id != _SELF_INDEX_NODE.id:
@@ -68,9 +66,7 @@ def get_self() -> StandaloneIndexNode:
             host = f"{hn}.{ns}"
         else:
             host = gethostname()
-        _SELF_INDEX_NODE = StandaloneIndexNode(
-            id=node_id, address=host, shard_count=0, available_disk=0
-        )
+        _SELF_INDEX_NODE = StandaloneIndexNode(id=node_id, address=host, shard_count=0, available_disk=0)
     try:
         _, _, available_disk = shutil.disk_usage(cluster_settings.data_path)
         _SELF_INDEX_NODE.available_disk = available_disk

nucliadb/common/cluster/utils.py

@@ -37,7 +37,8 @@ from nucliadb.common.cluster.standalone.service import (
     start_grpc as start_standalone_grpc,
 )
 from nucliadb.common.cluster.standalone.utils import is_index_node
-from nucliadb_protos import noderesources_pb2, writer_pb2
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb_protos import nodereader_pb2, writer_pb2
 from nucliadb_utils import const
 from nucliadb_utils.settings import is_onprem_nucliadb
 from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
@@ -119,38 +120,44 @@ async def wait_for_node(app_context: ApplicationContext, node_id: str) -> None:
         await asyncio.sleep(sleep)
 
 
-@backoff.on_exception(
-    backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8
-)
+async def get_resource(kbid: str, resource_id: str) -> Optional[Resource]:
+    async with datamanagers.with_ro_transaction() as txn:
+        return await datamanagers.resources.get_resource(txn, kbid=kbid, rid=resource_id)
+
+
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
+async def get_resource_index_message(kbid: str, resource_id: str) -> Optional[nodereader_pb2.Resource]:
+    async with datamanagers.with_ro_transaction() as txn:
+        resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=resource_id)
+        if resource is None:
+            logger.warning(
+                "Resource not found while indexing, skipping",
+                extra={"kbid": kbid, "resource_id": resource_id},
+            )
+            return None
+        resource_index_message = (await resource.generate_index_message(reindex=False)).brain
+        return resource_index_message
+
+
+@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=8)
 async def index_resource_to_shard(
     app_context: ApplicationContext,
     kbid: str,
     resource_id: str,
     shard: writer_pb2.ShardObject,
-) -> Optional[noderesources_pb2.Resource]:
+    resource_index_message: Optional[nodereader_pb2.Resource] = None,
+) -> None:
     logger.info("Indexing resource", extra={"kbid": kbid, "resource_id": resource_id})
-
     sm = app_context.shard_manager
     partitioning = app_context.partitioning
 
-    async with datamanagers.with_transaction() as txn:
-        resource_index_message = (
-            await datamanagers.resources.get_resource_index_message(
-                txn, kbid=kbid, rid=resource_id
-            )
-        )
-
     if resource_index_message is None:
-        logger.warning(
-            "Resource index message not found while indexing, skipping",
-            extra={"kbid": kbid, "resource_id": resource_id},
-        )
-        return None
+        resource_index_message = await get_resource_index_message(kbid, resource_id)
+        if resource_index_message is None:
+            return
+
     partition = partitioning.generate_partition(kbid, resource_id)
-    await sm.add_resource(
-        shard, resource_index_message, txid=-1, partition=str(partition), kb=kbid
-    )
-    return resource_index_message
+    await sm.add_resource(shard, resource_index_message, txid=-1, partition=str(partition), kb=kbid)
 
 
 async def delete_resource_from_shard(

nucliadb/common/constants.py

@@ -0,0 +1,20 @@
+# 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/>.
+#
+AVG_PARAGRAPH_SIZE_BYTES = 10_000

nucliadb/common/context/__init__.py

@@ -24,6 +24,7 @@ from nucliadb.common.cluster.settings import in_standalone_mode
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
 from nucliadb.common.maindb.utils import setup_driver, teardown_driver
+from nucliadb.common.nidx import start_nidx_utility, stop_nidx_utility
 from nucliadb_utils.indexing import IndexingUtility
 from nucliadb_utils.nats import NatsConnectionManager
 from nucliadb_utils.partition import PartitionUtility
@@ -78,11 +79,13 @@ class ApplicationContext:
             )
             self.indexing = await start_indexing_utility()
         self.transaction = await start_transaction_utility(self.service_name)
+        self.nidx = await start_nidx_utility()
 
     async def finalize(self) -> None:
         if not self._initialized:
             return
 
+        await stop_nidx_utility()
         await stop_transaction_utility()
         if not in_standalone_mode():
             await stop_indexing_utility()

nucliadb/common/context/fastapi.py

@@ -18,25 +18,28 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
+from contextlib import asynccontextmanager
+
 from fastapi import FastAPI
 from starlette.routing import Mount
 
 from nucliadb.common.context import ApplicationContext
 
 
-def set_app_context(app: FastAPI):
+@asynccontextmanager
+async def inject_app_context(app: FastAPI):
     context = ApplicationContext()
 
     app.state.context = context
-    app.add_event_handler("startup", context.initialize)
-    app.add_event_handler("shutdown", context.finalize)
 
     # Need to add app context in all sub-applications
     for route in app.router.routes:
         if isinstance(route, Mount) and isinstance(route.app, FastAPI):
             route.app.state.context = context
-            route.app.add_event_handler("startup", context.initialize)
-            route.app.add_event_handler("shutdown", context.finalize)
+
+    await context.initialize()
+    yield context
+    await context.finalize()
 
 
 def get_app_context(application: FastAPI) -> ApplicationContext:

nucliadb/{tests/knowledgeboxes/__init__.py → common/counters.py}

@@ -18,5 +18,11 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from .philosophy_books import *  # noqa
-from .ten_dummy_resources import *  # noqa
+from dataclasses import dataclass
+
+
+@dataclass
+class IndexCounts:
+    fields: int
+    paragraphs: int
+    sentences: int

nucliadb/common/datamanagers/__init__.py

@@ -33,25 +33,31 @@ from . import (
     cluster,
     entities,
     exceptions,
+    fields,
     kb,
     labels,
     processing,
     resources,
     rollover,
     synonyms,
+    vectorsets,
 )
-from .utils import with_transaction
+from .utils import with_ro_transaction, with_rw_transaction, with_transaction
 
 __all__ = (
     "atomic",
     "cluster",
     "entities",
     "exceptions",
+    "fields",
     "kb",
     "labels",
     "processing",
     "resources",
     "rollover",
     "synonyms",
+    "vectorsets",
     "with_transaction",
+    "with_rw_transaction",
+    "with_ro_transaction",
 )

nucliadb/common/datamanagers/atomic.py

@@ -39,8 +39,10 @@ import sys
 from functools import wraps
 
 from . import kb as kb_dm
+from . import labels as labels_dm
 from . import resources as resources_dm
-from .utils import with_transaction
+from . import synonyms as synonyms_dm
+from .utils import with_ro_transaction, with_transaction
 
 # XXX: we are using the not exported _ParamSpec to support 3.9. Whenever we
 # upgrade to >= 3.10 we'll be able to use ParamSpecKwargs and improve the
@@ -49,7 +51,7 @@ from .utils import with_transaction
 
 __python_version = (sys.version_info.major, sys.version_info.minor)
 if __python_version == (3, 9):
-    from typing import _ParamSpec as ParamSpec  # type: ignore
+    from typing_extensions import ParamSpec
 else:
     from typing import ParamSpec  # type: ignore
 
@@ -59,7 +61,7 @@ P = ParamSpec("P")
 def ro_txn_wrap(fun: P) -> P:  # type: ignore
     @wraps(fun)
     async def wrapper(**kwargs: P.kwargs):
-        async with with_transaction(read_only=True) as txn:
+        async with with_ro_transaction() as txn:
             return await fun(txn, **kwargs)
 
     return wrapper
@@ -69,16 +71,32 @@ def rw_txn_wrap(fun: P) -> P:  # type: ignore
     @wraps(fun)
     async def wrapper(**kwargs: P.kwargs):
         async with with_transaction() as txn:
-            return await fun(txn, **kwargs)
+            result = await fun(txn, **kwargs)
+            await txn.commit()
+            return result
 
     return wrapper
 
 
 class kb:
     exists_kb = ro_txn_wrap(kb_dm.exists_kb)
+    get_config = ro_txn_wrap(kb_dm.get_config)
+    get_external_index_provider_metadata = ro_txn_wrap(kb_dm.get_external_index_provider_metadata)
 
 
 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)
+
+
+class labelset:
+    get = ro_txn_wrap(labels_dm.get_labelset)
+    set = rw_txn_wrap(labels_dm.set_labelset)
+    delete = rw_txn_wrap(labels_dm.delete_labelset)
+    get_all = ro_txn_wrap(labels_dm.get_labels)
+
+
+class synonyms:
+    get = ro_txn_wrap(synonyms_dm.get)
+    set = rw_txn_wrap(synonyms_dm.set)

nucliadb/common/datamanagers/cluster.py

@@ -31,13 +31,13 @@ logger = logging.getLogger(__name__)
 KB_SHARDS = "/kbs/{kbid}/shards"
 
 
-async def get_kb_shards(txn: Transaction, *, kbid: str) -> Optional[writer_pb2.Shards]:
+async def get_kb_shards(
+    txn: Transaction, *, kbid: str, for_update: bool = False
+) -> Optional[writer_pb2.Shards]:
     key = KB_SHARDS.format(kbid=kbid)
-    return await get_kv_pb(txn, key, writer_pb2.Shards)
+    return await get_kv_pb(txn, key, writer_pb2.Shards, for_update=for_update)
 
 
-async def update_kb_shards(
-    txn: Transaction, *, kbid: str, shards: writer_pb2.Shards
-) -> None:
+async def update_kb_shards(txn: Transaction, *, kbid: str, shards: writer_pb2.Shards) -> None:
     key = KB_SHARDS.format(kbid=kbid)
     await txn.set(key, shards.SerializeToString())

nucliadb/common/datamanagers/entities.py

@@ -85,11 +85,9 @@ async def set_entities_group(
     await txn.set(key, entities.SerializeToString())
 
 
-async def iterate_entities_groups(
-    txn: Transaction, *, kbid: str
-) -> AsyncGenerator[str, None]:
+async def iterate_entities_groups(txn: Transaction, *, kbid: str) -> AsyncGenerator[str, None]:
     entities_key = KB_ENTITIES.format(kbid=kbid)
-    async for key in txn.keys(entities_key, count=-1):
+    async for key in txn.keys(entities_key):
         group = key.split("/")[-1]
         yield group
 
@@ -106,9 +104,7 @@ async def get_entities_group(
     return eg
 
 
-async def get_deleted_groups(
-    txn: Transaction, *, kbid: str
-) -> kb_pb2.DeletedEntitiesGroups:
+async def get_deleted_groups(txn: Transaction, *, kbid: str) -> kb_pb2.DeletedEntitiesGroups:
     deleted_groups_key = KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid)
     payload = await txn.get(deleted_groups_key)
     deg = kb_pb2.DeletedEntitiesGroups()
@@ -122,18 +118,14 @@ async def mark_group_as_deleted(txn: Transaction, *, kbid: str, group: str) -> N
     deg = await get_deleted_groups(txn, kbid=kbid)
     if group not in deg.entities_groups:
         deg.entities_groups.append(group)
-        await txn.set(
-            KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString()
-        )
+        await txn.set(KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString())
 
 
 async def unmark_group_as_deleted(txn: Transaction, *, kbid: str, group: str) -> None:
     deg = await get_deleted_groups(txn, kbid=kbid)
     if group in deg.entities_groups:
         deg.entities_groups.remove(group)
-        await txn.set(
-            KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString()
-        )
+        await txn.set(KB_DELETED_ENTITIES_GROUPS.format(kbid=kbid), deg.SerializeToString())
 
 
 async def get_entities_meta_cache(txn: Transaction, *, kbid: str) -> EntitiesMetaCache:
@@ -143,7 +135,5 @@ async def get_entities_meta_cache(txn: Transaction, *, kbid: str) -> EntitiesMet
     return pickle.loads(value)
 
 
-async def set_entities_meta_cache(
-    txn: Transaction, kbid: str, cache: EntitiesMetaCache
-) -> None:
+async def set_entities_meta_cache(txn: Transaction, kbid: str, cache: EntitiesMetaCache) -> None:
     await txn.set(KB_ENTITIES_CACHE.format(kbid=kbid), pickle.dumps(cache, protocol=5))