feat: improve load-from-logs CLI command

hathor/cli/load_from_logs.py

@@ -12,56 +12,46 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import re
 import sys
 from argparse import ArgumentParser, FileType
 
+from twisted.internet.defer import Deferred
+from twisted.internet.task import deferLater
+
 from hathor.cli.run_node import RunNode
 
 
 class LoadFromLogs(RunNode):
-    def start_manager(self) -> None:
-        pass
-
-    def register_signal_handlers(self) -> None:
-        pass
-
     @classmethod
     def create_parser(cls) -> ArgumentParser:
         parser = super().create_parser()
         parser.add_argument('--log-dump', type=FileType('r', encoding='UTF-8'), default=sys.stdin, nargs='?',
-                            help='Where to read logs from, defaults to stdin.')
+                            help='Where to read logs from, defaults to stdin. Should be pre-parsed with parse-logs.')
         return parser
 
-    def prepare(self, *, register_resources: bool = True) -> None:
-        super().prepare(register_resources=False)
-
     def run(self) -> None:
+        self.reactor.callLater(0, lambda: Deferred.fromCoroutine(self._load_from_logs()))
+        super().run()
+
+    async def _load_from_logs(self) -> None:
         from hathor.conf.get_settings import get_global_settings
         from hathor.transaction.vertex_parser import VertexParser
         settings = get_global_settings()
         parser = VertexParser(settings=settings)
 
-        pattern = r'new (tx|block)    .*bytes=([^ ]*) '
-        pattern = r'new (tx|block)    .*bytes=([^ ]*) '
-        compiled_pattern = re.compile(pattern)
-
         while True:
             line_with_break = self._args.log_dump.readline()
             if not line_with_break:
                 break
-            line = line_with_break.strip()
-
-            matches = compiled_pattern.findall(line)
-            if len(matches) == 0:
+            if line_with_break.startswith('//'):
                 continue
-
-            assert len(matches) == 1
-            _, vertex_bytes_hex = matches[0]
-
-            vertex_bytes = bytes.fromhex(vertex_bytes_hex)
+            line = line_with_break.strip()
+            vertex_bytes = bytes.fromhex(line)
             vertex = parser.deserialize(vertex_bytes)
-            self.manager.on_new_tx(vertex)
+            await deferLater(self.reactor, 0, self.manager.on_new_tx, vertex)
+
+        self.manager.connections.disconnect_all_peers(force=True)
+        self.reactor.fireSystemEvent('shutdown')
 
 
 def main():

hathor/cli/main.py

@@ -48,6 +48,7 @@ def __init__(self) -> None:
             oracle_create_key,
             oracle_encode_data,
             oracle_get_pubkey,
+            parse_logs,
             peer_id,
             quick_test,
             replay_logs,
@@ -98,8 +99,10 @@ def __init__(self) -> None:
         self.add_cmd('dev', 'events_simulator', events_simulator, 'Simulate integration events via websocket')
         self.add_cmd('dev', 'x-export', db_export, 'EXPERIMENTAL: Export database to a simple format.')
         self.add_cmd('dev', 'x-import', db_import, 'EXPERIMENTAL: Import database from exported format.')
-        self.add_cmd('dev', 'replay-logs', replay_logs, 'EXPERIMENTAL: re-play json logs as console printted')
-        self.add_cmd('dev', 'load-from-logs', load_from_logs, 'Load vertices as they are found in a log dump')
+        self.add_cmd('dev', 'replay-logs', replay_logs, 'EXPERIMENTAL: re-play json logs as console printed')
+        self.add_cmd('dev', 'load-from-logs', load_from_logs,
+                     'Load vertices as they are found in a log dump that was parsed with parse-logs')
+        self.add_cmd('dev', 'parse-logs', parse_logs, 'Parse a log dump to use it with load-from-logs')
 
     def add_cmd(self, group: str, cmd: str, module: ModuleType, short_description: Optional[str] = None) -> None:
         self.command_list[cmd] = module

hathor/cli/parse_logs.py

@@ -0,0 +1,99 @@
+# Copyright 2024 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import re
+import sys
+from argparse import FileType
+from io import TextIOWrapper
+from typing import Iterator
+
+
+def main() -> None:
+    """
+    Parse logs from a dump file (either as json or plain logs) into a file with only vertex hex bytes.
+    The logs must be generated with --log-vertex-bytes. Then, use load-from-logs to run a full node from this file.
+    """
+    from hathor.cli.util import create_parser
+    parser = create_parser()
+    file_args = parser.add_mutually_exclusive_group(required=True)
+    file_args.add_argument(
+        '--json-logs-file',
+        type=FileType('r', encoding='UTF-8'),
+        help='Where to read json logs from.',
+    )
+    file_args.add_argument(
+        '--plain-logs-file',
+        type=FileType('r', encoding='UTF-8'),
+        help='Where to read plain logs from.',
+    )
+    parser.add_argument(
+        '--output-file',
+        type=FileType('w', encoding='UTF-8'),
+        required=True,
+        help='Output file.',
+    )
+    args = parser.parse_args(sys.argv[1:])
+    assert isinstance(args.output_file, TextIOWrapper)
+
+    vertex_iter: Iterator[str]
+    if args.json_logs_file is not None:
+        assert isinstance(args.json_logs_file, TextIOWrapper)
+        print('parsing json logs file...')
+        vertex_iter = _parse_json_logs(args.json_logs_file)
+    else:
+        assert isinstance(args.plain_logs_file, TextIOWrapper)
+        print('parsing plain logs file...')
+        vertex_iter = _parse_plain_logs(args.plain_logs_file)
+
+    print('writing to output file...')
+    for vertex in vertex_iter:
+        args.output_file.write(vertex + '\n')
+    print('done')
+
+
+def _parse_json_logs(file: TextIOWrapper) -> Iterator[str]:
+    while True:
+        line = file.readline()
+        if not line:
+            break
+
+        json_dict = json.loads(line)
+        event = json_dict.get('event')
+        if not event:
+            return
+
+        if event in ('new block', 'new tx'):
+            vertex_bytes = json_dict.get('bytes')
+            assert vertex_bytes is not None, 'logs should be generated with --log-vertex-bytes'
+            yield vertex_bytes
+
+
+def _parse_plain_logs(file: TextIOWrapper) -> Iterator[str]:
+    pattern = r'new (tx|block)    .*bytes=([^ ]*) '
+    compiled_pattern = re.compile(pattern)
+
+    while True:
+        line_with_break = file.readline()
+        if not line_with_break:
+            break
+        line = line_with_break.strip()
+
+        matches = compiled_pattern.findall(line)
+        if len(matches) == 0:
+            continue
+
+        assert len(matches) == 1
+        _, vertex_bytes_hex = matches[0]
+        yield vertex_bytes_hex