nucliadb 6.2.0.post2679__py3-none-any.whl → 6.2.1__py3-none-any.whl

nucliadb/common/cluster/standalone/utils.py

@@ -19,13 +19,10 @@
 
 import logging
 import os
-import shutil
 import uuid
-from socket import gethostname
 
 from nucliadb.common.cluster.settings import StandaloneNodeRole
 from nucliadb.common.cluster.settings import settings as cluster_settings
-from nucliadb.common.cluster.standalone.index_node import StandaloneIndexNode
 
 logger = logging.getLogger(__name__)
 
@@ -46,46 +43,6 @@ def get_standalone_node_id() -> str:
         return str(uuid.UUID(bytes=f.read()))
 
 
-_SELF_INDEX_NODE = None
-
-
-def get_self() -> StandaloneIndexNode:
-    """
-    This returns an instance of the standalone index node
-    so when API requests come into this mode, we don't
-    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.")
-    global _SELF_INDEX_NODE
-    node_id = get_standalone_node_id()
-    if _SELF_INDEX_NODE is None or node_id != _SELF_INDEX_NODE.id:
-        if "NUCLIADB_SERVICE_HOST" in os.environ:
-            hn = os.environ["HOSTNAME"]
-            ns = os.environ.get("NAMESPACE", "nucliadb")
-            host = f"{hn}.{ns}"
-        else:
-            host = gethostname()
-        _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
-    except FileNotFoundError:  # pragma: no cover
-        ...
-    try:
-        _shards_dir = os.path.join(cluster_settings.data_path, "shards")
-        _SELF_INDEX_NODE.shard_count = len(
-            [
-                shard_dir
-                for shard_dir in os.listdir(_shards_dir)
-                if os.path.isdir(os.path.join(_shards_dir, shard_dir))
-            ]
-        )
-    except FileNotFoundError:  # pragma: no cover
-        ...
-    return _SELF_INDEX_NODE
-
-
 def is_index_node() -> bool:
     return cluster_settings.standalone_node_role in (
         StandaloneNodeRole.ALL,

nucliadb/common/cluster/utils.py

@@ -23,20 +23,11 @@ from typing import TYPE_CHECKING, Optional, Union
 import backoff
 
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.discovery.utils import (
-    setup_cluster_discovery,
-    teardown_cluster_discovery,
-)
 from nucliadb.common.cluster.manager import (
     KBShardManager,
     StandaloneKBShardManager,
-    clear_index_nodes,
 )
 from nucliadb.common.cluster.settings import settings
-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.ingest.orm.resource import Resource
 from nucliadb_protos import nodereader_pb2, writer_pb2
 from nucliadb_utils import const
@@ -62,12 +53,8 @@ async def setup_cluster() -> Union[KBShardManager, StandaloneKBShardManager]:
             # already setup
             return get_utility(Utility.SHARD_MANAGER)
 
-        await setup_cluster_discovery()
         mng: Union[KBShardManager, StandaloneKBShardManager]
         if settings.standalone_mode:
-            if is_index_node():
-                server = await start_standalone_grpc()
-                set_utility(_STANDALONE_SERVER, server)
             mng = StandaloneKBShardManager()
         else:
             mng = KBShardManager()
@@ -76,7 +63,6 @@ async def setup_cluster() -> Union[KBShardManager, StandaloneKBShardManager]:
 
 
 async def teardown_cluster():
-    await teardown_cluster_discovery()
     if get_utility(Utility.SHARD_MANAGER):
         clean_utility(Utility.SHARD_MANAGER)
 
@@ -85,8 +71,6 @@ async def teardown_cluster():
         await std_server.stop(None)
         clean_utility(_STANDALONE_SERVER)
 
-    clear_index_nodes()
-
 
 def get_shard_manager() -> KBShardManager:
     return get_utility(Utility.SHARD_MANAGER)  # type: ignore

nucliadb/common/counters.py

@@ -26,3 +26,4 @@ class IndexCounts:
     fields: int
     paragraphs: int
     sentences: int
+    size_bytes: int

nucliadb/common/datamanagers/fields.py

@@ -23,11 +23,13 @@ from typing import Optional
 from google.protobuf.message import Message
 
 from nucliadb.common.datamanagers.utils import get_kv_pb
+from nucliadb.common.ids import FIELD_TYPE_PB_TO_STR
 from nucliadb.common.maindb.driver import Transaction
 from nucliadb_protos import writer_pb2
 
 KB_RESOURCE_FIELD = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}"
 KB_RESOURCE_FIELD_ERROR = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/error"
+KB_RESOURCE_FIELD_STATUS = "/kbs/{kbid}/r/{uuid}/f/{type}/{field}/status"
 
 
 async def get_raw(
@@ -52,13 +54,7 @@ async def set(
 
 async def delete(txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str):
     base_key = KB_RESOURCE_FIELD.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
-    # Make sure we explicitly delete the field and any nested key
-    keys_to_delete = []
-    async for key in txn.keys(base_key):
-        keys_to_delete.append(key)
-
-    for key in keys_to_delete:
-        await txn.delete(key)
+    await txn.delete_by_prefix(base_key)
 
 
 # Error
@@ -82,3 +78,48 @@ async def set_error(
 ):
     key = KB_RESOURCE_FIELD_ERROR.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
     await txn.set(key, error.SerializeToString())
+
+
+# Status, replaces error
+
+
+async def get_status(
+    txn: Transaction, *, kbid: str, rid: str, field_type: str, field_id: str
+) -> Optional[writer_pb2.FieldStatus]:
+    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)
+
+
+async def get_statuses(
+    txn: Transaction, *, kbid: str, rid: str, fields: list[writer_pb2.FieldID]
+) -> list[writer_pb2.FieldStatus]:
+    keys = [
+        KB_RESOURCE_FIELD_STATUS.format(
+            kbid=kbid, uuid=rid, type=FIELD_TYPE_PB_TO_STR[fid.field_type], field=fid.field
+        )
+        for fid in fields
+    ]
+    serialized = await txn.batch_get(keys, for_update=False)
+    statuses = []
+    for serialized_status in serialized:
+        pb = writer_pb2.FieldStatus()
+        if serialized_status is not None:
+            pb.ParseFromString(serialized_status)
+        else:
+            pb = writer_pb2.FieldStatus()
+        statuses.append(pb)
+
+    return statuses
+
+
+async def set_status(
+    txn: Transaction,
+    *,
+    kbid: str,
+    rid: str,
+    field_type: str,
+    field_id: str,
+    status: writer_pb2.FieldStatus,
+):
+    key = KB_RESOURCE_FIELD_STATUS.format(kbid=kbid, uuid=rid, type=field_type, field=field_id)
+    await txn.set(key, status.SerializeToString())

nucliadb/common/datamanagers/vectorsets.py

@@ -58,6 +58,11 @@ async def iter(
         yield config.vectorset_id, config
 
 
+async def count(txn: Transaction, *, kbid: str) -> int:
+    kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=False)
+    return len(kb_vectorsets.vectorsets)
+
+
 async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
     """Create or update a vectorset configuration"""
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
@@ -73,16 +78,20 @@ async def set(txn: Transaction, *, kbid: str, config: knowledgebox_pb2.VectorSet
     await txn.set(key, kb_vectorsets.SerializeToString())
 
 
-async def delete(txn: Transaction, *, kbid: str, vectorset_id: str):
+async def delete(
+    txn: Transaction, *, kbid: str, vectorset_id: str
+) -> Optional[knowledgebox_pb2.VectorSetConfig]:
     kb_vectorsets = await _get_or_default(txn, kbid=kbid, for_update=True)
     index = _find_vectorset(kb_vectorsets, vectorset_id)
     if index is None:
         # already deleted
-        return
+        return None
 
+    deleted = kb_vectorsets.vectorsets[index]
     del kb_vectorsets.vectorsets[index]
     key = KB_VECTORSETS.format(kbid=kbid)
     await txn.set(key, kb_vectorsets.SerializeToString())
+    return deleted
 
 
 # XXX At some point in the vectorset epic, we should make this key mandatory and

nucliadb/common/external_index_providers/base.py

@@ -28,7 +28,7 @@ from nucliadb.common.counters import IndexCounts
 from nucliadb.common.external_index_providers.exceptions import ExternalIndexingError
 from nucliadb.common.ids import ParagraphId
 from nucliadb_models.external_index_providers import ExternalIndexProviderType
-from nucliadb_models.search import SCORE_TYPE, TextPosition
+from nucliadb_models.search import SCORE_TYPE, Relations, TextPosition
 from nucliadb_protos.knowledgebox_pb2 import (
     CreateExternalIndexProviderMetadata,
     StoredExternalIndexProviderMetadata,
@@ -73,6 +73,7 @@ class TextBlockMatch(BaseModel):
     paragraph_labels: list[str] = []
     field_labels: list[str] = []
     text: Optional[str] = None
+    relevant_relations: Optional[Relations] = None
 
 
 class QueryResults(BaseModel):

nucliadb/common/external_index_providers/pinecone.py

@@ -441,6 +441,7 @@ class PineconeIndexManager(ExternalIndexManager):
 
     def get_prefixes_to_delete(self, index_data: Resource) -> set[str]:
         prefixes_to_delete = set()
+        # TODO: migrate to vector_prefixes_to_delete
         for field_id in index_data.sentences_to_delete:
             try:
                 delete_vid = VectorId.from_string(field_id)
@@ -706,11 +707,7 @@ class PineconeIndexManager(ExternalIndexManager):
         if self.kbid in COUNTERS_CACHE:
             # Cache hit
             return COUNTERS_CACHE[self.kbid]
-        total = IndexCounts(
-            fields=0,
-            paragraphs=0,
-            sentences=0,
-        )
+        total = IndexCounts(fields=0, paragraphs=0, sentences=0, size_bytes=0)
         tasks = []
         vectorset_results: dict[str, IndexCounts] = {}
 
@@ -738,6 +735,7 @@ class PineconeIndexManager(ExternalIndexManager):
                 fields=0,
                 paragraphs=index_stats.totalVectorCount,
                 sentences=index_stats.totalVectorCount,
+                size_bytes=0,
             )
         except Exception:
             logger.exception(

nucliadb/common/ids.py

@@ -111,13 +111,11 @@ class FieldId:
         parts = value.split("/")
         if len(parts) == 3:
             rid, _type, key = parts
-            if _type not in FIELD_TYPE_STR_TO_PB:
-                raise ValueError(f"Invalid FieldId: {value}")
+            _type = cls.parse_field_type(_type)
             return cls(rid=rid, type=_type, key=key)
         elif len(parts) == 4:
             rid, _type, key, subfield_id = parts
-            if _type not in FIELD_TYPE_STR_TO_PB:
-                raise ValueError(f"Invalid FieldId: {value}")
+            _type = cls.parse_field_type(_type)
             return cls(
                 rid=rid,
                 type=_type,
@@ -127,6 +125,22 @@ class FieldId:
         else:
             raise ValueError(f"Invalid FieldId: {value}")
 
+    @classmethod
+    def parse_field_type(cls, _type: str) -> str:
+        if _type not in FIELD_TYPE_STR_TO_PB:
+            # Try to parse the enum value
+            # XXX: This is to support field types that are integer values of FieldType
+            # Which is how legacy processor relations reported the paragraph_id
+            try:
+                type_pb = FieldType.ValueType(int(_type))
+            except ValueError:
+                raise ValueError(f"Invalid FieldId: {_type}")
+            if type_pb in FIELD_TYPE_PB_TO_STR:
+                return FIELD_TYPE_PB_TO_STR[type_pb]
+            else:
+                raise ValueError(f"Invalid FieldId: {_type}")
+        return _type
+
 
 @dataclass
 class ParagraphId: