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

nucliadb/common/catalog/utils.py

@@ -0,0 +1,56 @@
+# 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/>.
+#
+from nidx_protos.noderesources_pb2 import Resource as IndexMessage
+
+from nucliadb.common.catalog.interface import CatalogResourceData
+from nucliadb.ingest.orm.resource import Resource
+
+
+def build_catalog_resource_data(resource: Resource, index_message: IndexMessage) -> CatalogResourceData:
+    if resource.basic is None:
+        raise ValueError("Cannot index into the catalog a resource without basic metadata ")
+
+    created_at = resource.basic.created.ToDatetime()
+    modified_at = resource.basic.modified.ToDatetime()
+    if modified_at < created_at:
+        modified_at = created_at
+
+    # Do not index canceled labels
+    cancelled_labels = {
+        f"/l/{clf.labelset}/{clf.label}"
+        for clf in resource.basic.usermetadata.classifications
+        if clf.cancelled_by_user
+    }
+
+    # Labels from the resource and classification labels from each field
+    labels = {label for label in index_message.labels}
+    for classification in resource.basic.computedmetadata.field_classifications:
+        for clf in classification.classifications:
+            label = f"/l/{clf.labelset}/{clf.label}"
+            if label not in cancelled_labels:
+                labels.add(label)
+
+    return CatalogResourceData(
+        title=resource.basic.title,
+        created_at=created_at,
+        modified_at=modified_at,
+        labels=list(labels),
+        slug=resource.basic.slug,
+    )

nucliadb/common/cluster/manager.py

@@ -20,7 +20,8 @@
 import asyncio
 import logging
 import uuid
-from typing import Any, Awaitable, Callable, Optional
+from collections.abc import Awaitable, Callable
+from typing import Any
 
 from nidx_protos import noderesources_pb2, nodewriter_pb2
 from nidx_protos.nodewriter_pb2 import (
@@ -43,8 +44,6 @@ from nucliadb_protos import knowledgebox_pb2, writer_pb2
 from nucliadb_telemetry import errors
 from nucliadb_utils.utilities import get_storage
 
-from .settings import settings
-
 logger = logging.getLogger(__name__)
 
 
@@ -98,7 +97,7 @@ class KBShardManager:
     # TODO: move to data manager
     async def get_current_active_shard(
         self, txn: Transaction, kbid: str
-    ) -> Optional[writer_pb2.ShardObject]:
+    ) -> writer_pb2.ShardObject | None:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
         if kb_shards is None:
             return None
@@ -113,6 +112,8 @@ class KBShardManager:
         self,
         txn: Transaction,
         kbid: str,
+        *,
+        prewarm_enabled: bool,
     ) -> writer_pb2.ShardObject:
         kb_shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
         if kb_shards is None:
@@ -133,6 +134,7 @@ class KBShardManager:
             req = NewShardRequest(
                 kbid=kbid,
                 vectorsets_configs=vectorsets,
+                prewarm_enabled=prewarm_enabled,
             )
 
             resp = await nidx_api.NewShard(req)  # type: ignore
@@ -195,7 +197,7 @@ class KBShardManager:
         txid: int,
         partition: str,
         kb: str,
-        reindex_id: Optional[str] = None,
+        reindex_id: str | None = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
         """
@@ -232,23 +234,6 @@ class KBShardManager:
         indexpb.shard = shard.nidx_shard_id
         await nidx.index(indexpb)
 
-    def should_create_new_shard(self, num_paragraphs: int) -> bool:
-        return num_paragraphs > settings.max_shard_paragraphs
-
-    async def maybe_create_new_shard(
-        self,
-        kbid: str,
-        num_paragraphs: int,
-    ):
-        if not self.should_create_new_shard(num_paragraphs):
-            return
-
-        logger.info({"message": "Adding shard", "kbid": kbid})
-
-        async with datamanagers.with_transaction() as txn:
-            await self.create_shard_by_kbid(txn, kbid)
-            await txn.commit()
-
     async def create_vectorset(self, kbid: str, config: knowledgebox_pb2.VectorSetConfig):
         """Create a new vectorset in all KB shards."""
 
@@ -322,7 +307,7 @@ class StandaloneKBShardManager(KBShardManager):
         txid: int,
         partition: str,
         kb: str,
-        reindex_id: Optional[str] = None,
+        reindex_id: str | None = None,
         source: IndexMessageSource.ValueType = IndexMessageSource.PROCESSOR,
     ) -> None:
         """