feat: allow database reset

tod_attack_miner/cli.py

@@ -30,6 +30,11 @@ def main():
         default=None,
         help="If passed, filter TOD candidates that are {window-size} or more blocks apart",
     )
+    parser.add_argument(
+        "--reset-db",
+        action="store_true",
+        help="Delete data from previous runs before starting to mine",
+    )
     parser.add_argument("--postgres-user", type=str, default="postgres")
     parser.add_argument("--postgres-password", type=str, default="password")
     parser.add_argument("--postgres-host", type=str, default="localhost")
@@ -49,6 +54,8 @@ def main():
         if args.stats_only:
             print(json.dumps(miner.get_stats()))
         else:
+            if args.reset_db:
+                miner.reset_db()
             miner.fetch(int(args.from_block), int(args.to_block))
             miner.compute_skelcodes()
             miner.find_collisions()

tod_attack_miner/db/db.py

@@ -60,6 +60,25 @@ def _setup_tables(self):
                 )
             self._con.commit()
 
+    def reset(self):
+        with self._con.cursor() as cursor:
+            for name in _INDEXES:
+                cursor.execute(
+                    cast(
+                        psycopg.sql.SQL,
+                        f"DROP INDEX IF EXISTS {name}",
+                    )
+                )
+            for name in _TABLES:
+                cursor.execute(
+                    cast(
+                        psycopg.sql.SQL,
+                        f"DROP TABLE IF EXISTS {name}",
+                    )
+                )
+            self._con.commit()
+        self._setup_tables()
+
     def insert_prestate(self, block_number: int, tx_index: int, prestate: TxPrestate):
         with self._con.cursor() as cursor:
             accesses: list[tuple[int, int, str, ACCESS_TYPE, str, str]] = []

tod_attack_miner/miner/miner.py

@@ -11,6 +11,9 @@ def __init__(self, rpc: RPC, db: DB) -> None:
         self._filter_stats = {"candidates": {}, "filtered": {}}
         self._original_collisions = {}
 
+    def reset_db(self) -> None:
+        self.db.reset()
+
     def fetch(self, start: int, end: int) -> None:
         fetch_block_range(self.rpc, self.db, BlockRange(start, end))