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

migrations/0003_allfields_key.py

@@ -17,45 +17,11 @@
 # 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 typing import Optional
 
-from nucliadb_protos.resources_pb2 import AllFieldIDs, FieldID
-
-from nucliadb.common import datamanagers
 from nucliadb.migrator.context import ExecutionContext
-from nucliadb.migrator.migrator import logger
 
 
 async def migrate(context: ExecutionContext) -> None: ...
 
 
-async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    async for resource_id in datamanagers.resources.iterate_resource_ids(kbid=kbid):
-        async with context.kv_driver.transaction() as txn:
-            resource = await datamanagers.resources.get_resource(
-                txn, kbid=kbid, rid=resource_id
-            )
-            if resource is None:
-                logger.warning(
-                    f"kb={kbid} rid={resource_id}: resource not found. Skipping..."
-                )
-                continue
-
-            all_fields: Optional[AllFieldIDs] = await resource.get_all_field_ids()
-            if all_fields is not None:
-                logger.warning(
-                    f"kb={kbid} rid={resource_id}: already has all fields key. Skipping..."
-                )
-                continue
-
-            # Migrate resource
-            logger.warning(f"kb={kbid} rid={resource_id}: migrating...")
-            all_fields = AllFieldIDs()
-            async for (
-                field_type,
-                field_id,
-            ) in resource._deprecated_scan_fields_ids():
-                fid = FieldID(field_type=field_type, field=field_id)
-                all_fields.fields.append(fid)
-            await resource.set_all_field_ids(all_fields)
-            await txn.commit()
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None: ...

migrations/0009_upgrade_relations_and_texts_to_v2.py

@@ -26,7 +26,6 @@ rollover will do the upgrade automatically.
 
 """
 
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 
@@ -34,4 +33,7 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    await rollover_kb_shards(context, kbid)
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0010_fix_corrupt_indexes.py

@@ -27,23 +27,23 @@ index data loss. Rollover affected KBs
 
 import logging
 
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
 
-AFFECTED_KBS = [
-    "1efc5a33-bc5a-490c-8b47-b190beee212d",
-    "f11d6eb9-da5e-4519-ac3d-e304bfa5c354",
-    "096d9070-f7be-40c8-a24c-19c89072e3ff",
-    "848f01bc-341a-4346-b473-6b11b76b26eb",
-]
+# AFFECTED_KBS = [
+#     "1efc5a33-bc5a-490c-8b47-b190beee212d",
+#     "f11d6eb9-da5e-4519-ac3d-e304bfa5c354",
+#     "096d9070-f7be-40c8-a24c-19c89072e3ff",
+#     "848f01bc-341a-4346-b473-6b11b76b26eb",
+# ]
 
 
 async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    if kbid in AFFECTED_KBS:
-        logger.info(f"Rolling over affected KB: {kbid}")
-        await rollover_kb_shards(context, kbid)
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0011_materialize_labelset_ids.py

@@ -26,7 +26,6 @@ Tikv doesn't really like scanning a lot of keys, so we need to materialize the l
 
 import logging
 
-from nucliadb.common import datamanagers
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -35,18 +34,4 @@ logger = logging.getLogger(__name__)
 async def migrate(context: ExecutionContext) -> None: ...
 
 
-async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    async with context.kv_driver.transaction() as txn:
-        labelset_list = await datamanagers.labels._get_labelset_ids(txn, kbid=kbid)
-        if labelset_list is not None:
-            logger.info("No need for labelset list migration", extra={"kbid": kbid})
-            return
-
-        labelset_list = await datamanagers.labels._deprecated_scan_labelset_ids(
-            txn, kbid=kbid
-        )
-        await datamanagers.labels._set_labelset_ids(
-            txn, kbid=kbid, labelsets=labelset_list
-        )
-        logger.info("Labelset list migrated", extra={"kbid": kbid})
-        await txn.commit()
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None: ...

migrations/0012_rollover_shards.py

@@ -23,25 +23,20 @@ Allow targeted rollover of KBs
 """
 
 import logging
-import os
 
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
 
 
-AFFECTED_KBS = [
-    kbid.strip()
-    for kbid in os.environ.get("ROLLOVER_KBS", "").split(",")
-    if kbid.strip()
-]
+# AFFECTED_KBS = [kbid.strip() for kbid in os.environ.get("ROLLOVER_KBS", "").split(",") if kbid.strip()]
 
 
 async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    if kbid in AFFECTED_KBS:
-        logger.info(f"Rolling over affected KB: {kbid}")
-        await rollover_kb_shards(context, kbid)
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0014_rollover_shards.py

@@ -23,9 +23,7 @@ Allow targeted rollover of KBs
 """
 
 import logging
-import os
 
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -35,6 +33,7 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    if os.environ.get("RUNNING_ENVIRONMENT", os.environ.get("ENVIRONMENT")) == "stage":
-        logger.info(f"Rolling over affected KB: {kbid}")
-        await rollover_kb_shards(context, kbid)
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0015_targeted_rollover.py

@@ -24,25 +24,20 @@ Targeted rollover for a specific KB
 """
 
 import logging
-import os
 
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
 
 
-AFFECTED_KBS = [
-    kbid.strip()
-    for kbid in os.environ.get("ROLLOVER_KBS", "").split(",")
-    if kbid.strip()
-]
+# AFFECTED_KBS = [kbid.strip() for kbid in os.environ.get("ROLLOVER_KBS", "").split(",") if kbid.strip()]
 
 
 async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    if kbid in AFFECTED_KBS:
-        logger.info(f"Rolling over affected KB: {kbid}")
-        await rollover_kb_shards(context, kbid)
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0016_upgrade_to_paragraphs_v2.py

@@ -25,10 +25,6 @@ Targeted rollover for a specific KBs which still don't have the latest version o
 
 import logging
 
-from nucliadb_protos.noderesources_pb2 import ShardCreated
-
-from nucliadb.common import datamanagers
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -41,29 +37,30 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    try:
-        if await has_old_paragraphs_index(context, kbid):
-            logger.info("Rolling over affected KB", extra={"kbid": kbid})
-            await rollover_kb_shards(context, kbid)
-        else:
-            logger.info(
-                "KB already has the latest version of the paragraphs index, skipping rollover",
-                extra={"kbid": kbid},
-            )
-    except ShardsObjectNotFound:
-        logger.warning("KB not found, skipping rollover", extra={"kbid": kbid})
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """
+    # try:
+    #     if await has_old_paragraphs_index(context, kbid):
+    #         logger.info("Rolling over affected KB", extra={"kbid": kbid})
+    #         await rollover_kb_index(context, kbid)
+    #     else:
+    #         logger.info(
+    #             "KB already has the latest version of the paragraphs index, skipping rollover",
+    #             extra={"kbid": kbid},
+    #         )
+    # except ShardsObjectNotFound:
+    #     logger.warning("KB not found, skipping rollover", extra={"kbid": kbid})
 
 
-async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
-    async with context.kv_driver.transaction(read_only=True) as txn:
-        shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
-        if not shards_object:
-            raise ShardsObjectNotFound()
-        for shard in shards_object.shards:
-            for replica in shard.replicas:
-                if (
-                    replica.shard.paragraph_service
-                    != ShardCreated.ParagraphService.PARAGRAPH_V2
-                ):
-                    return True
-        return False
+# async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
+#     async with context.kv_driver.transaction(read_only=True) as txn:
+#         shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=False)
+#         if not shards_object:
+#             raise ShardsObjectNotFound()
+#         for shard in shards_object.shards:
+#             for replica in shard.replicas:
+#                 if replica.shard.paragraph_service != ShardCreated.ParagraphService.PARAGRAPH_V2:
+#                     return True
+#         return False

migrations/0017_multiple_writable_shards.py

@@ -42,7 +42,7 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
     async with context.kv_driver.transaction() as txn:
-        shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
+        shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid, for_update=True)
         if shards is None:
             logger.error("KB without shards", extra={"kbid": kbid})
             return
@@ -52,9 +52,7 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
         shards.shards[shards.actual].read_only = False
 
         # just ensure we're writing it correctly
-        assert [shard_object.read_only for shard_object in shards.shards].count(
-            False
-        ) == 1
+        assert [shard_object.read_only for shard_object in shards.shards].count(False) == 1
 
         await datamanagers.cluster.update_kb_shards(txn, kbid=kbid, shards=shards)
         await txn.commit()

migrations/0018_purge_orphan_kbslugs.py

@@ -18,17 +18,17 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-"""Migration #17
+"""Migration #18
 
 Due to a bug on backend services, some kbslugs were not properly deleted and got
 orphan. Let's delete them!
 
 """
+
 import logging
 
 from nucliadb.common import datamanagers
 from nucliadb.common.datamanagers.kb import KB_SLUGS_BASE
-from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -36,14 +36,12 @@ logger = logging.getLogger(__name__)
 
 async def migrate(context: ExecutionContext) -> None:
     async with context.kv_driver.transaction() as txn:
-        async for key in txn.keys(KB_SLUGS_BASE, count=-1):
+        async for key in txn.keys(KB_SLUGS_BASE):
             slug = key.replace(KB_SLUGS_BASE, "")
-            value = await txn.get(key)
+            value = await txn.get(key, for_update=False)
             if value is None:
                 # KB with slug but without uuid? Seems wrong, let's remove it too
-                logger.info(
-                    "Removing /kbslugs with empty value", extra={"maindb_key": key}
-                )
+                logger.info("Removing /kbslugs with empty value", extra={"maindb_key": key})
                 await txn.delete(key)
                 continue
 

migrations/0019_upgrade_to_paragraphs_v3.py

@@ -25,10 +25,6 @@ Targeted rollover for a specific KBs which still don't have the latest version o
 
 import logging
 
-from nucliadb_protos.noderesources_pb2 import ShardCreated
-
-from nucliadb.common import datamanagers
-from nucliadb.common.cluster.rollover import rollover_kb_shards
 from nucliadb.migrator.context import ExecutionContext
 
 logger = logging.getLogger(__name__)
@@ -41,29 +37,30 @@ async def migrate(context: ExecutionContext) -> None: ...
 
 
 async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
-    try:
-        if await has_old_paragraphs_index(context, kbid):
-            logger.info("Rolling over affected KB", extra={"kbid": kbid})
-            await rollover_kb_shards(context, kbid)
-        else:
-            logger.info(
-                "KB already has the latest version of the paragraphs index, skipping rollover",
-                extra={"kbid": kbid},
-            )
-    except ShardsObjectNotFound:
-        logger.warning("KB not found, skipping rollover", extra={"kbid": kbid})
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """
+    # try:
+    #     if await has_old_paragraphs_index(context, kbid):
+    #         logger.info("Rolling over affected KB", extra={"kbid": kbid})
+    #         await rollover_kb_index(context, kbid)
+    #     else:
+    #         logger.info(
+    #             "KB already has the latest version of the paragraphs index, skipping rollover",
+    #             extra={"kbid": kbid},
+    #         )
+    # except ShardsObjectNotFound:
+    #     logger.warning("KB not found, skipping rollover", extra={"kbid": kbid})
 
 
-async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
-    async with context.kv_driver.transaction(read_only=True) as txn:
-        shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
-        if not shards_object:
-            raise ShardsObjectNotFound()
-        for shard in shards_object.shards:
-            for replica in shard.replicas:
-                if (
-                    replica.shard.paragraph_service
-                    != ShardCreated.ParagraphService.PARAGRAPH_V3
-                ):
-                    return True
-        return False
+# async def has_old_paragraphs_index(context: ExecutionContext, kbid: str) -> bool:
+#     async with context.kv_driver.transaction(read_only=True) as txn:
+#         shards_object = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
+#         if not shards_object:
+#             raise ShardsObjectNotFound()
+#         for shard in shards_object.shards:
+#             for replica in shard.replicas:
+#                 if replica.shard.paragraph_service != ShardCreated.ParagraphService.PARAGRAPH_V3:
+#                     return True
+#         return False

migrations/0020_drain_nodes_from_cluster.py

@@ -29,7 +29,7 @@ create new shards in the remaining nodes.
 import logging
 
 from nucliadb.common import datamanagers
-from nucliadb.common.cluster.rollover import rollover_kb_shards
+from nucliadb.common.cluster.rollover import rollover_kb_index
 from nucliadb.common.cluster.settings import settings as cluster_settings
 from nucliadb.migrator.context import ExecutionContext
 
@@ -56,11 +56,11 @@ async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
         return
 
     logger.info("Rolling over affected KB", extra={"kbid": kbid})
-    await rollover_kb_shards(context, kbid, drain_nodes=drain_node_ids)
+    await rollover_kb_index(context, kbid, drain_nodes=drain_node_ids)
 
 
 async def kb_has_shards_on_drain_nodes(kbid: str, drain_node_ids: list[str]) -> bool:
-    async with datamanagers.with_transaction(read_only=True) as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         shards = await datamanagers.cluster.get_kb_shards(txn, kbid=kbid)
         if not shards:
             logger.warning("Shards object not found", extra={"kbid": kbid})

nucliadb/standalone/tests/unit/test_run.py → migrations/0021_overwrite_vectorsets_key.py

@@ -17,31 +17,28 @@
 # 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 unittest import mock
 
-import pytest
+"""Migration #21
 
-from nucliadb.standalone.run import run, run_async_nucliadb
-from nucliadb.standalone.settings import Settings
+With the new vectorsets implementation, we need to store some information on
+maindb. As the key "/kbs/{kbid}/vectorsets" was already used at some point, this
+migration will ensure to overwrite the key and set the new value
 
-STANDALONE_RUN = "nucliadb.standalone.run"
+"""
 
+import logging
 
-@pytest.fixture(scope="function", autouse=True)
-def mocked_deps():
-    with (
-        mock.patch("uvicorn.Server.run"),
-        mock.patch("nucliadb.standalone.run.parser", return_value=Settings()),
-        mock.patch(f"{STANDALONE_RUN}.get_latest_nucliadb", return_value="1.0.0"),
-        mock.patch("uvicorn.Server.startup"),
-        mock.patch(f"{STANDALONE_RUN}.run_migrations"),
-    ):
-        yield
+from nucliadb.common import datamanagers
+from nucliadb.migrator.context import ExecutionContext
 
+logger = logging.getLogger(__name__)
 
-def test_run():
-    run()
 
+async def migrate(context: ExecutionContext) -> None: ...
 
-async def test_run_async_nucliadb():
-    await run_async_nucliadb(Settings())
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    async with context.kv_driver.transaction() as txn:
+        logger.info(f"Overwriting vectorsets key", extra={"kbid": kbid})
+        await datamanagers.vectorsets.initialize(txn, kbid=kbid)
+        await txn.commit()

nucliadb/tests/unit/test_openapi.py → migrations/0022_fix_paragraph_deletion_bug.py

@@ -17,22 +17,27 @@
 # 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 unittest.mock import Mock
 
-from starlette.routing import Mount
+"""Migration #22
 
-from nucliadb.openapi import extract_openapi, is_versioned_route
-from nucliadb.search.app import application
+There was a bug while ingesting/indexing that made paragraphs not being properly
+removed in some cases. This rollover migration ensures data is consistently
+indexed.
 
+"""
 
-def get_route(path):
-    return Mount(path=path, app=Mock())
+import logging
 
+from nucliadb.migrator.context import ExecutionContext
 
-def test_is_versioned_route():
-    assert is_versioned_route(get_route(path="/api/v1/search"))
-    assert not is_versioned_route(get_route(path="/metrics"))
+logger = logging.getLogger(__name__)
 
 
-def test_extract_openapi():
-    assert extract_openapi(application, "1", "commitid", "nucliadb_search")
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    """
+    We only need 1 rollover migration defined at a time; otherwise, we will
+    possibly run many for a kb when we only ever need to run one
+    """

migrations/0023_backfill_pg_catalog.py

@@ -0,0 +1,80 @@
+# 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 #23
+
+Backfill the data into the PG catalog
+
+"""
+
+import logging
+from typing import cast
+
+from nucliadb.common import datamanagers
+from nucliadb.common.maindb.pg import PGDriver, PGTransaction
+from nucliadb.ingest.orm.processor.pgcatalog import pgcatalog_update
+from nucliadb.migrator.context import ExecutionContext
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None: ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    if not isinstance(context.kv_driver, PGDriver):
+        return
+
+    BATCH_SIZE = 100
+    async with context.kv_driver.transaction() as txn:
+        txn = cast(PGTransaction, txn)
+        continue_sql = ""
+        while True:
+            async with txn.connection.cursor() as cur:
+                # Get list of resources except those already in the catalog
+                await cur.execute(
+                    f"""
+                    SELECT SPLIT_PART(key, '/', 5)::UUID FROM resources
+                    LEFT JOIN catalog ON kbid = %s AND SPLIT_PART(key, '/', 5)::UUID = rid
+                    WHERE key SIMILAR TO %s
+                    AND rid IS NULL
+                    {continue_sql}
+                    ORDER BY key
+                    LIMIT %s
+                    """,
+                    (kbid, f"/kbs/{kbid}/r/[a-f0-9]*", BATCH_SIZE),
+                )
+                resources_to_index = [r[0] for r in await cur.fetchall()]
+                if len(resources_to_index) == 0:
+                    return
+
+                # Index each resource
+                for rid in resources_to_index:
+                    rid = str(rid).replace("-", "")
+                    resource = await datamanagers.resources.get_resource(txn, kbid=kbid, rid=rid)
+                    if resource is None:
+                        logger.warning(f"Could not load resource {rid} for kbid {kbid}")
+                        continue
+
+                    await resource.compute_global_tags(resource.indexer)
+                    await pgcatalog_update(txn, kbid, resource)
+
+                await txn.commit()
+                continue_sql = f"AND key > '/kbs/{kbid}/r/{rid}'"