nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/common/nidx.py

@@ -0,0 +1,260 @@
+# 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/>.
+#
+
+import os
+from typing import Optional
+
+from nidx_protos.nidx_pb2_grpc import NidxApiStub, NidxSearcherStub
+
+from nucliadb.common.cluster.base import AbstractIndexNode
+from nucliadb.common.cluster.settings import settings
+from nucliadb.ingest.settings import DriverConfig
+from nucliadb.ingest.settings import settings as ingest_settings
+from nucliadb_protos.nodewriter_pb2 import (
+    IndexMessage,
+)
+from nucliadb_utils import logger
+from nucliadb_utils.grpc import get_traced_grpc_channel
+from nucliadb_utils.nats import NatsConnectionManager
+from nucliadb_utils.settings import FileBackendConfig, indexing_settings, storage_settings
+from nucliadb_utils.storages.settings import settings as extended_storage_settings
+from nucliadb_utils.utilities import Utility, clean_utility, get_utility, set_utility
+
+NIDX_ENABLED = bool(os.environ.get("NIDX_ENABLED"))
+
+
+class NidxUtility:
+    api_client = None
+    searcher_client = None
+
+    async def initialize(self):
+        raise NotImplementedError()
+
+    async def finalize(self):
+        raise NotImplementedError()
+
+    async def index(self, msg: IndexMessage) -> int:
+        raise NotImplementedError()
+
+    def wait_for_sync(self):
+        pass
+
+
+def _storage_config(prefix: str, bucket: Optional[str]) -> dict[str, str]:
+    config = {}
+    if storage_settings.file_backend == FileBackendConfig.LOCAL:
+        local_bucket = bucket or storage_settings.local_indexing_bucket
+        file_path = f"{storage_settings.local_files}/{local_bucket}"
+        os.makedirs(file_path, exist_ok=True)
+
+        config[f"{prefix}__OBJECT_STORE"] = "file"
+        config[f"{prefix}__FILE_PATH"] = file_path
+    elif storage_settings.file_backend == FileBackendConfig.GCS:
+        gcs_bucket = bucket or extended_storage_settings.gcs_indexing_bucket
+        config[f"{prefix}__OBJECT_STORE"] = "gcs"
+        if gcs_bucket:
+            config[f"{prefix}__BUCKET"] = gcs_bucket
+        if storage_settings.gcs_base64_creds:
+            config[f"{prefix}__BASE64_CREDS"] = storage_settings.gcs_base64_creds
+        if storage_settings.gcs_endpoint_url:
+            config[f"{prefix}__ENDPOINT"] = storage_settings.gcs_endpoint_url
+    elif storage_settings.file_backend == FileBackendConfig.S3:
+        s3_bucket = bucket or extended_storage_settings.s3_indexing_bucket
+        config[f"{prefix}__OBJECT_STORE"] = "s3"
+        if s3_bucket:
+            config[f"{prefix}__BUCKET"] = s3_bucket
+        config[f"{prefix}__CLIENT_ID"] = storage_settings.s3_client_id or ""
+        config[f"{prefix}__CLIENT_SECRET"] = storage_settings.s3_client_secret or ""
+        config[f"{prefix}__REGION_NAME"] = storage_settings.s3_region_name or ""
+        if storage_settings.s3_endpoint:
+            config[f"{prefix}__ENDPOINT"] = storage_settings.s3_endpoint
+
+    return config
+
+
+class NidxBindingUtility(NidxUtility):
+    """Implements Nidx utility using the binding"""
+
+    def __init__(self):
+        if ingest_settings.driver != DriverConfig.PG:
+            raise ValueError("nidx_binding requires DRIVER=pg")
+
+        self.config = {
+            "METADATA__DATABASE_URL": ingest_settings.driver_pg_url,
+            **_storage_config("INDEXER", None),
+            **_storage_config("STORAGE", "nidx"),
+        }
+
+    async def initialize(self):
+        import nidx_binding  # type: ignore
+
+        self.binding = nidx_binding.NidxBinding(self.config)
+        self.api_client = NidxApiStub(
+            get_traced_grpc_channel(f"localhost:{self.binding.api_port}", "nidx_api")
+        )
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(f"localhost:{self.binding.searcher_port}", "nidx_searcher")
+        )
+
+    async def finalize(self):
+        del self.binding
+
+    async def index(self, msg: IndexMessage) -> int:
+        return self.binding.index(msg.SerializeToString())
+
+    def wait_for_sync(self):
+        self.binding.wait_for_sync()
+
+
+class NidxServiceUtility(NidxUtility):
+    """Implements Nidx utility connecting to the network service"""
+
+    def __init__(self):
+        if indexing_settings.index_nidx_subject is None:
+            raise ValueError("INDEX_NIDX_SUBJECT needed for nidx utility")
+
+        if not settings.nidx_api_address or not settings.nidx_searcher_address:
+            raise ValueError("NIDX_API_ADDRESS and NIDX_SEARCHER_ADDRESS are required")
+
+        self.nats_connection_manager = NatsConnectionManager(
+            service_name="NidxIndexer",
+            nats_servers=indexing_settings.index_jetstream_servers,
+            nats_creds=indexing_settings.index_jetstream_auth,
+        )
+        self.subject = indexing_settings.index_nidx_subject
+
+    async def initialize(self):
+        await self.nats_connection_manager.initialize()
+        self.api_client = NidxApiStub(get_traced_grpc_channel(settings.nidx_api_address, "nidx_api"))
+        self.searcher_client = NidxSearcherStub(
+            get_traced_grpc_channel(settings.nidx_searcher_address, "nidx_searcher")
+        )
+
+    async def finalize(self):
+        await self.nats_connection_manager.finalize()
+
+    async def index(self, writer: IndexMessage) -> int:
+        res = await self.nats_connection_manager.js.publish(self.subject, writer.SerializeToString())
+        logger.info(
+            f" = Pushed message to nidx shard: {writer.shard}, txid: {writer.txid}  seqid: {res.seq}"  # noqa
+        )
+        return res.seq
+
+
+async def start_nidx_utility() -> Optional[NidxUtility]:
+    if not NIDX_ENABLED:
+        return None
+
+    nidx = get_nidx()
+    if nidx:
+        return nidx
+
+    nidx_utility: NidxUtility
+    if settings.standalone_mode:
+        nidx_utility = NidxBindingUtility()
+    else:
+        nidx_utility = NidxServiceUtility()
+
+    await nidx_utility.initialize()
+    set_utility(Utility.NIDX, nidx_utility)
+    return nidx_utility
+
+
+async def stop_nidx_utility():
+    nidx_utility = get_nidx()
+    if nidx_utility:
+        clean_utility(Utility.NIDX)
+        await nidx_utility.finalize()
+
+
+def get_nidx() -> Optional[NidxUtility]:
+    return get_utility(Utility.NIDX)
+
+
+def get_nidx_api_client() -> Optional["NidxApiStub"]:
+    nidx = get_nidx()
+    if nidx:
+        return nidx.api_client
+    else:
+        return None
+
+
+def get_nidx_searcher_client() -> Optional["NidxSearcherStub"]:
+    nidx = get_nidx()
+    if nidx:
+        return nidx.searcher_client
+    else:
+        return None
+
+
+# TODO: Remove the index node abstraction
+class NodeNidxAdapter:
+    def __init__(self, api_client, searcher_client):
+        # API methods
+        self.GetShard = api_client.GetShard
+        self.NewShard = api_client.NewShard
+        self.DeleteShard = api_client.DeleteShard
+        self.ListShards = api_client.ListShards
+        self.AddVectorSet = api_client.AddVectorSet
+        self.RemoveVectorSet = api_client.RemoveVectorSet
+        self.ListVectorSets = api_client.ListVectorSets
+        self.GetMetadata = api_client.GetMetadata
+
+        # Searcher methods
+        self.Search = searcher_client.Search
+        self.Suggest = searcher_client.Suggest
+        self.Paragraphs = searcher_client.Paragraphs
+        self.Documents = searcher_client.Documents
+
+
+class FakeNode(AbstractIndexNode):
+    def __init__(self, api_client, searcher_client):
+        self.client = NodeNidxAdapter(api_client, searcher_client)
+
+    @property
+    def reader(self):
+        return self.client
+
+    @property
+    def writer(self):
+        return self.client
+
+    def is_read_replica(_):
+        return False
+
+    @property
+    def id(self):
+        return "nidx"
+
+    @property
+    def address(self):
+        return "nidx"
+
+    @property
+    def primary_id(self):
+        return "nidx"
+
+
+def get_nidx_fake_node() -> Optional[FakeNode]:
+    nidx = get_nidx()
+    if nidx:
+        return FakeNode(nidx.api_client, nidx.searcher_client)
+    else:
+        return None

nucliadb/export_import/datamanager.py

@@ -17,6 +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/>.
 #
+import json
 from datetime import datetime
 from typing import AsyncGenerator, Union
 
@@ -53,25 +54,38 @@ class ExportImportDataManager:
 
     async def get_metadata(self, type: str, kbid: str, id: str) -> Metadata:
         key = self._get_maindb_metadata_key(type, kbid, id)
-        async with self.driver.transaction() as txn:
+        async with self.driver.transaction(read_only=True) as txn:
             data = await txn.get(key)
         if data is None or data == b"":
             raise MetadataNotFound()
         decoded = data.decode("utf-8")
-        model_type = {
-            "export": ExportMetadata,
-            "import": ImportMetadata,
-        }[type]
-        return model_type.parse_raw(decoded)  # type: ignore
+        if type == "export":
+            model_type = ExportMetadata
+        elif type == "import":
+            model_type = ImportMetadata  # type: ignore
+        else:
+            raise ValueError(f"Invalid type: {type}")
+        json_decoded = json.loads(decoded)
+
+        # For some reason, the total and processed fields are not always present in the metadata.
+        # This is to unblock already created exports that hit this bug.
+        if json_decoded.get("total") is None:
+            json_decoded["total"] = 0
+        if json_decoded.get("processed") is None:
+            json_decoded["processed"] = 0
+
+        return model_type.model_validate(json_decoded)
 
     async def set_metadata(
         self,
         type: str,
         metadata: Metadata,
     ):
+        metadata.processed = metadata.processed or 0
+        metadata.total = metadata.total or 0
         metadata.modified = datetime.utcnow()
         key = self._get_maindb_metadata_key(type, metadata.kbid, metadata.id)
-        data = metadata.json().encode("utf-8")
+        data = metadata.model_dump_json().encode("utf-8")
         async with self.driver.transaction() as txn:
             await txn.set(key, data)
             await txn.commit()
@@ -97,9 +111,7 @@ class ExportImportDataManager:
         await self.storage.uploaditerator(export_bytes, field, cf)
         return cf.size
 
-    async def download_export(
-        self, kbid: str, export_id: str
-    ) -> AsyncGenerator[bytes, None]:
+    async def download_export(self, kbid: str, export_id: str) -> AsyncGenerator[bytes, None]:
         key = STORAGE_EXPORT_KEY.format(export_id=export_id)
         bucket = self.storage.get_bucket_name(kbid)
         async for chunk in self.storage.download(bucket, key):
@@ -125,13 +137,9 @@ class ExportImportDataManager:
         async for chunk in self.storage.download(bucket, key):
             yield chunk
 
-    def _get_storage_field(
-        self, kbid: str, key: str, cf: resources_pb2.CloudFile
-    ) -> StorageField:
+    def _get_storage_field(self, kbid: str, key: str, cf: resources_pb2.CloudFile) -> StorageField:
         bucket = self.storage.get_bucket_name(kbid)
-        return self.storage.field_klass(
-            storage=self.storage, bucket=bucket, fullkey=key, field=cf
-        )
+        return self.storage.field_klass(storage=self.storage, bucket=bucket, fullkey=key, field=cf)
 
     async def delete_import(self, kbid: str, import_id: str):
         key = STORAGE_IMPORT_KEY.format(import_id=import_id)
@@ -151,6 +159,4 @@ class ExportImportDataManager:
             await func(kbid, id)
         except Exception as ex:
             errors.capture_exception(ex)
-            logger.exception(
-                f"Could not delete {type} {id} from storage", extra={"kbid": kbid}
-            )
+            logger.exception(f"Could not delete {type} {id} from storage", extra={"kbid": kbid})

nucliadb/export_import/exceptions.py

@@ -70,3 +70,11 @@ class WrongExportStreamFormat(Exception):
     """
 
     pass
+
+
+class IncompatibleExport(Exception):
+    """
+    Raised when trying to import an export file that is incompatible with the destination knowledgebox.
+    """
+
+    pass

nucliadb/export_import/exporter.py

@@ -35,6 +35,7 @@ from nucliadb.export_import.utils import (
     get_cloud_files,
     get_entities,
     get_labels,
+    get_learning_config,
     iter_kb_resource_uuids,
 )
 from nucliadb_protos import writer_pb2
@@ -51,6 +52,9 @@ async def export_kb(
 
     If a metadata object is provided, uses it to resume the export if it was interrupted.
     """
+    async for chunk in export_learning_config(kbid):
+        yield chunk
+
     resources_iterator = export_resources(context, kbid)
     if metadata is not None:
         assert metadata.kbid == kbid
@@ -66,9 +70,7 @@ async def export_kb(
         yield chunk
 
 
-async def export_kb_to_blob_storage(
-    context: ApplicationContext, msg: NatsTaskMessage
-) -> None:
+async def export_kb_to_blob_storage(context: ApplicationContext, msg: NatsTaskMessage) -> None:
     """
     Exports the data of a knowledgebox to the blob storage service.
     """
@@ -86,7 +88,7 @@ async def export_kb_to_blob_storage(
     export_size = await upload_export_retried(iterator, kbid, export_id)
 
     # Store export size
-    metadata.total = metadata.processed = export_size
+    metadata.total = metadata.processed = export_size or 0
     await dm.set_metadata("export", metadata)
 
 
@@ -103,9 +105,7 @@ async def export_resources(
             yield chunk
 
 
-async def export_resources_resumable(
-    context, metadata: ExportMetadata
-) -> AsyncGenerator[bytes, None]:
+async def export_resources_resumable(context, metadata: ExportMetadata) -> AsyncGenerator[bytes, None]:
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
 
     kbid = metadata.kbid
@@ -189,3 +189,16 @@ async def export_labels(
         yield ExportedItemType.LABELS.encode("utf-8")
         yield len(data).to_bytes(4, byteorder="big")
         yield data
+
+
+async def export_learning_config(
+    kbid: str,
+) -> AsyncGenerator[bytes, None]:
+    lconfig = await get_learning_config(kbid)
+    if lconfig is None:
+        logger.warning(f"No learning configuration found for kbid", extra={"kbid": kbid})
+        return
+    data = lconfig.model_dump_json().encode("utf-8")
+    yield ExportedItemType.LEARNING_CONFIG.encode("utf-8")
+    yield len(data).to_bytes(4, byteorder="big")
+    yield data

nucliadb/export_import/importer.py

@@ -28,9 +28,7 @@ from nucliadb.export_import.models import (
     NatsTaskMessage,
 )
 from nucliadb.export_import.utils import (
-    ExportStream,
     ExportStreamReader,
-    IteratorExportStream,
     TaskRetryHandler,
     import_binary,
     import_broker_message,
@@ -47,7 +45,7 @@ BinaryStreamGenerator = Callable[[int], BinaryStream]
 async def import_kb(
     context: ApplicationContext,
     kbid: str,
-    stream: ExportStream,
+    stream: AsyncGenerator[bytes, None],
     metadata: Optional[ImportMetadata] = None,
 ) -> None:
     """
@@ -91,28 +89,25 @@ async def import_kb(
             await dm.set_metadata("import", metadata)
 
     if metadata is not None:
-        metadata.processed = stream_reader.read_bytes
+        metadata.processed = stream_reader.read_bytes or 0
         await dm.set_metadata("import", metadata)
 
 
-async def import_kb_from_blob_storage(
-    context: ApplicationContext, msg: NatsTaskMessage
-):
+async def import_kb_from_blob_storage(context: ApplicationContext, msg: NatsTaskMessage):
     """
     Imports to a knowledgebox from an export stored in the blob storage service.
     """
     kbid, import_id = msg.kbid, msg.id
     dm = ExportImportDataManager(context.kv_driver, context.blob_storage)
     metadata = await dm.get_metadata(type="import", kbid=kbid, id=import_id)
-    iterator = dm.download_import(kbid, import_id)
-    stream = IteratorExportStream(iterator)
 
     retry_handler = TaskRetryHandler("import", dm, metadata)
 
     @retry_handler.wrap
-    async def import_kb_retried(context, kbid, stream, metadata):
+    async def import_kb_retried(context, kbid, metadata):
+        stream = dm.download_import(kbid, import_id)
         await import_kb(context, kbid, stream, metadata)
 
-    await import_kb_retried(context, kbid, stream, metadata)  # type: ignore
+    await import_kb_retried(context, kbid, metadata)
 
     await dm.try_delete_from_storage("import", kbid, import_id)

nucliadb/export_import/models.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 datetime import datetime
+import datetime
 from enum import Enum
 from typing import Any
 
@@ -36,6 +36,7 @@ class ExportedItemType(str, Enum):
     LABELS = "LAB"
     ENTITIES = "ENT"
     BINARY = "BIN"
+    LEARNING_CONFIG = "LEA"
 
 
 ExportItem = tuple[ExportedItemType, Any]
@@ -56,8 +57,8 @@ class Metadata(BaseModel):
     task: TaskMetadata = TaskMetadata(status=Status.SCHEDULED)
     total: int = 0
     processed: int = 0
-    created: datetime = datetime.utcnow()
-    modified: datetime = datetime.utcnow()
+    created: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
+    modified: datetime.datetime = datetime.datetime.now(datetime.timezone.utc)
 
 
 class ExportMetadata(Metadata):
@@ -65,8 +66,7 @@ class ExportMetadata(Metadata):
     exported_resources: list[str] = list()
 
 
-class ImportMetadata(Metadata):
-    ...
+class ImportMetadata(Metadata): ...
 
 
 class NatsTaskMessage(BaseModel):

nucliadb/export_import/tasks.py

@@ -32,7 +32,7 @@ def get_exports_consumer() -> NatsTaskConsumer:
         name="exports_consumer",
         stream=const.Streams.KB_EXPORTS,  # type: ignore
         callback=export_kb_to_blob_storage,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
@@ -41,7 +41,7 @@ async def get_exports_producer(context: ApplicationContext) -> NatsTaskProducer:
     producer = create_producer(
         name="exports_producer",
         stream=const.Streams.KB_EXPORTS,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)
     return producer
@@ -52,7 +52,7 @@ def get_imports_consumer() -> NatsTaskConsumer:
         name="imports_consumer",
         stream=const.Streams.KB_IMPORTS,  # type: ignore
         callback=import_kb_from_blob_storage,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
         max_concurrent_messages=10,
     )
 
@@ -61,7 +61,7 @@ async def get_imports_producer(context: ApplicationContext) -> NatsTaskProducer:
     producer = create_producer(
         name="imports_producer",
         stream=const.Streams.KB_IMPORTS,  # type: ignore
-        msg_type=NatsTaskMessage,  # type: ignore
+        msg_type=NatsTaskMessage,
     )
     await producer.initialize(context)
     return producer