silverstream/tools/migrate_database.py
Casper V. Kristensen 49eb78d5c0
Publish.
2019-04-19 02:33:31 +02:00

44 lines
1.4 KiB
Python

import dbm.gnu as dbm
import os
import sys
from pathlib import Path
from typing import List, Tuple, Iterator
sys.path.insert(0, os.path.dirname(os.path.dirname((os.path.abspath(__file__))))) # magic import-fixer
from silverstream.bittorrent.database import TorrentDatabase
from silverstream.bittorrent.peer import TorrentPeer
from silverstream.util import grouper
class OldPeerDatabaseStub:
def __init__(self, filepath) -> None:
self.dbm = dbm.open(filepath, "rf") # read/fast mode
def __iter__(self) -> Iterator[Tuple[bytes, List[TorrentPeer]]]:
info_hash = self.dbm.firstkey()
while info_hash is not None:
yield info_hash, self.get(info_hash)
info_hash = self.dbm.nextkey(info_hash)
def get(self, info_hash: bytes) -> List[TorrentPeer]:
compact_peers = grouper(self.dbm.get(info_hash), 6)
return [TorrentPeer.from_compact(bytes(compact_peer)) for compact_peer in compact_peers]
try:
old_filepath = Path(sys.argv[1])
except IndexError:
exit(f"Usage: {__file__} <path to old database>")
old = OldPeerDatabaseStub(str(old_filepath))
new = TorrentDatabase(filepath=old_filepath.with_suffix(".sqlite"))
num_keys = len(old.dbm.keys())
num_hashes = num_peers = 0
for info_hash, peers in old:
num_hashes += 1
for peer in peers:
num_peers += 1
new.add_peer(info_hash, peer)
print(f"{num_hashes} hashes, {num_peers} peers done ({num_hashes / num_keys:.2%})")