Fix rollover migration (#1740)

nucliadb/migrations/0008_cleanup_leftover_rollover_metadata.py

@@ -0,0 +1,36 @@
+# 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 [email protected].
+#
+# 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/>.
+#
+
+"""Upgrade relations and texts indices to v2.
+
+This migration is leaning up leftover junk metadata from previous rollover migration from all KB shards objects.
+
+"""
+
+from nucliadb.common.cluster.rollover import clean_rollover_status
+from nucliadb.migrator.context import ExecutionContext
+
+
+async def migrate(context: ExecutionContext) -> None:
+    ...
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None:
+    await clean_rollover_status(context, kbid)

nucliadb/nucliadb/common/cluster/rollover.py

@@ -58,6 +58,11 @@ def _set_rollover_status(rollover_shards: writer_pb2.Shards, status: RolloverSta
     rollover_shards.extra[status.value] = "true"
 
 
+def _clear_rollover_status(rollover_shards: writer_pb2.Shards):
+    for status in RolloverStatus:
+        rollover_shards.extra.pop(status.value, None)
+
+
 class UnexpectedRolloverError(Exception):
     pass
 
@@ -312,6 +317,7 @@ async def cutover_shards(app_context: ApplicationContext, kbid: str) -> None:
     if previously_active_shards is None or rollover_shards is None:
         raise UnexpectedRolloverError("Shards for kb not found")
 
+    _clear_rollover_status(rollover_shards)
     await cluster_datamanager.update_kb_shards(kbid, rollover_shards)
     await rollover_datamanager.delete_kb_rollover_shards(kbid)
 
@@ -436,6 +442,20 @@ async def clean_indexed_data(app_context: ApplicationContext, kbid: str) -> None
         await rollover_datamanager.remove_indexed(kbid, batch)
 
 
+async def clean_rollover_status(app_context: ApplicationContext, kbid: str) -> None:
+    cluster_datamanager = ClusterDataManager(app_context.kv_driver)
+    kb_shards = await cluster_datamanager.get_kb_shards(kbid)
+    if kb_shards is None:
+        logger.warning(
+            "No shards found for KB, skipping clean rollover status",
+            extra={"kbid": kbid},
+        )
+        return
+
+    _clear_rollover_status(kb_shards)
+    await cluster_datamanager.update_kb_shards(kbid, kb_shards)
+
+
 async def rollover_kb_shards(app_context: ApplicationContext, kbid: str) -> None:
     """
     Rollover a shard is the process of creating new shard replicas for every
@@ -471,6 +491,7 @@ async def rollover_kb_shards(app_context: ApplicationContext, kbid: str) -> None
     # we need to cut over BEFORE we validate the data
     await validate_indexed_data(app_context, kbid)
     await clean_indexed_data(app_context, kbid)
+    await clean_rollover_status(app_context, kbid)
 
     logger.warning("Finished rolling over shards", extra={"kbid": kbid})
 

nucliadb/nucliadb/tests/integration/common/cluster/test_rollover.py

@@ -22,6 +22,7 @@
 
 from nucliadb.common.cluster import rollover
 from nucliadb.common.context import ApplicationContext
+from nucliadb.common.datamanagers.cluster import ClusterDataManager
 
 pytestmark = pytest.mark.asyncio
 
@@ -79,3 +80,22 @@ async def test_rollover_kb_shards(
     assert resp.status_code == 200
     body = resp.json()
     assert len(body["resources"]) == count
+
+
+@pytest.mark.parametrize("knowledgebox", ("EXPERIMENTAL", "STABLE"), indirect=True)
+async def test_rollover_kb_shards_does_a_clean_cutover(
+    app_context,
+    knowledgebox,
+):
+    async def get_kb_shards(kbid: str):
+        driver = app_context.kv_driver
+        cluster_data_manager = ClusterDataManager(driver)
+        return await cluster_data_manager.get_kb_shards(kbid)
+
+    shards1 = await get_kb_shards(knowledgebox)
+    assert shards1.extra == {}
+
+    await rollover.rollover_kb_shards(app_context, knowledgebox)
+
+    shards2 = await get_kb_shards(knowledgebox)
+    assert shards2.extra == {}