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

migrations/0023_backfill_pg_catalog.py

@@ -27,10 +27,11 @@ Backfill the data into the PG catalog
 import logging
 from typing import cast
 
-from nucliadb.common import datamanagers
+from nucliadb.common.catalog import catalog_update, get_catalog
+from nucliadb.common.catalog.pg import PGCatalog
 from nucliadb.common.maindb.pg import PGDriver, PGTransaction
 from nucliadb.ingest.orm.index_message import get_resource_index_message
-from nucliadb.ingest.orm.processor.pgcatalog import pgcatalog_update
+from nucliadb.ingest.orm.resource import Resource
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -43,6 +44,9 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     if not isinstance(context.kv_driver, PGDriver):
         return
 
+    if not isinstance(get_catalog(), PGCatalog):
+        return
+
     BATCH_SIZE = 100
     async with context.kv_driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
@@ -69,13 +73,13 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
                 # Index each resource
                 for rid in resources_to_index:
                     rid = str(rid).replace("-", "")
-                    resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=rid)
+                    resource = await Resource.get(txn, kbid=kbid, rid=rid)
                     if resource is None:
                         logger.warning(f"Could not load resource {rid} for kbid {kbid}")
                         continue
 
                     index_message = await get_resource_index_message(resource, reindex=False)
-                    await pgcatalog_update(txn, kbid, resource, index_message)
+                    await catalog_update(txn, kbid, resource, index_message)
 
                 await txn.commit()
                 continue_sql = f"AND key > '/kbs/{kbid}/r/{rid}'"

migrations/0028_extracted_vectors_reference.py

@@ -39,7 +39,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     async with datamanagers.with_rw_transaction() as txn:
         vectorsets = [vs async for (_vid, vs) in datamanagers.vectorsets.iter(txn, kbid=kbid)]
 
-        if len(vectorsets) == 0:  # pragma: nocover
+        if len(vectorsets) == 0:  # pragma: no cover
             # should never happen, everyone should have at least one
             logger.warning(f"KB has no vectorsets!", extra={"kbid": kbid})
             return

migrations/0029_backfill_field_status.py

@@ -24,7 +24,6 @@ Backfill field status (from error)
 """
 
 import logging
-from typing import Optional
 
 from nucliadb.migrator.context import ExecutionContext
 from nucliadb_protos import resources_pb2, writer_pb2
@@ -33,7 +32,7 @@ logger = logging.getLogger(__name__)
 
 
 async def migrate(context: ExecutionContext) -> None:
-    start: Optional[str] = ""
+    start: str | None = ""
     while True:
         if start is None:
             break
@@ -43,7 +42,7 @@ async def migrate(context: ExecutionContext) -> None:
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None: ...
 
 
-async def do_batch(context: ExecutionContext, start: str) -> Optional[str]:
+async def do_batch(context: ExecutionContext, start: str) -> str | None:
     logger.info(f"Running batch from {start}")
     async with context.kv_driver.rw_transaction() as txn:
         async with txn.connection.cursor() as cur:  # type: ignore
@@ -64,7 +63,7 @@ async def do_batch(context: ExecutionContext, start: str) -> Optional[str]:
             field_keys = [r[0] for r in records]
 
             # Retrieve resources basic (to check status)
-            resource_keys = set(["/".join(f.split("/")[:5]) for f in field_keys])
+            resource_keys = {"/".join(f.split("/")[:5]) for f in field_keys}
             await cur.execute(
                 """
                 SELECT key, value FROM resources

migrations/0032_remove_old_relations.py

@@ -26,7 +26,6 @@ is stored in object storage.
 """
 
 import logging
-from typing import Optional
 
 from nucliadb.migrator.context import ExecutionContext
 
@@ -34,7 +33,7 @@ logger = logging.getLogger(__name__)
 
 
 async def migrate(context: ExecutionContext) -> None:
-    start: Optional[str] = ""
+    start: str | None = ""
     while True:
         if start is None:
             break
@@ -45,7 +44,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     pass
 
 
-async def do_batch(context: ExecutionContext, start: str) -> Optional[str]:
+async def do_batch(context: ExecutionContext, start: str) -> str | None:
     logger.info(f"Running batch from {start}")
     async with context.kv_driver.rw_transaction() as txn:
         async with txn.connection.cursor() as cur:  # type: ignore

migrations/0038_backfill_catalog_field_labels.py

@@ -27,10 +27,11 @@ Backfill the catalog with labels from fields metadata
 import logging
 from typing import cast
 
-from nucliadb.common import datamanagers
+from nucliadb.common.catalog import catalog_update, get_catalog
+from nucliadb.common.catalog.pg import PGCatalog
 from nucliadb.common.maindb.pg import PGDriver, PGTransaction
 from nucliadb.ingest.orm.index_message import get_resource_index_message
-from nucliadb.ingest.orm.processor.pgcatalog import pgcatalog_update
+from nucliadb.ingest.orm.resource import Resource
 from nucliadb.migrator.context import ExecutionContext
 from nucliadb_protos import resources_pb2
 
@@ -44,6 +45,9 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     if not isinstance(context.kv_driver, PGDriver):
         return
 
+    if not isinstance(get_catalog(), PGCatalog):
+        return
+
     BATCH_SIZE = 100
     async with context.kv_driver.rw_transaction() as txn:
         txn = cast(PGTransaction, txn)
@@ -78,13 +82,13 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
                 # Index each resource
                 for key in to_index:
                     rid = key.split("/")[4]
-                    resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=rid)
+                    resource = await Resource.get(txn, kbid=kbid, rid=rid)
                     if resource is None:
                         logger.warning(f"Could not load resource {rid} for kbid {kbid}")
                         continue
 
                     index_message = await get_resource_index_message(resource, reindex=False)
-                    await pgcatalog_update(txn, kbid, resource, index_message)
+                    await catalog_update(txn, kbid, resource, index_message)
 
                 if to_index:
                     await txn.commit()

migrations/0039_backfill_converation_splits_metadata.py

@@ -0,0 +1,106 @@
+# 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/>.
+#
+
+"""Migration #39
+
+Backfill splits metadata on conversation fields
+
+"""
+
+import logging
+from typing import cast
+
+from nucliadb.common.maindb.driver import Transaction
+from nucliadb.common.maindb.pg import PGTransaction
+from nucliadb.ingest.fields.conversation import (
+    CONVERSATION_SPLITS_METADATA,
+    Conversation,
+)
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox as KnowledgeBoxORM
+from nucliadb.migrator.context import ExecutionContext
+from nucliadb_protos import resources_pb2
+from nucliadb_protos.resources_pb2 import SplitsMetadata
+from nucliadb_utils.storages.storage import Storage
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    BATCH_SIZE = 100
+    start = ""
+    while True:
+        to_fix: list[tuple[str, str]] = []
+        async with context.kv_driver.rw_transaction() as txn:
+            txn = cast(PGTransaction, txn)
+            async with txn.connection.cursor() as cur:
+                # Retrieve a bunch of conversation fields
+                await cur.execute(
+                    """
+                    SELECT key FROM resources
+                    WHERE key ~ ('^/kbs/' || %s || '/r/[^/]*/f/c/[^/]*$')
+                    AND key > %s
+                    ORDER BY key
+                    LIMIT %s""",
+                    (kbid, start, BATCH_SIZE),
+                )
+                rows = await cur.fetchall()
+                if len(rows) == 0:
+                    return
+                for row in rows:
+                    key = row[0]
+                    start = key
+                    rid = key.split("/")[4]
+                    field_id = key.split("/")[7]
+                    to_fix.append((rid, field_id))
+
+        for rid, field_id in to_fix:
+            async with context.kv_driver.rw_transaction() as txn2:
+                splits_metadata = await build_splits_metadata(
+                    txn2, context.blob_storage, kbid, rid, field_id
+                )
+                splits_metadata_key = CONVERSATION_SPLITS_METADATA.format(
+                    kbid=kbid, uuid=rid, type="c", field=field_id
+                )
+                await txn2.set(splits_metadata_key, splits_metadata.SerializeToString())
+                await txn2.commit()
+
+
+async def build_splits_metadata(
+    txn: Transaction, storage: Storage, kbid: str, rid: str, field_id: str
+) -> SplitsMetadata:
+    splits_metadata = SplitsMetadata()
+    kb_orm = KnowledgeBoxORM(txn, storage, kbid)
+    resource_obj = await kb_orm.get(rid)
+    if resource_obj is None:
+        return splits_metadata
+    field_obj: Conversation = await resource_obj.get_field(
+        field_id, resources_pb2.FieldType.CONVERSATION, load=False
+    )
+    conv_metadata = await field_obj.get_metadata()
+    for i in range(1, conv_metadata.pages + 1):
+        page = await field_obj.get_value(page=i)
+        if page is None:
+            continue
+        for message in page.messages:
+            splits_metadata.metadata.get_or_create(message.ident)
+    return splits_metadata

migrations/0040_migrate_search_configurations.py

@@ -0,0 +1,79 @@
+# 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/>.
+#
+
+"""Migration #40
+
+Replaces deprecated and removed generative models from search configurations
+
+"""
+
+import logging
+from typing import cast
+
+from nucliadb.common import datamanagers
+from nucliadb.migrator.context import ExecutionContext
+from nucliadb_models.configuration import SearchConfiguration
+from nucliadb_models.search import AskRequest, FindRequest
+
+logger = logging.getLogger(__name__)
+
+REPLACEMENTS = {
+    "claude-3-5-small": "claude-4-5-sonnet",
+    "gcp-claude-3-5-sonnet-v2": "gcp-claude-4-5-sonnet",
+}
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    affected = await get_affected_search_configurations(kbid)
+    if not affected:
+        return
+
+    async with datamanagers.with_rw_transaction() as txn:
+        for name, config in affected.items():
+            logger.info(
+                "Migrating search config for kb",
+                extra={
+                    "kbid": kbid,
+                    "search_config": name,
+                    "generative_model": config.config.generative_model,  # type: ignore
+                },
+            )
+            config.config.generative_model = REPLACEMENTS[config.config.generative_model]  # type: ignore
+            await datamanagers.search_configurations.set(txn, kbid=kbid, name=name, config=config)
+        await txn.commit()
+
+
+async def get_affected_search_configurations(kbid: str) -> dict[str, SearchConfiguration]:
+    result: dict[str, SearchConfiguration] = {}
+    async with datamanagers.with_ro_transaction() as txn:
+        search_configs = await datamanagers.search_configurations.list(txn, kbid=kbid)
+        for name, config in search_configs.items():
+            if config.kind == "find":
+                find_config = cast(FindRequest, config.config)
+                if find_config.generative_model in REPLACEMENTS:
+                    result[name] = config
+            elif config.kind == "ask":
+                ask_config = cast(AskRequest, config.config)
+                if ask_config.generative_model in REPLACEMENTS:
+                    result[name] = config
+    return result

migrations/0041_reindex_conversations.py

@@ -0,0 +1,137 @@
+# 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 logging
+import uuid
+from collections.abc import AsyncIterator
+from typing import cast
+
+from nucliadb.common import datamanagers
+from nucliadb.common.maindb.pg import PGTransaction
+from nucliadb.ingest.orm.index_message import get_resource_index_message
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.migrator.context import ExecutionContext
+from nucliadb_protos.writer_pb2 import ShardObject, Shards
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    """
+    Reindex resources that have conversation fields
+    """
+    kb_shards = await datamanagers.atomic.cluster.get_kb_shards(kbid=kbid, for_update=False)
+    if kb_shards is not None:
+        async for rid in iter_affected_resource_ids(context, kbid):
+            await reindex_resource(context, kbid, rid, kb_shards)
+    else:
+        logger.warning(
+            "Migration 41: KB shards not found, skipping reindexing",
+            extra={"kbid": kbid},
+        )
+
+
+async def reindex_resource(
+    context: ExecutionContext,
+    kbid: str,
+    rid: str,
+    kb_shards: Shards,
+) -> None:
+    """
+    Reindex a single resource
+    """
+    async with datamanagers.with_ro_transaction() as rs_txn:
+        # Fetch the resource
+        resource = await Resource.get(rs_txn, kbid=kbid, rid=rid)
+        if resource is None:
+            logger.warning(
+                "Migration 41: Resource not found, skipping reindexing",
+                extra={"kbid": kbid, "rid": rid},
+            )
+            return
+
+        # Get the shard for the resource
+        shard: ShardObject | None = None
+        shard_id = await datamanagers.resources.get_resource_shard_id(
+            rs_txn, kbid=kbid, rid=rid, for_update=False
+        )
+        if shard_id is not None:
+            shard = next((shard for shard in kb_shards.shards if shard.shard == shard_id), None)
+        if shard is None:
+            logger.warning(
+                "Migration 41: Shard not found for resource, skipping reindexing",
+                extra={"kbid": kbid, "rid": rid, "shard_id": shard_id},
+            )
+            return
+
+        # Create the index message and reindex the resource
+        index_message = await get_resource_index_message(resource, reindex=True)
+        await context.shard_manager.add_resource(
+            shard,
+            index_message,
+            0,
+            partition="0",
+            kb=kbid,
+            reindex_id=uuid.uuid4().hex,
+        )
+        logger.info(
+            "Migration 41: Resource reindexed",
+            extra={"kbid": kbid, "rid": rid},
+        )
+
+
+async def iter_affected_resource_ids(context: ExecutionContext, kbid: str) -> AsyncIterator[str]:
+    start = ""
+    while True:
+        keys_batch = await get_batch(context, kbid, start)
+        if keys_batch is None:
+            break
+        start = keys_batch[-1]
+        for key in keys_batch:
+            # The keys have the format /kbs/{kbid}/r/{rid}/f/c/{field_id}
+            rid = key.split("/")[4]
+            yield rid
+
+
+async def get_batch(context: ExecutionContext, kbid: str, start: str) -> list[str] | None:
+    """
+    Get a batch of resource keys that hold conversation fields for the given KB.
+    Starting after the given start key.
+    Returns None if no more keys are found.
+    """
+    batch_size = 100
+    async with context.kv_driver.rw_transaction() as txn:
+        txn = cast(PGTransaction, txn)
+        async with txn.connection.cursor() as cur:
+            await cur.execute(
+                """
+                SELECT key FROM resources
+                WHERE key ~ ('^/kbs/' || %s || '/r/[^/]*/f/c/[^/]*$')
+                AND key > %s
+                ORDER BY key
+                LIMIT %s""",
+                (kbid, start, batch_size),
+            )
+            rows = await cur.fetchall()
+            if len(rows) == 0:
+                return None
+            return [row[0] for row in rows]

migrations/pg/0010_shards_index.py

@@ -0,0 +1,34 @@
+# 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 nucliadb.common.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    # Concurrent index must be created outside of a transaction but psycopg automatically
+    # creates transactions. We temporarily disable this for building indexes.
+    await txn.connection.commit()
+    try:
+        await txn.connection.set_autocommit(True)
+        await txn.connection.execute(
+            "CREATE INDEX CONCURRENTLY ON resources (key, value) WHERE key ~ '/kbs/[^/]*/r/[^/]*/shard$';"
+        )
+    finally:
+        await txn.connection.set_autocommit(False)

nucliadb/search/api/v1/resource/utils.py → migrations/pg/0011_catalog_statistics.py

@@ -18,11 +18,10 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-from typing import Optional
+from nucliadb.common.maindb.pg import PGTransaction
 
-from nucliadb.common import datamanagers
 
-
-async def get_resource_uuid_by_slug(kbid: str, slug: str) -> Optional[str]:
-    async with datamanagers.with_ro_transaction() as txn:
-        return await datamanagers.resources.get_resource_uuid_from_slug(txn, kbid=kbid, slug=slug)
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute("CREATE STATISTICS catalog_kbid_labels ON kbid, labels FROM catalog;")
+        await cur.execute("ANALYZE catalog;")

migrations/pg/0012_catalog_statistics_undo.py

@@ -0,0 +1,26 @@
+# 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 nucliadb.common.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute("DROP STATISTICS catalog_kbid_labels;")

nucliadb/backups/create.py

@@ -21,8 +21,8 @@ import asyncio
 import json
 import logging
 import tarfile
+from collections.abc import AsyncIterator
 from datetime import datetime, timezone
-from typing import AsyncIterator, Optional
 
 from nucliadb.backups.const import (
     BackupFinishedStream,
@@ -37,7 +37,6 @@ from nucliadb.export_import.utils import (
     download_binary,
     get_broker_message,
     get_cloud_files,
-    get_entities,
     get_labels,
     get_search_configurations,
     get_synonyms,
@@ -76,7 +75,6 @@ async def backup_kb(context: ApplicationContext, kbid: str, backup_id: str):
     """
     await backup_resources(context, kbid, backup_id)
     await backup_labels(context, kbid, backup_id)
-    await backup_entities(context, kbid, backup_id)
     await backup_synonyms(context, kbid, backup_id)
     await backup_search_configurations(context, kbid, backup_id)
     await notify_backup_completed(context, kbid, backup_id)
@@ -235,15 +233,6 @@ async def backup_labels(context: ApplicationContext, kbid: str, backup_id: str):
     )
 
 
-async def backup_entities(context: ApplicationContext, kbid: str, backup_id: str):
-    entities = await get_entities(context, kbid)
-    await context.blob_storage.upload_object(
-        bucket=settings.backups_bucket,
-        key=StorageKeys.ENTITIES.format(backup_id=backup_id),
-        data=entities.SerializeToString(),
-    )
-
-
 async def backup_synonyms(context: ApplicationContext, kbid: str, backup_id: str):
     synonyms = await get_synonyms(context, kbid)
     await context.blob_storage.upload_object(
@@ -266,9 +255,7 @@ async def backup_search_configurations(context: ApplicationContext, kbid: str, b
     )
 
 
-async def get_metadata(
-    context: ApplicationContext, kbid: str, backup_id: str
-) -> Optional[BackupMetadata]:
+async def get_metadata(context: ApplicationContext, kbid: str, backup_id: str) -> BackupMetadata | None:
     async with context.kv_driver.ro_transaction() as txn:
         metadata_raw = await txn.get(MaindbKeys.METADATA.format(kbid=kbid, backup_id=backup_id))
         if metadata_raw is None:

nucliadb/backups/restore.py

@@ -24,7 +24,8 @@ import functools
 import json
 import logging
 import tarfile
-from typing import Any, AsyncIterator, Callable, Optional, Union
+from collections.abc import AsyncIterator, Callable
+from typing import Any
 
 from pydantic import TypeAdapter
 
@@ -35,7 +36,6 @@ from nucliadb.common.context import ApplicationContext
 from nucliadb.export_import.utils import (
     import_binary,
     restore_broker_message,
-    set_entities_groups,
     set_labels,
     set_search_configurations,
     set_synonyms,
@@ -74,7 +74,6 @@ async def restore_kb(context: ApplicationContext, kbid: str, backup_id: str):
     """
     await restore_resources(context, kbid, backup_id)
     await restore_labels(context, kbid, backup_id)
-    await restore_entities(context, kbid, backup_id)
     await restore_synonyms(context, kbid, backup_id)
     await restore_search_configurations(context, kbid, backup_id)
     await delete_last_restored(context, kbid, backup_id)
@@ -101,7 +100,7 @@ async def restore_resources(context: ApplicationContext, kbid: str, backup_id: s
         await set_last_restored(context, kbid, backup_id, key)
 
 
-async def get_last_restored(context: ApplicationContext, kbid: str, backup_id: str) -> Optional[str]:
+async def get_last_restored(context: ApplicationContext, kbid: str, backup_id: str) -> str | None:
     key = MaindbKeys.LAST_RESTORED.format(kbid=kbid, backup_id=backup_id)
     async with context.kv_driver.ro_transaction() as txn:
         raw = await txn.get(key)
@@ -191,7 +190,7 @@ class ResourceBackupReader:
         data = await self.read(tarinfo_size + padding_bytes)
         return data[:tarinfo_size]
 
-    async def read_item(self) -> Union[BrokerMessage, CloudFile, CloudFileBinary]:
+    async def read_item(self) -> BrokerMessage | CloudFile | CloudFileBinary:
         tarinfo = await self.read_tarinfo()
         if tarinfo.name.startswith("broker-message"):
             raw_bm = await self.read_data(tarinfo)
@@ -257,16 +256,6 @@ async def restore_labels(context: ApplicationContext, kbid: str, backup_id: str)
     await set_labels(context, kbid, labels)
 
 
-async def restore_entities(context: ApplicationContext, kbid: str, backup_id: str):
-    raw = await context.blob_storage.downloadbytes(
-        bucket=settings.backups_bucket,
-        key=StorageKeys.ENTITIES.format(backup_id=backup_id),
-    )
-    entities = kb_pb2.EntitiesGroups()
-    entities.ParseFromString(raw.getvalue())
-    await set_entities_groups(context, kbid, entities)
-
-
 async def restore_synonyms(context: ApplicationContext, kbid: str, backup_id: str):
     raw = await context.blob_storage.downloadbytes(
         bucket=settings.backups_bucket,