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

nucliadb/migrator/migrator.py

@@ -22,23 +22,20 @@ import logging
 from typing import Optional
 
 from nucliadb.common import locking
-from nucliadb.common.cluster.rollover import rollover_kb_shards
+from nucliadb.common.cluster.rollover import rollover_kb_index
 from nucliadb.common.cluster.settings import in_standalone_mode
+from nucliadb.common.maindb.pg import PGDriver
 from nucliadb.migrator.context import ExecutionContext
-from nucliadb.migrator.utils import get_migrations
+from nucliadb.migrator.utils import get_migrations, get_pg_migrations
 from nucliadb_telemetry import errors, metrics
 
-migration_observer = metrics.Observer(
-    "nucliadb_migrations", labels={"type": "kb", "target_version": ""}
-)
+migration_observer = metrics.Observer("nucliadb_migrations", labels={"type": "kb", "target_version": ""})
 
 
 logger = logging.getLogger(__name__)
 
 
-async def run_kb_migrations(
-    context: ExecutionContext, kbid: str, target_version: int
-) -> None:
+async def run_kb_migrations(context: ExecutionContext, kbid: str, target_version: int) -> None:
     async with locking.distributed_lock(f"migration-{kbid}"):
         kb_info = await context.data_manager.get_kb_info(kbid)
         if kb_info is None:
@@ -46,9 +43,7 @@ async def run_kb_migrations(
             await context.data_manager.delete_kb_migration(kbid=kbid)
             return
 
-        migrations = get_migrations(
-            from_version=kb_info.current_version, to_version=target_version
-        )
+        migrations = get_migrations(from_version=kb_info.current_version, to_version=target_version)
 
         for migration in migrations:
             migration_info = {
@@ -59,14 +54,10 @@ async def run_kb_migrations(
 
             try:
                 logger.info("Migrating KB", extra=migration_info)
-                with migration_observer(
-                    {"type": "kb", "target_version": str(migration.version)}
-                ):
-                    await migration.module.migrate_kb(context, kbid)  # type: ignore
+                with migration_observer({"type": "kb", "target_version": str(migration.version)}):
+                    await migration.module.migrate_kb(context, kbid)
                 logger.info("Finished KB Migration", extra=migration_info)
-                await context.data_manager.update_kb_info(
-                    kbid=kbid, current_version=migration.version
-                )
+                await context.data_manager.update_kb_info(kbid=kbid, current_version=migration.version)
             except Exception as exc:
                 errors.capture_exception(exc)
                 logger.exception("Failed to migrate KB", extra=migration_info)
@@ -74,9 +65,7 @@ async def run_kb_migrations(
 
         refreshed_kb_info = await context.data_manager.get_kb_info(kbid=kbid)
         if refreshed_kb_info is None:
-            logger.warning(
-                "KB not found. This should not happen.", extra={"kbid": kbid}
-            )
+            logger.warning("KB not found. This should not happen.", extra={"kbid": kbid})
             return
         assert refreshed_kb_info.current_version == target_version
 
@@ -88,7 +77,7 @@ async def run_all_kb_migrations(context: ExecutionContext, target_version: int)
     Schedule all KB migrations to run in parallel. Only a certain number of migrations will run at the same time.
     If any of the migrations fail, the whole process will fail.
     """
-    to_migrate = await context.data_manager.get_kb_migrations(limit=-1)
+    to_migrate = await context.data_manager.get_kb_migrations()
 
     if len(to_migrate) == 0:
         return
@@ -154,13 +143,9 @@ async def run_global_migrations(context: ExecutionContext, target_version: int)
         }
         try:
             logger.info("Migrating", extra=migration_info)
-            with migration_observer(
-                {"type": "global", "target_version": str(migration.version)}
-            ):
-                await migration.module.migrate(context)  # type: ignore
-            await context.data_manager.update_global_info(
-                current_version=migration.version
-            )
+            with migration_observer({"type": "global", "target_version": str(migration.version)}):
+                await migration.module.migrate(context)
+            await context.data_manager.update_global_info(current_version=migration.version)
             logger.info("Finished migration", extra=migration_info)
         except Exception as exc:
             errors.capture_exception(exc)
@@ -177,7 +162,7 @@ async def run_rollover_in_parallel(
 ) -> None:
     async with max_concurrent:
         try:
-            await rollover_kb_shards(context, kbid)
+            await rollover_kb_index(context, kbid)
             await context.data_manager.delete_kb_rollover(kbid=kbid)
         except Exception as exc:
             errors.capture_exception(exc)
@@ -221,7 +206,39 @@ async def run_rollovers(context: ExecutionContext) -> None:
         raise Exception(f"Failed to migrate KBs. Failures: {failures}")
 
 
+async def run_pg_schema_migrations(driver: PGDriver):
+    migrations = get_pg_migrations()
+
+    # The migration uses two transactions. The former is only used to get a lock (pg_advisory_lock)
+    # without having to worry about correctly unlocking it (postgres unlocks it when the transaction ends)
+    async with driver.transaction() as tx_lock, tx_lock.connection.cursor() as cur_lock:  # type: ignore[attr-defined]
+        await cur_lock.execute(
+            "CREATE TABLE IF NOT EXISTS migrations (version INT PRIMARY KEY, migrated_at TIMESTAMP NOT NULL DEFAULT NOW())"
+        )
+        await tx_lock.commit()
+        await cur_lock.execute("SELECT pg_advisory_xact_lock(3116614845278015934)")
+
+        await cur_lock.execute("SELECT version FROM migrations")
+        migrated = [r[0] for r in await cur_lock.fetchall()]
+
+        for version, migration in migrations:
+            if version in migrated:
+                continue
+
+            # Gets a new transaction for each migration, so if they get interrupted we at least
+            # save the state of the last finished transaction
+            async with driver.transaction() as tx, tx.connection.cursor() as cur:  # type: ignore[attr-defined]
+                await migration.migrate(tx)
+                await cur.execute("INSERT INTO migrations (version) VALUES (%s)", (version,))
+                await tx.commit()
+
+
 async def run(context: ExecutionContext, target_version: Optional[int] = None) -> None:
+    # Run schema migrations first, since they create the `resources` table needed for the lock below
+    # Schema migrations use their own locking system
+    if isinstance(context.kv_driver, PGDriver):
+        await run_pg_schema_migrations(context.kv_driver)
+
     async with locking.distributed_lock(locking.MIGRATIONS_LOCK):
         # before we move to managed migrations, see if there are any rollovers
         # scheduled and run them

nucliadb/migrator/utils.py

@@ -17,13 +17,12 @@
 # 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 importlib
 import logging
 import os
 import types
 from functools import lru_cache
 
-import migrations
-
 from .models import Migration
 
 logger = logging.getLogger(__name__)
@@ -33,14 +32,27 @@ MIGRATION_DIR = os.path.sep.join(
 )
 
 
+def get_pg_migrations() -> list[tuple[int, types.ModuleType]]:
+    output = []
+    for filename in os.listdir(os.path.join(MIGRATION_DIR, "pg")):
+        if filename.endswith(".py") and filename != "__init__.py":
+            module_name = filename[:-3]
+            version = int(module_name.split("_")[0])
+            module = importlib.import_module(f"migrations.pg.{module_name}")
+            if not hasattr(module, "migrate"):
+                raise Exception(f"Missing `migrate` function in {module_name}")
+            output.append((version, module))
+    output.sort()
+    return output
+
+
 def get_migration_modules() -> list[tuple[types.ModuleType, int]]:
     output = []
     for filename in os.listdir(MIGRATION_DIR):
         if filename.endswith(".py") and filename != "__init__.py":
             module_name = filename[:-3]
             version = int(module_name.split("_")[0])
-            __import__(f"migrations.{module_name}")
-            module = getattr(migrations, module_name)
+            module = importlib.import_module(f"migrations.{module_name}")
             if not hasattr(module, "migrate"):
                 raise Exception(f"Missing `migrate` function in {module_name}")
             if not hasattr(module, "migrate_kb"):
@@ -49,17 +61,13 @@ def get_migration_modules() -> list[tuple[types.ModuleType, int]]:
     return output
 
 
-def get_migrations(
-    from_version: int = 0, to_version: int = 99999999
-) -> list[Migration]:
+def get_migrations(from_version: int = 0, to_version: int = 99999999) -> list[Migration]:
     migrations: list[Migration] = []
     for module, version in get_migration_modules():
         migrations.append(Migration(version=version, module=module))
 
     migrations.sort(key=lambda m: m.version)
-    return [
-        m for m in migrations if m.version > from_version and m.version <= to_version
-    ]
+    return [m for m in migrations if m.version > from_version and m.version <= to_version]
 
 
 @lru_cache(maxsize=None)

nucliadb/purge/__init__.py

@@ -18,10 +18,9 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
+import importlib.metadata
 from typing import AsyncGenerator
 
-import pkg_resources
-
 from nucliadb.common.cluster.exceptions import NodeError, ShardNotFound
 from nucliadb.common.cluster.utils import setup_cluster, teardown_cluster
 from nucliadb.common.maindb.driver import Driver
@@ -31,6 +30,9 @@ from nucliadb.ingest.orm.knowledgebox import (
     KB_TO_DELETE,
     KB_TO_DELETE_BASE,
     KB_TO_DELETE_STORAGE_BASE,
+    KB_VECTORSET_TO_DELETE,
+    KB_VECTORSET_TO_DELETE_BASE,
+    RESOURCE_TO_DELETE_STORAGE_BASE,
     KnowledgeBox,
 )
 from nucliadb_telemetry import errors
@@ -41,7 +43,7 @@ from nucliadb_utils.utilities import get_storage
 
 async def _iter_keys(driver: Driver, match: str) -> AsyncGenerator[str, None]:
     async with driver.transaction(read_only=True) as keys_txn:
-        async for key in keys_txn.keys(match=match, count=-1):
+        async for key in keys_txn.keys(match=match):
             yield key
 
 
@@ -52,9 +54,7 @@ async def purge_kb(driver: Driver):
         try:
             kbid = key.split("/")[2]
         except Exception:
-            logger.warning(
-                f"  X Skipping purge {key}, wrong key format, expected {KB_TO_DELETE_BASE}"
-            )
+            logger.warning(f"  X Skipping purge {key}, wrong key format, expected {KB_TO_DELETE_BASE}")
             continue
 
         try:
@@ -62,15 +62,11 @@ async def purge_kb(driver: Driver):
             logger.info(f"  √ Successfully Purged {kbid}")
         except ShardNotFound as exc:
             errors.capture_exception(exc)
-            logger.error(
-                f"  X At least one shard was unavailable while purging {kbid}, skipping"
-            )
+            logger.error(f"  X At least one shard was unavailable while purging {kbid}, skipping")
             continue
         except NodeError as exc:
             errors.capture_exception(exc)
-            logger.error(
-                f"  X At least one node was unavailable while purging {kbid}, skipping"
-            )
+            logger.error(f"  X At least one node was unavailable while purging {kbid}, skipping")
             continue
 
         except Exception as exc:
@@ -82,10 +78,10 @@ async def purge_kb(driver: Driver):
 
         # Now delete the tikv delete mark
         try:
-            txn = await driver.begin()
-            key_to_purge = KB_TO_DELETE.format(kbid=kbid)
-            await txn.delete(key_to_purge)
-            await txn.commit()
+            async with driver.transaction() as txn:
+                key_to_purge = KB_TO_DELETE.format(kbid=kbid)
+                await txn.delete(key_to_purge)
+                await txn.commit()
             logger.info(f"  √ Deleted {key_to_purge}")
         except Exception as exc:
             errors.capture_exception(exc)
@@ -112,16 +108,12 @@ async def purge_kb_storage(driver: Driver, storage: Storage):
 
         delete_marker = False
         if conflict:
-            logger.info(
-                f"  . Nothing was deleted for {key}, (Bucket not yet empty), will try next time"
-            )
+            logger.info(f"  . Nothing was deleted for {key}, (Bucket not yet empty), will try next time")
             # Just in case something failed while setting a lifecycle policy to
             # remove all elements from the bucket, reschedule it
             await storage.schedule_delete_kb(kbid)
         elif not deleted:
-            logger.info(
-                f"  ! Expected bucket for {key} was not found, will delete marker"
-            )
+            logger.info(f"  ! Expected bucket for {key} was not found, will delete marker")
             delete_marker = True
         elif deleted:
             logger.info("  √ Bucket successfully deleted")
@@ -129,19 +121,122 @@ async def purge_kb_storage(driver: Driver, storage: Storage):
 
         if delete_marker:
             try:
-                txn = await driver.begin()
-                await txn.delete(key)
+                async with driver.transaction() as txn:
+                    await txn.delete(key)
+                    await txn.commit()
                 logger.info(f"  √ Deleted storage deletion marker {key}")
             except Exception as exc:
                 errors.capture_exception(exc)
                 logger.info(f"  X Error while deleting key {key}")
-                await txn.abort()
-            else:
-                await txn.commit()
 
     logger.info("FINISH PURGING KB STORAGE")
 
 
+async def purge_deleted_resource_storage(driver: Driver, storage: Storage) -> None:
+    """
+    Remove from storage all resources marked as deleted.
+
+    Returns the number of resources purged.
+    """
+    logger.info("Starting purge of deleted resource storage")
+    to_purge = await _count_resources_storage_to_purge(driver)
+    logger.info(f"Found {to_purge} resources to purge")
+    while True:
+        try:
+            purged = await _purge_resources_storage_batch(driver, storage, batch_size=100)
+            if not purged:
+                logger.info("No more resources to purge found")
+                return
+            logger.info(f"Purged {purged} resources")
+
+        except asyncio.CancelledError:
+            logger.info("Purge of deleted resource storage was cancelled")
+            return
+
+
+async def _count_resources_storage_to_purge(driver: Driver) -> int:
+    """
+    Count the number of resources marked as deleted in storage.
+    """
+    async with driver.transaction(read_only=True) as txn:
+        return await txn.count(match=RESOURCE_TO_DELETE_STORAGE_BASE)
+
+
+async def _purge_resources_storage_batch(driver: Driver, storage: Storage, batch_size: int = 100) -> int:
+    """
+    Remove from storage a batch of resources marked as deleted. Returns the
+    number of resources purged.
+    """
+    # Get the keys of the resources to delete in batches of 100
+    to_delete_batch = []
+    async with driver.transaction(read_only=True) as txn:
+        async for key in txn.keys(match=RESOURCE_TO_DELETE_STORAGE_BASE, count=batch_size):
+            to_delete_batch.append(key)
+
+    if not to_delete_batch:
+        return 0
+
+    # Delete the resources blobs from storage
+    logger.info(f"Purging {len(to_delete_batch)} deleted resources")
+    tasks = []
+    for key in to_delete_batch:
+        kbid, resource_id = key.split("/")[-2:]
+        tasks.append(asyncio.create_task(storage.delete_resource(kbid, resource_id)))
+    await asyncio.gather(*tasks)
+
+    # Delete the schedule-to-delete keys
+    async with driver.transaction() as txn:
+        for key in to_delete_batch:
+            await txn.delete(key)
+        await txn.commit()
+
+    return len(to_delete_batch)
+
+
+async def purge_kb_vectorsets(driver: Driver, storage: Storage):
+    """Vectors for a vectorset are stored in a key inside each resource. Iterate
+    through all resources of the KB and remove any storage object containing
+    vectors for the specific vectorset to purge.
+
+    """
+    logger.info("START PURGING KB VECTORSETS")
+
+    purged = []
+    async for key in _iter_keys(driver, KB_VECTORSET_TO_DELETE_BASE):
+        logger.info(f"Purging vectorsets {key}")
+        try:
+            _base, kbid, vectorset = key.lstrip("/").split("/")
+        except ValueError:
+            logger.info(f"  X Skipping purge {key}, wrong key format, expected {KB_VECTORSET_TO_DELETE}")
+            continue
+
+        try:
+            async with driver.transaction(read_only=True) as txn:
+                kb = KnowledgeBox(txn, storage, kbid)
+                async for resource in kb.iterate_resources():
+                    fields = await resource.get_fields(force=True)
+            # we don't need the maindb transaction anymore to remove vectors from storage
+            for field in fields.values():
+                await field.delete_vectors(vectorset)
+        except Exception as exc:
+            errors.capture_exception(exc)
+            logger.error(
+                f"  X ERROR while executing KB vectorset purge, skipping",
+                exc_info=exc,
+                extra={"kbid": kbid},
+            )
+            continue
+
+        purged.append(key)
+
+    async with driver.transaction() as txn:
+        for key in purged:
+            await txn.delete(key)
+        await txn.commit()
+
+    logger.info("FINISH PURGING KB VECTORSETS")
+
+
 async def main():
     """
     This script will purge all knowledge boxes marked to be deleted in maindb.
@@ -153,17 +248,28 @@ async def main():
         service_name=SERVICE_NAME,
     )
     try:
+        purge_resources_storage_task = asyncio.create_task(
+            purge_deleted_resource_storage(driver, storage)
+        )
         await purge_kb(driver)
         await purge_kb_storage(driver, storage)
+        await purge_kb_vectorsets(driver, storage)
+        await purge_resources_storage_task
+    except Exception as ex:  # pragma: no cover
+        logger.exception("Unhandled exception on purge command")
+        errors.capture_exception(ex)
     finally:
-        await storage.finalize()
-        await teardown_driver()
-        await teardown_cluster()
+        try:
+            purge_resources_storage_task.cancel()
+            await storage.finalize()
+            await teardown_driver()
+            await teardown_cluster()
+        except Exception:  # pragma: no cover
+            logger.exception("Error tearing down utilities on purge command")
+            pass
 
 
 def run() -> int:  # pragma: no cover
     setup_logging()
-
-    errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
-
+    errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
     return asyncio.run(main())

nucliadb/purge/orphan_shards.py

@@ -19,11 +19,11 @@
 
 import argparse
 import asyncio
+import importlib.metadata
 from dataclasses import dataclass
 from typing import Optional
 
-import pkg_resources
-from grpc.aio import AioRpcError  # type: ignore
+from grpc.aio import AioRpcError
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster import manager
@@ -86,7 +86,7 @@ async def detect_orphan_shards(driver: Driver) -> dict[str, ShardLocation]:
     orphan_shard_ids = indexed_shards.keys() - stored_shards.keys()
     orphan_shards: dict[str, ShardLocation] = {}
     unavailable_nodes: set[str] = set()
-    async with datamanagers.with_transaction() as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         for shard_id in orphan_shard_ids:
             node_id = indexed_shards[shard_id].node_id
             node = manager.get_index_node(node_id)  # type: ignore
@@ -99,9 +99,7 @@ async def detect_orphan_shards(driver: Driver) -> dict[str, ShardLocation]:
             # Shards with knwon KB ids can be checked and ignore those comming from
             # an ongoing migration/rollover
             if kbid != UNKNOWN_KB:
-                skip = await datamanagers.rollover.is_rollover_shard(
-                    txn, kbid=kbid, shard_id=shard_id
-                )
+                skip = await datamanagers.rollover.is_rollover_shard(txn, kbid=kbid, shard_id=shard_id)
                 if skip:
                     continue
 
@@ -133,18 +131,14 @@ async def _get_stored_shards(driver: Driver) -> dict[str, ShardLocation]:
             try:
                 kb_shards = await shards_manager.get_shards_by_kbid(kbid)
             except ShardsNotFound:
-                logger.warning(
-                    "KB not found while looking for orphan shards", extra={"kbid": kbid}
-                )
+                logger.warning("KB not found while looking for orphan shards", extra={"kbid": kbid})
                 continue
             else:
                 for shard_object_pb in kb_shards:
                     for shard_replica_pb in shard_object_pb.replicas:
                         shard_replica_id = shard_replica_pb.shard.id
                         node_id = shard_replica_pb.node
-                        stored_shards[shard_replica_id] = ShardLocation(
-                            kbid=kbid, node_id=node_id
-                        )
+                        stored_shards[shard_replica_id] = ShardLocation(kbid=kbid, node_id=node_id)
     return stored_shards
 
 
@@ -264,6 +258,6 @@ async def main():
 def run() -> int:  # pragma: no cover
     setup_logging()
 
-    errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
+    errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
     return asyncio.run(main())

nucliadb/reader/__init__.py

@@ -29,9 +29,7 @@ API_PREFIX = "api"
 class EndpointFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
         return (
-            record.args is not None
-            and len(record.args) >= 3
-            and record.args[2] not in ("/", "/metrics")  # type: ignore
+            record.args is not None and len(record.args) >= 3 and record.args[2] not in ("/", "/metrics")  # type: ignore
         )
 
 

nucliadb/reader/api/models.py

@@ -22,15 +22,12 @@ from typing import TYPE_CHECKING, Any, Optional, Union
 from pydantic import BaseModel
 
 import nucliadb_models as models
-from nucliadb_models.common import FIELD_TYPES_MAP, FieldTypeName
+from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import (
     ConversationFieldExtractedData,
-    DatetimeFieldExtractedData,
     Error,
     ExtractedDataType,
     FileFieldExtractedData,
-    KeywordsetFieldExtractedData,
-    LayoutFieldExtractedData,
     LinkFieldExtractedData,
     TextFieldExtractedData,
 )
@@ -41,10 +38,7 @@ if TYPE_CHECKING:  # pragma: no cover
             models.FieldText,
             models.FieldFile,
             models.FieldLink,
-            models.FieldLayout,
             models.Conversation,
-            models.FieldKeywordset,
-            models.FieldDatetime,
         ]
     ]
 else:
@@ -60,14 +54,9 @@ class ResourceField(BaseModel):
     error: Optional[Error] = None
 
 
-FIELD_NAMES_TO_PB_TYPE_MAP = {v: k for k, v in FIELD_TYPES_MAP.items()}
-
 FIELD_NAME_TO_EXTRACTED_DATA_FIELD_MAP: dict[FieldTypeName, Any] = {
     FieldTypeName.TEXT: TextFieldExtractedData,
     FieldTypeName.FILE: FileFieldExtractedData,
     FieldTypeName.LINK: LinkFieldExtractedData,
-    FieldTypeName.DATETIME: DatetimeFieldExtractedData,
-    FieldTypeName.KEYWORDSET: KeywordsetFieldExtractedData,
-    FieldTypeName.LAYOUT: LayoutFieldExtractedData,
     FieldTypeName.CONVERSATION: ConversationFieldExtractedData,
 }

nucliadb/reader/api/v1/__init__.py

@@ -20,7 +20,6 @@
 from . import download  # noqa
 from . import export_import  # noqa
 from . import knowledgebox  # noqa
-from . import learning_collector  # noqa
 from . import learning_config  # noqa
 from . import resource  # noqa
 from . import services  # noqa