File size: 2,100 Bytes
a1a9773 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import argparse
import os
import sqlite3
import lmdb
def get_map_size(lmdb_path):
data_path = os.path.join(lmdb_path, "data.mdb")
try:
size = os.path.getsize(data_path)
except OSError:
return 1048576 * 2
return max(size * 2, 1048576 * 2)
def open_sqlite(path):
conn = sqlite3.connect(path)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS pixif_cache (
post_id TEXT PRIMARY KEY,
url TEXT
)
"""
)
conn.commit()
return conn
def upsert_rows(conn, rows):
if not rows:
return
conn.executemany(
"""
INSERT INTO pixif_cache (post_id, url)
VALUES (?, ?)
ON CONFLICT(post_id) DO UPDATE SET url = excluded.url
""",
rows,
)
def migrate(lmdb_path, sqlite_path, batch_size):
map_size = get_map_size(lmdb_path)
lmdb_env = lmdb.open(
lmdb_path,
subdir=True,
readonly=True,
lock=False,
map_size=map_size,
)
conn = open_sqlite(sqlite_path)
rows = []
with lmdb_env.begin() as txn:
cursor = txn.cursor()
for key, value in cursor:
post_id = key.decode("utf-8")
url = value.decode("utf-8") if value else ""
rows.append((post_id, url))
if len(rows) >= batch_size:
with conn:
upsert_rows(conn, rows)
rows.clear()
if rows:
with conn:
upsert_rows(conn, rows)
conn.close()
lmdb_env.close()
def main():
parser = argparse.ArgumentParser(description="Migrate LMDB cache to SQLite.")
parser.add_argument("--lmdb", default="db", help="Path to LMDB directory.")
parser.add_argument("--sqlite", default="db.sqlite", help="Path to SQLite file.")
parser.add_argument("--batch-size", type=int, default=1000, help="Rows per batch insert.")
args = parser.parse_args()
os.chdir(os.path.dirname(os.path.abspath(__file__)))
migrate(args.lmdb, args.sqlite, args.batch_size)
if __name__ == "__main__":
main()
|