improve migration

Author: glevco

hathor/transaction/storage/cache_storage.py

@@ -15,6 +15,7 @@
 from collections import OrderedDict
 from typing import Any, Iterator, Optional
 
+from structlog.stdlib import BoundLogger
 from twisted.internet import threads
 from typing_extensions import override
 
@@ -25,7 +26,6 @@
 from hathor.transaction.storage.migrations import MigrationState
 from hathor.transaction.storage.transaction_storage import BaseTransactionStorage
 from hathor.transaction.storage.tx_allow_scope import TxAllowScope
-from hathor.types import VertexId
 
 
 class TransactionCacheStorage(BaseTransactionStorage):
@@ -253,5 +253,5 @@ def flush(self):
         self._flush_to_storage(self.dirty_txs.copy())
 
     @override
-    def iter_all_raw_metadata(self) -> Iterator[tuple[VertexId, dict[str, Any]]]:
-        return self.store.iter_all_raw_metadata()
+    def migrate_static_metadata(self, log: BoundLogger) -> None:
+        return self.store.migrate_static_metadata(log)

hathor/transaction/storage/memory_storage.py

@@ -14,6 +14,7 @@
 
 from typing import Any, Iterator, Optional, TypeVar
 
+from structlog.stdlib import BoundLogger
 from typing_extensions import override
 
 from hathor.conf.settings import HathorSettings
@@ -23,7 +24,6 @@
 from hathor.transaction.storage.migrations import MigrationState
 from hathor.transaction.storage.transaction_storage import BaseTransactionStorage
 from hathor.transaction.transaction_metadata import TransactionMetadata
-from hathor.types import VertexId
 
 _Clonable = TypeVar('_Clonable', BaseTransaction, TransactionMetadata)
 
@@ -128,7 +128,7 @@ def get_value(self, key: str) -> Optional[str]:
         return self.attributes.get(key)
 
     @override
-    def iter_all_raw_metadata(self) -> Iterator[tuple[VertexId, dict[str, Any]]]:
-        # This method is only ever used by the `migrate_static_metadata` migration, and therefore it is not necessary
-        # to implement for the memory storage.
+    def migrate_static_metadata(self, log: BoundLogger) -> None:
+        # This method is only ever used by the `migrate_static_metadata` migration, and therefore must not be
+        # implemented for the memory storage.
         raise NotImplementedError

hathor/transaction/storage/migrations/migrate_static_metadata.py

@@ -13,13 +13,12 @@
 #  limitations under the License.
 
 from typing import TYPE_CHECKING
-from unittest.mock import Mock
 
 from structlog import get_logger
 
 from hathor.conf.get_settings import get_global_settings
-from hathor.transaction import BaseTransaction
-from hathor.transaction.static_metadata import BlockStaticMetadata, TransactionStaticMetadata, VertexStaticMetadata
+from hathor.transaction import Block, Transaction
+from hathor.transaction.static_metadata import BlockStaticMetadata, TransactionStaticMetadata
 from hathor.transaction.storage.migrations import BaseMigration
 from hathor.util import progress
 
@@ -40,48 +39,29 @@ def run(self, storage: 'TransactionStorage') -> None:
         """This migration takes attributes from existing vertex metadata and saves them as static metadata."""
         log = logger.new()
         settings = get_global_settings()
-        # We have to iterate over metadata instead of vertices because the storage doesn't allow us to get a vertex if
-        # its static metadata is not set. We also use raw dict metadata because `metadata.create_from_json()` doesn't
-        # include attributes that should be static, which are exactly the ones we need for this migration.
-        metadata_iter = storage.iter_all_raw_metadata()
 
-        for vertex_id, raw_metadata in progress(metadata_iter, log=log, total=None):
-            height = raw_metadata['height']
-            min_height = raw_metadata['min_height']
-            bit_counts = raw_metadata.get('feature_activation_bit_counts')
+        # First we migrate static metadata using the storage itself since it uses internal structures.
+        log.info('creating static metadata...')
+        storage.migrate_static_metadata(log)
 
-            assert isinstance(height, int)
-            assert isinstance(min_height, int)
+        # Now that static metadata is set, we can use the topological iterator normally
+        log.info('removing old metadata and validating...')
+        topological_iter = storage.topological_iterator()
 
-            static_metadata: VertexStaticMetadata
-            is_block = (vertex_id == settings.GENESIS_BLOCK_HASH or height != 0)
-
-            if is_block:
-                assert isinstance(bit_counts, list)
-                for item in bit_counts:
-                    assert isinstance(item, int)
+        for vertex in progress(topological_iter, log=log, total=None):
+            # We re-save the vertex's metadata so it's serialized with the new `to_bytes()` method, excluding fields
+            # that were migrated.
+            storage.save_transaction(vertex, only_metadata=True)
 
-                static_metadata = BlockStaticMetadata(
-                    height=height,
-                    min_height=min_height,
-                    feature_activation_bit_counts=bit_counts,
-                    feature_states={},  # This will be populated in the next PR
+            # We re-create the static metadata from scratch and compare it with the value that was created by the
+            # migration above, as a sanity check.
+            if isinstance(vertex, Block):
+                assert vertex.static_metadata == BlockStaticMetadata.create_from_storage(
+                    vertex, settings, storage
                 )
-            else:
-                assert bit_counts is None or bit_counts == []
-                static_metadata = TransactionStaticMetadata(
-                    min_height=min_height
+            elif isinstance(vertex, Transaction):
+                assert vertex.static_metadata == TransactionStaticMetadata.create_from_storage(
+                    vertex, settings, storage
                 )
-
-            # We create a fake vertex with just the hash and static metadata, so we can use the existing
-            # `storage._save_static_metadata()` instead of having to create an unsafe storage API that takes those
-            # two arguments.
-            mock_vertex = Mock(spec_set=BaseTransaction)
-            mock_vertex.hash = vertex_id
-            mock_vertex.static_metadata = static_metadata
-            storage._save_static_metadata(mock_vertex)
-
-            # Now we can take the actual vertex from the storage and save its metadata. It'll be serialized with the
-            # new `to_bytes()` method, excluding fields that were migrated.
-            vertex = storage.get_transaction(vertex_id)
-            storage.save_transaction(vertex, only_metadata=True)
+            else:
+                raise NotImplementedError

hathor/transaction/storage/rocksdb_storage.py

@@ -12,21 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING, Any, Iterator, Optional
+from typing import TYPE_CHECKING, Iterator, Optional
 
 from structlog import get_logger
+from structlog.stdlib import BoundLogger
 from typing_extensions import override
 
 from hathor.conf.settings import HathorSettings
 from hathor.indexes import IndexesManager
 from hathor.storage import RocksDBStorage
-from hathor.transaction.static_metadata import VertexStaticMetadata
+from hathor.transaction.static_metadata import BlockStaticMetadata, TransactionStaticMetadata, VertexStaticMetadata
 from hathor.transaction.storage.exceptions import TransactionDoesNotExist
 from hathor.transaction.storage.migrations import MigrationState
 from hathor.transaction.storage.transaction_storage import BaseTransactionStorage
 from hathor.transaction.vertex_parser import VertexParser
-from hathor.types import VertexId
-from hathor.util import json_loadb
+from hathor.util import json_loadb, progress
 
 if TYPE_CHECKING:
     import rocksdb
@@ -235,9 +235,41 @@ def get_value(self, key: str) -> Optional[str]:
             return data.decode()
 
     @override
-    def iter_all_raw_metadata(self) -> Iterator[tuple[VertexId, dict[str, Any]]]:
-        items = self._db.iteritems(self._cf_meta)
-        items.seek_to_first()
+    def migrate_static_metadata(self, log: BoundLogger) -> None:
+        metadata_iter = self._db.iteritems(self._cf_meta)
+        metadata_iter.seek_to_first()
+
+        # We have to iterate over metadata instead of vertices because the storage doesn't allow us to get a vertex if
+        # its static metadata is not set. We also use raw dict metadata because `metadata.create_from_json()` doesn't
+        # include attributes that should be static, which are exactly the ones we need for this migration.
+        for (_, vertex_id), metadata_bytes in progress(metadata_iter, log=log, total=None):
+            raw_metadata = json_loadb(metadata_bytes)
+            height = raw_metadata['height']
+            min_height = raw_metadata['min_height']
+            bit_counts = raw_metadata.get('feature_activation_bit_counts')
+
+            assert isinstance(height, int)
+            assert isinstance(min_height, int)
+
+            static_metadata: VertexStaticMetadata
+            is_block = (vertex_id == self._settings.GENESIS_BLOCK_HASH or height != 0)
+
+            if is_block:
+                assert isinstance(bit_counts, list)
+                for item in bit_counts:
+                    assert isinstance(item, int)
+
+                static_metadata = BlockStaticMetadata(
+                    height=height,
+                    min_height=min_height,
+                    feature_activation_bit_counts=bit_counts,
+                    feature_states={},  # This will be populated in the next PR
+                )
+            else:
+                assert bit_counts is None or bit_counts == []
+                static_metadata = TransactionStaticMetadata(
+                    min_height=min_height
+                )
 
-        for (_, vertex_id), metadata_bytes in items:
-            yield vertex_id, json_loadb(metadata_bytes)
+            # Save it manually to the CF
+            self._db.put((self._cf_static_meta, vertex_id), static_metadata.to_bytes())

hathor/transaction/storage/transaction_storage.py

@@ -22,6 +22,7 @@
 
 from intervaltree.interval import Interval
 from structlog import get_logger
+from structlog.stdlib import BoundLogger
 
 from hathor.conf.settings import HathorSettings
 from hathor.execution_manager import ExecutionManager
@@ -1126,10 +1127,9 @@ def get_block(self, block_id: VertexId) -> Block:
         return block
 
     @abstractmethod
-    def iter_all_raw_metadata(self) -> Iterator[tuple[VertexId, dict[str, Any]]]:
+    def migrate_static_metadata(self, log: BoundLogger) -> None:
         """
-        Iterate over all vertex metadata from this storage, as raw dicts. This is only used for the
-        `migrate_static_metadata` migration.
+        Migrate metadata attributes to static metadata. This is only used for the `migrate_static_metadata` migration.
         """
         raise NotImplementedError
 

hathor/transaction/transaction_metadata.py

@@ -298,11 +298,15 @@ def to_bytes(self) -> bytes:
 
         # The `to_json()` method includes these fields for backwards compatibility with APIs, but since they're not
         # part of metadata, they should not be serialized.
-        del json_dict['height']
-        del json_dict['min_height']
-        del json_dict['feature_activation_bit_counts']
+        if 'height' in json_dict:
+            del json_dict['height']
+        if 'min_height' in json_dict:
+            del json_dict['min_height']
+        if 'feature_activation_bit_counts' in json_dict:
+            del json_dict['feature_activation_bit_counts']
         # TODO: This one has not been migrated yet, but will be in the next PR
-        # del json_dict['feature_states']
+        # if 'feature_states' in json_dict:
+        #     del json_dict['feature_states']
 
         return json_dumpb(json_dict)