fix: aimrt_cli trans will first repair the aimrtbag (AimRT#112)

* fix(bag): optimize bag file processing workflow - Add handling for missing files in PlaybackAction, ignore non-existent files - Add bag_recover.py script for repairing corrupted bag files - Modify rosbag_trans.py to add file existence check and repair steps - Add index creation and transaction handling in AimrtbagToRos2 to improve conversion efficiency * build(bag_recover): update version * fix: rename the file * fix: opt the code
postrantor · Nov 25, 2024 · b41416c · b41416c
1 parent 18d45db
commit b41416c
Show file tree

Hide file tree

Showing 4 changed files with 66 additions and 12 deletions.
diff --git a/src/plugins/record_playback_plugin/playback_action.cc b/src/plugins/record_playback_plugin/playback_action.cc
@@ -163,13 +163,15 @@ void PlaybackAction::Initialize(YAML::Node options_node) {
   AIMRT_CHECK_ERROR_THROW(!metadata_.files.empty(),
                           "Empty bag! bag path: {}", options_.bag_path);
 
-  for (auto& item : metadata_.files) {
-    const auto db_file_path = std::filesystem::path(options_.bag_path) / item.path;
-
-    AIMRT_CHECK_ERROR_THROW(
-        std::filesystem::exists(db_file_path) && std::filesystem::is_regular_file(db_file_path),
-        "Can not find bag file '{}' in bag path '{}'.", db_file_path.string(), options_.bag_path);
-  }
+  metadata_.files.erase(std::remove_if(metadata_.files.begin(), metadata_.files.end(), [this](const auto& item) {
+                          const auto db_file_path = std::filesystem::path(options_.bag_path) / item.path;
+                          if (!std::filesystem::exists(db_file_path) || !std::filesystem::is_regular_file(db_file_path)) {
+                            AIMRT_WARN("Can not find bag file '{}' in bag path '{}', this file will be ignored.", db_file_path.string(), options_.bag_path);
+                            return true;
+                          }
+                          return false;
+                        }),
+                        metadata_.files.end());
 
   options_node = options_;
 }

diff --git a/src/tools/aimrt_cli/aimrt_cli/trans/bag_repair.py b/src/tools/aimrt_cli/aimrt_cli/trans/bag_repair.py
@@ -0,0 +1,39 @@
+import subprocess
+import shutil
+import os
+
+
+def repair_bag(bag_path: str):
+    try:
+        journal_files = [f for f in os.listdir(bag_path) if f.endswith("-journal")]
+        if not journal_files:
+            print(f"there is no journal file in  {bag_path}")
+            return
+
+        print(f"detect {len(journal_files)} journal files, start to repair")
+
+        for journal_file in journal_files:
+            journal_file_name = journal_file.split("-")[0]
+            journal_path = os.path.join(bag_path, journal_file_name)
+            print(f"journal_path: {journal_path}")
+
+            backup_path = f"{journal_path}.bak"
+            recover_path = os.path.join(bag_path, "recovered.db3")
+            shutil.copy2(journal_path, backup_path)
+
+            try:
+                cmd = f'sqlite3 "{journal_path}" ".recover" | sqlite3 "{recover_path}"'
+                subprocess.run(cmd, shell=True, check=True)
+
+                if os.path.exists(recover_path):
+                    os.replace(recover_path, journal_path)
+                    print(f"{journal_path} repair done \n")
+
+            except subprocess.CalledProcessError as e:
+                print(f"database repair failed: {e.stderr.decode()}")
+            finally:
+                if os.path.exists(backup_path):
+                    os.remove(backup_path)
+
+    except Exception as e:
+        print(f"repair failed: {str(e)}\n")
diff --git a/src/tools/aimrt_cli/aimrt_cli/trans/rosbag_trans.py b/src/tools/aimrt_cli/aimrt_cli/trans/rosbag_trans.py
@@ -2,6 +2,7 @@
 # AimRT is licensed under Mulan PSL v2.
 
 from aimrt_cli.trans import TransBase
+from aimrt_cli.trans.bag_repair import repair_bag
 import os
 import sqlite3
 import shutil
@@ -181,9 +182,14 @@ def trans_single_bag(self, topic_map: dict):
         self.parse_yaml()
         print(f"there are {len(self.files_list)} db files in {self.input_dir}")
         for db_path in self.files_list:
-            trans_path = Path(self.output_dir) / db_path['path']
+            input_trans_path = Path(self.input_dir) / db_path['path']
+            if not input_trans_path.exists():
+                print(f"    trans_path: {input_trans_path} not found, skip")
+                continue
             self.trans_single_db(Path(self.input_dir) / db_path['path'], topic_map)
-            print(f"    trans_path: {trans_path} done")
+            output_trans_path = Path(self.output_dir) / db_path['path']
+            print(f"    trans_path: {output_trans_path} done")
+
         print(f"all db files in {self.input_dir} done\n")
 
 
@@ -357,6 +363,9 @@ def sort_db_data(self):
                 data        BLOB NOT NULL)
             """)
 
+            self.cursor.execute("CREATE INDEX idx_timestamp ON messages(timestamp)")
+            self.cursor.execute("BEGIN TRANSACTION")
+
             self.cursor.execute("""
             INSERT INTO messages_temp (topic_id, timestamp, data)
             SELECT topic_id, timestamp, data
@@ -368,6 +377,7 @@ def sort_db_data(self):
 
             self.cursor.execute("ALTER TABLE messages_temp RENAME TO messages")
 
+            self.cursor.execute("COMMIT")
             self.conn.commit()
 
         except Exception as e:
@@ -383,6 +393,9 @@ def trans(self):
             self.insert_topics_table()
 
             for input_dir in self.input_dir:
+                print("start repairing broken packets")
+                repair_bag(input_dir)
+
                 single_bag_trans = SingleBagTrans(
                     input_dir,
                     self.output_dir,

diff --git a/src/tools/aimrt_cli/setup.py b/src/tools/aimrt_cli/setup.py
@@ -8,16 +8,16 @@
 setup(
     name=package_name,
     app='aimrt_cli',
-    version='0.0.4',
+    version='0.0.5',
     packages=find_packages(exclude=['test']),
     install_requires=[
         'PyYAML>=5.4.1',
         'Jinja2>=3.0.3',
         'pyinstaller>=6.1.0',
         'autopep8>=1.6.0',
     ],
-    author='Yu Xi',
-    author_email='yuxi@zhiyuan-robot.com',
+    author=['Yu Xi', 'Yu Guanlin'],
+    author_email=['yuxi@zhiyuan-robot.com', 'yuguanlin@agibot.com'],
     description='AimRT application python tools',
     license='',
     entry_points={