Return set_after_successful_execution_version

src/tribler/core/upgrade/tribler_db/decorator.py

@@ -1,6 +1,6 @@
 import functools
 import logging
-from typing import Callable
+from typing import Callable, Optional
 
 from tribler.core.components.database.db.tribler_database import TriblerDatabase
 from tribler.core.utilities.pony_utils import db_session
@@ -10,17 +10,15 @@
 logger = logging.getLogger('Migration (TriblerDB)')
 
 
-def migration(execute_only_if_version: int):
+def migration(execute_only_if_version: int, set_after_successful_execution_version: Optional[int] = None):
     """ Decorator for migration functions.
-
     The migration executes in the single transaction. If the migration fails, the transaction is rolled back.
     The decorator also sets the metadata attribute to the decorated function. It could be checked by
     calling the `has_migration_metadata` function.
-
-    After the successful migration, db version is set to the value equal to `execute_only_if_version + 1`.
-
     Args:
         execute_only_if_version: Execute the migration only if the current db version is equal to this value.
+        set_after_successful_execution_version: Set the db version to this value after the migration is executed.
+            If it is not specified, then `set_after_successful_execution_version = execute_only_if_version + 1`
     """
 
     def decorator(func):
@@ -36,7 +34,11 @@ def wrapper(db: TriblerDatabase, **kwargs):
                 return None
 
             result = func(db, **kwargs)
-            db.version = target_version + 1
+
+            next_version = set_after_successful_execution_version
+            if next_version is None:
+                next_version = target_version + 1
+            db.version = next_version
 
             return result
 

src/tribler/core/upgrade/tribler_db/tests/test_decorator.py

@@ -16,22 +16,40 @@ def test(_: Mock):
     assert not test(Mock(version=2))
 
 
-def test_migration_set_next_version():
-    """ Test that the version of the database is set to the next one after the successful migration."""
+def test_set_after_successful_execution_version():
+    """ Test that the version of the database is set to the specified one after the migration is successfully
+    executed.
+    """
+
+    @migration(execute_only_if_version=1, set_after_successful_execution_version=33)
+    def test(_: Mock):
+        ...
+
+    db = Mock(version=1)
+    test(db)
+
+    assert db.version == 33
+
+
+def test_set_after_successful_execution_version_not_specified():
+    """ Test that if the version is not specified, the version of the database will be set to
+    execute_only_if_version + 1
+    """
 
     @migration(execute_only_if_version=1)
     def test(_: Mock):
-        return True
+        ...
 
     db = Mock(version=1)
-    assert test(db)
+    test(db)
+
     assert db.version == 2
 
 
-def test_migration_set_next_version_raise_an_exception():
+def test_set_after_successful_execution_raise_an_exception():
     """ Test that if an exception is raised during the migration, the version of the database is not changed."""
 
-    @migration(execute_only_if_version=1)
+    @migration(execute_only_if_version=1, set_after_successful_execution_version=33)
     def test(_: Mock):
         raise TypeError