Files
mautrix-telegram/mautrix_telegram/scripts/dbms_migrate/__main__.py
T
2019-03-21 13:48:53 +02:00

82 lines
2.8 KiB
Python

import argparse
import sqlalchemy as sql
from sqlalchemy import orm
from sqlalchemy.ext.declarative import declarative_base
from alchemysession import AlchemySessionContainer
parser = argparse.ArgumentParser(description="mautrix-telegram dbms migration script",
prog="python -m mautrix_telegram.scripts.dbms_migrate")
parser.add_argument("-f", "--from-url", type=str, required=True, metavar="<url>",
help="the old database path")
parser.add_argument("-t", "--to-url", type=str, required=True, metavar="<url>",
help="the new database path")
parser.add_argument("-v", "--verbose", action="store_true", help="Verbose logs while migrating")
args = parser.parse_args()
verbose = args.verbose or False
def log(message, end="\n"):
if verbose:
print(message, end=end, flush=True)
def connect(to):
import mautrix_telegram.db.base as base
base.Base = declarative_base(cls=base.BaseBase)
from mautrix_telegram.db import (Portal, Message, UserPortal, User, RoomState, UserProfile,
Contact, Puppet, BotChat, TelegramFile)
db_engine = sql.create_engine(to)
db_factory = orm.sessionmaker(bind=db_engine)
db_session = orm.scoped_session(db_factory) # type: orm.Session
base.Base.metadata.bind = db_engine
session_container = AlchemySessionContainer(engine=db_engine, session=db_session,
table_base=base.Base, table_prefix="telethon_",
manage_tables=False)
return db_session, {
"Version": session_container.Version,
"Session": session_container.Session,
"Entity": session_container.Entity,
"SentFile": session_container.SentFile,
"UpdateState": session_container.UpdateState,
"Portal": Portal,
"Message": Message,
"Puppet": Puppet,
"User": User,
"UserPortal": UserPortal,
"RoomState": RoomState,
"UserProfile": UserProfile,
"Contact": Contact,
"BotChat": BotChat,
"TelegramFile": TelegramFile,
}
log("Connecting to old database")
session, tables = connect(args.from_url)
data = {}
for name, table in tables.items():
log("Reading table {name}...".format(name=name), end=" ")
data[name] = session.query(table).all()
log("Done!")
log("Connecting to new database")
session, tables = connect(args.to_url)
for name, table in tables.items():
log("Writing table {name}".format(name=name), end="")
length = len(data[name])
n = 0
for row in data[name]:
session.merge(row)
n += 5
if n >= length:
log(".", end="")
n = 0
log(" Done!")
log("Committing changes to database...", end=" ")
session.commit()
log("Done!")