Fix port script after storage changes. Add very simple (off by default) curses UI to see progress.

1 files changed, 409 insertions, 171 deletions

diff --git a/scripts/port_from_sqlite_to_postgres.py b/scripts/port_from_sqlite_to_postgres.py
index 4b3fd9e529..fc1603c1c9 100644
--- a/scripts/port_from_sqlite_to_postgres.py
+++ b/scripts/port_from_sqlite_to_postgres.py
@@ -20,38 +20,17 @@ from synapse.storage._base import LoggingTransaction, SQLBaseStore
 from synapse.storage.engines import create_engine
 
 import argparse
-import itertools
+import curses
 import logging
-import types
+import sys
+import time
+import traceback
 import yaml
 
 
 logger = logging.getLogger("port_from_sqlite_to_postgres")
 
 
-BINARY_COLUMNS = {
-    "event_content_hashes": ["hash"],
-    "event_reference_hashes": ["hash"],
-    "event_signatures": ["signature"],
-    "event_edge_hashes": ["hash"],
-    "events": ["content", "unrecognized_keys"],
-    "event_json": ["internal_metadata", "json"],
-    "application_services_txns": ["event_ids"],
-    "received_transactions": ["response_json"],
-    "sent_transactions": ["response_json"],
-    "server_tls_certificates": ["tls_certificate"],
-    "server_signature_keys": ["verify_key"],
-    "pushers": ["pushkey", "data"],
-    "user_filters": ["filter_json"],
-}
-
-UNICODE_COLUMNS = {
-    "events": ["content", "unrecognized_keys"],
-    "event_json": ["internal_metadata", "json"],
-    "users": ["password_hash"],
-}
-
-
 BOOLEAN_COLUMNS = {
     "events": ["processed", "outlier"],
     "rooms": ["is_public"],
@@ -91,7 +70,15 @@ APPEND_ONLY_TABLES = [
 ]
 
 
+end_error_exec_info = None
+
+
 class Store(object):
+    """This object is used to pull out some of the convenience API from the
+    Storage layer.
+
+    *All* database interactions should go through this object.
+    """
     def __init__(self, db_pool, engine):
         self.db_pool = db_pool
         self.database_engine = engine
@@ -130,11 +117,14 @@ class Store(object):
                                 continue
                         raise
             except Exception as e:
-                logger.debug("[TXN FAIL] {%s}", desc, e)
+                logger.debug("[TXN FAIL] {%s} %s", desc, e)
                 raise
 
         return self.db_pool.runWithConnection(r)
 
+    def execute(self, f):
+        return self.runInteraction(f.__name__, f)
+
     def insert_many_txn(self, txn, table, headers, rows):
         sql = "INSERT INTO %s (%s) VALUES (%s)" % (
             table,
@@ -152,205 +142,435 @@ class Store(object):
             raise
 
 
+class Progress(object):
+    """Used to report progress of the port
+    """
+    def __init__(self):
+        self.tables = {}
+
+        self.start_time = int(time.time())
+
+    def add_table(self, table, cur, size):
+        self.tables[table] = {
+            "start": cur,
+            "num_done": cur,
+            "total": size,
+            "perc": int(cur * 100 / size),
+        }
+
+    def update(self, table, num_done):
+        data = self.tables[table]
+        data["num_done"] = num_done
+        data["perc"] = int(num_done * 100 / data["total"])
+
+    def done(self):
+        pass
+
 
-def chunks(n):
-    for i in itertools.count(0, n):
-        yield range(i, i+n)
+class CursesProgress(Progress):
+    """Reports progress to a curses window
+    """
+    def __init__(self, stdscr):
+        self.stdscr = stdscr
 
+        curses.use_default_colors()
+        curses.curs_set(0)
 
-@defer.inlineCallbacks
-def handle_table(table, sqlite_store, postgres_store):
-    if table in APPEND_ONLY_TABLES:
-        # It's safe to just carry on inserting.
-        next_chunk = yield postgres_store._simple_select_one_onecol(
-            table="port_from_sqlite3",
-            keyvalues={"table_name": table},
-            retcol="rowid",
-            allow_none=True,
+        curses.init_pair(1, curses.COLOR_RED, -1)
+        curses.init_pair(2, curses.COLOR_GREEN, -1)
+
+        self.last_update = 0
+
+        self.finished = False
+
+        super(CursesProgress, self).__init__()
+
+    def update(self, table, num_done):
+        super(CursesProgress, self).update(table, num_done)
+
+        self.render()
+
+    def render(self, force=False):
+        now = time.time()
+
+        if not force and now - self.last_update < 0.2:
+            # reactor.callLater(1, self.render)
+            return
+
+        self.stdscr.clear()
+
+        rows, cols = self.stdscr.getmaxyx()
+
+        duration = int(now) - int(self.start_time)
+
+        minutes, seconds = divmod(duration, 60)
+        duration_str = '%02dm %02ds' % (minutes, seconds,)
+
+        if self.finished:
+            status = "Time spent: %s (Done!)" % (duration_str,)
+        else:
+            min_perc = min(
+                (v["num_done"] - v["start"]) * 100. / (v["total"] - v["start"])
+                if v["total"] - v["start"] else 100
+                for v in self.tables.values()
+            )
+            if min_perc > 0:
+                est_remaining = (int(now) - self.start_time) * 100 / min_perc
+                est_remaining_str = '%02dm %02ds remaining' % divmod(est_remaining, 60)
+            else:
+                est_remaining_str = "Unknown"
+            status = (
+                "Time spent: %s (est. remaining: %s)"
+                % (duration_str, est_remaining_str,)
+            )
+
+        self.stdscr.addstr(
+            0, 0,
+            status,
+            curses.A_BOLD,
         )
 
-        if next_chunk is None:
-            yield postgres_store._simple_insert(
-                table="port_from_sqlite3",
-                values={"table_name": table, "rowid": 0}
+        max_len = max([len(t) for t in self.tables.keys()])
+
+        left_margin = 5
+        middle_space = 1
+
+        items = self.tables.items()
+        items.sort(
+            key=lambda i: (i[1]["perc"], i[0]),
+        )
+
+        for i, (table, data) in enumerate(items):
+            if i + 2 >= rows:
+                break
+
+            perc = data["perc"]
+
+            color = curses.color_pair(2) if perc == 100 else curses.color_pair(1)
+
+            self.stdscr.addstr(
+                i+2, left_margin + max_len - len(table),
+                table,
+                curses.A_BOLD | color,
             )
 
-            next_chunk = 0
-    else:
-        def delete_all(txn):
-            txn.execute(
-                "DELETE FROM port_from_sqlite3 WHERE table_name = %s",
-                (table,)
+            size = 20
+
+            progress = "[%s%s]" % (
+                "#" * int(perc*size/100),
+                " " * (size - int(perc*size/100)),
             )
-            txn.execute("TRUNCATE %s CASCADE" % (table,))
-            postgres_store._simple_insert_txn(
-                txn,
-                table="port_from_sqlite3",
-                values={"table_name": table, "rowid": 0}
+
+            self.stdscr.addstr(
+                i+2, left_margin + max_len + middle_space,
+                "%s %3d%% (%d/%d)" % (progress, perc, data["num_done"], data["total"]),
             )
 
-        yield postgres_store.runInteraction(
-            "delete_non_append_only", delete_all
-        )
+        if self.finished:
+            self.stdscr.addstr(
+                self.rows-1, 0 ,
+                "Press any key to exit...",
+            )
 
-        next_chunk = 0
+        self.stdscr.refresh()
+        self.last_update = time.time()
 
-    logger.info("next_chunk for %s: %d", table, next_chunk)
+    def done(self):
+        self.finished = True
+        self.render(True)
+        self.stdscr.getch()
 
-    N = 5000
+    def on_prepare_sqlite(self):
+        self.stdscr.clear()
+        self.stdscr.addstr(
+            0, 0,
+            "Preparing SQLite database...",
+            curses.A_BOLD,
+        )
+        self.stdscr.refresh()
+
+    def on_prepare_postgres(self):
+        self.stdscr.clear()
+        self.stdscr.addstr(
+            0, 0,
+            "Preparing PostgreSQL database...",
+            curses.A_BOLD,
+        )
+        self.stdscr.refresh()
+
+    def fetching_tables(self):
+        self.stdscr.clear()
+        self.stdscr.addstr(
+            0, 0,
+            "Fetching tables...",
+            curses.A_BOLD,
+        )
+        self.stdscr.refresh()
+
+    def preparing_tables(self):
+        self.stdscr.clear()
+        self.stdscr.addstr(
+            0, 0,
+            "Preparing tables...",
+            curses.A_BOLD,
+        )
+        self.stdscr.refresh()
 
-    select = "SELECT rowid, * FROM %s WHERE rowid >= ? ORDER BY rowid LIMIT ?" % (table,)
 
-    uni_col_names = UNICODE_COLUMNS.get(table, [])
-    bool_col_names = BOOLEAN_COLUMNS.get(table, [])
-    bin_col_names = BINARY_COLUMNS.get(table, [])
+class TerminalProgress(Progress):
+    """Just prints progress to the terminal
+    """
+    def update(self, table, num_done):
+        super(TerminalProgress, self).update(table, num_done)
 
-    while True:
-        def r(txn):
-            txn.execute(select, (next_chunk, N,))
-            rows = txn.fetchall()
-            headers = [column[0] for column in txn.description]
+        data = self.tables[table]
 
-            return headers, rows
+        print "%s: %d%% (%d/%d)" % (
+            table, data["perc"],
+            data["num_done"], data["total"],
+        )
 
-        headers, rows = yield sqlite_store.runInteraction("select", r)
+    def on_prepare_sqlite(self):
+        print "Preparing SQLite database..."
 
-        logger.info("Got %d rows for %s", len(rows), table)
+    def on_prepare_postgres(self):
+        print "Preparing PostgreSQL database..."
 
-        if rows:
-            uni_cols = [i for i, h in enumerate(headers) if h in uni_col_names]
-            bool_cols = [i for i, h in enumerate(headers) if h in bool_col_names]
-            bin_cols = [i for i, h in enumerate(headers) if h in bin_col_names]
-            next_chunk = rows[-1][0] + 1
+    def fetching_tables(self):
+        print "Fetching tables..."
 
-            def conv(j, col):
-                if j in uni_cols:
-                    col = sqlite_store.database_engine.load_unicode(col)
-                if j in bool_cols:
-                    return bool(col)
+    def preparing_tables(self):
+        print "Preparing tables..."
 
-                if j in bin_cols:
-                    if isinstance(col, types.UnicodeType):
-                        col = buffer(col.encode("utf8"))
 
-                return col
+class Porter(object):
+    def __init__(self, **kwargs):
+        self.__dict__.update(kwargs)
 
-            for i, row in enumerate(rows):
-                rows[i] = tuple(
-                    postgres_store.database_engine.encode_parameter(
-                        conv(j, col)
-                    )
-                    for j, col in enumerate(row)
-                    if j > 0
-                )
+    @defer.inlineCallbacks
+    def handle_table(self, table):
+        if table in APPEND_ONLY_TABLES:
+            # It's safe to just carry on inserting.
+            next_chunk = yield self.postgres_store._simple_select_one_onecol(
+                table="port_from_sqlite3",
+                keyvalues={"table_name": table},
+                retcol="rowid",
+                allow_none=True,
+            )
 
-            def ins(txn):
-                postgres_store.insert_many_txn(txn, table, headers[1:], rows)
+            if next_chunk is None:
+                yield self.postgres_store._simple_insert(
+                    table="port_from_sqlite3",
+                    values={"table_name": table, "rowid": 1}
+                )
 
-                postgres_store._simple_update_one_txn(
+                next_chunk = 1
+        else:
+            def delete_all(txn):
+                txn.execute(
+                    "DELETE FROM port_from_sqlite3 WHERE table_name = %s",
+                    (table,)
+                )
+                txn.execute("TRUNCATE %s CASCADE" % (table,))
+                self.postgres_store._simple_insert_txn(
                     txn,
                     table="port_from_sqlite3",
-                    keyvalues={"table_name": table},
-                    updatevalues={"rowid": next_chunk},
+                    values={"table_name": table, "rowid": 0}
                 )
+            yield self.postgres_store.execute(delete_all)
 
-            yield postgres_store.runInteraction("insert_many", ins)
-        else:
+            next_chunk = 1
+
+        def get_table_size(txn):
+            txn.execute("SELECT count(*) FROM %s" % (table,))
+            size, = txn.fetchone()
+            return int(size)
+
+        table_size = yield self.sqlite_store.execute(get_table_size)
+        postgres_size = yield self.postgres_store.execute(get_table_size)
+
+        if not table_size:
             return
 
+        self.progress.add_table(table, postgres_size, table_size)
 
-def setup_db(db_config, database_engine):
-    db_conn = database_engine.module.connect(
-        **{
-            k: v for k, v in db_config.get("args", {}).items()
-            if not k.startswith("cp_")
-        }
-    )
+        select = (
+            "SELECT rowid, * FROM %s WHERE rowid >= ? ORDER BY rowid LIMIT ?"
+            % (table,)
+        )
 
-    database_engine.prepare_database(db_conn)
+        bool_col_names = BOOLEAN_COLUMNS.get(table, [])
 
-    db_conn.commit()
+        while True:
+            def r(txn):
+                txn.execute(select, (next_chunk, self.batch_size,))
+                rows = txn.fetchall()
+                headers = [column[0] for column in txn.description]
 
+                return headers, rows
 
-@defer.inlineCallbacks
-def main(sqlite_config, postgress_config):
-    try:
-        sqlite_db_pool = adbapi.ConnectionPool(
-            sqlite_config["name"],
-            **sqlite_config["args"]
-        )
+            headers, rows = yield self.sqlite_store.runInteraction("select", r)
 
-        postgres_db_pool = adbapi.ConnectionPool(
-            postgress_config["name"],
-            **postgress_config["args"]
-        )
+            if rows:
+                bool_cols = [
+                    i for i, h in enumerate(headers) if h in bool_col_names
+                ]
+                next_chunk = rows[-1][0] + 1
+
+                def conv(j, col):
+                    if j in bool_cols:
+                        return bool(col)
+                    return col
+
+                for i, row in enumerate(rows):
+                    rows[i] = tuple(
+                        self.postgres_store.database_engine.encode_parameter(
+                            conv(j, col)
+                        )
+                        for j, col in enumerate(row)
+                        if j > 0
+                    )
 
-        sqlite_engine = create_engine("sqlite3")
-        postgres_engine = create_engine("psycopg2")
+                def insert(txn):
+                    self.postgres_store.insert_many_txn(
+                        txn, table, headers[1:], rows
+                    )
 
-        sqlite_store = Store(sqlite_db_pool, sqlite_engine)
-        postgres_store = Store(postgres_db_pool, postgres_engine)
+                    self.postgres_store._simple_update_one_txn(
+                        txn,
+                        table="port_from_sqlite3",
+                        keyvalues={"table_name": table},
+                        updatevalues={"rowid": next_chunk},
+                    )
 
-        # Step 1. Set up databases.
-        logger.info("Preparing sqlite database...")
-        setup_db(sqlite_config, sqlite_engine)
+                yield self.postgres_store.execute(insert)
 
-        logger.info("Preparing postgres database...")
-        setup_db(postgress_config, postgres_engine)
+                postgres_size += len(rows)
 
-        # Step 2. Get tables.
-        logger.info("Fetching tables...")
-        tables = yield sqlite_store._simple_select_onecol(
-            table="sqlite_master",
-            keyvalues={
-                "type": "table",
-            },
-            retcol="name",
+                self.progress.update(table, postgres_size)
+            else:
+                return
+
+    def setup_db(self, db_config, database_engine):
+        db_conn = database_engine.module.connect(
+            **{
+                k: v for k, v in db_config.get("args", {}).items()
+                if not k.startswith("cp_")
+            }
         )
 
-        logger.info("Found %d tables", len(tables))
+        database_engine.prepare_database(db_conn)
 
-        def create_port_table(txn):
-            txn.execute(
-                "CREATE TABLE port_from_sqlite3 ("
-                " table_name varchar(100) NOT NULL UNIQUE,"
-                " rowid bigint NOT NULL"
-                ")"
-            )
+        db_conn.commit()
 
+    @defer.inlineCallbacks
+    def run(self):
         try:
-            yield postgres_store.runInteraction(
-                "create_port_table", create_port_table
+            sqlite_db_pool = adbapi.ConnectionPool(
+                self.sqlite_config["name"],
+                **self.sqlite_config["args"]
             )
-        except Exception as e:
-            logger.info("Failed to create port table: %s", e)
-
-        # Process tables.
-        yield defer.gatherResults(
-            [
-                handle_table(table, sqlite_store, postgres_store)
-                for table in tables
-                if table not in ["schema_version", "applied_schema_deltas"]
-                and not table.startswith("sqlite_")
-            ],
-            consumeErrors=True,
-        )
 
-    except:
-        logger.exception("")
-    finally:
-        reactor.stop()
+            postgres_db_pool = adbapi.ConnectionPool(
+                self.postgres_config["name"],
+                **self.postgres_config["args"]
+            )
+
+            sqlite_engine = create_engine("sqlite3")
+            postgres_engine = create_engine("psycopg2")
+
+            self.sqlite_store = Store(sqlite_db_pool, sqlite_engine)
+            self.postgres_store = Store(postgres_db_pool, postgres_engine)
+
+            # Step 1. Set up databases.
+            self.progress.on_prepare_sqlite()
+            self.setup_db(sqlite_config, sqlite_engine)
+
+            self.progress.on_prepare_postgres()
+            self.setup_db(postgres_config, postgres_engine)
+
+            # Step 2. Get tables.
+            self.progress.fetching_tables()
+            sqlite_tables = yield self.sqlite_store._simple_select_onecol(
+                table="sqlite_master",
+                keyvalues={
+                    "type": "table",
+                },
+                retcol="name",
+            )
+
+            postgres_tables = yield self.postgres_store._simple_select_onecol(
+                table="information_schema.tables",
+                keyvalues={
+                    "table_schema": "public",
+                },
+                retcol="distinct table_name",
+            )
+
+            tables = set(sqlite_tables) & set(postgres_tables)
+
+            self.progress.preparing_tables()
+
+            logger.info("Found %d tables", len(tables))
+
+            def create_port_table(txn):
+                txn.execute(
+                    "CREATE TABLE port_from_sqlite3 ("
+                    " table_name varchar(100) NOT NULL UNIQUE,"
+                    " rowid bigint NOT NULL"
+                    ")"
+                )
+
+            try:
+                yield self.postgres_store.runInteraction(
+                    "create_port_table", create_port_table
+                )
+            except Exception as e:
+                logger.info("Failed to create port table: %s", e)
+
+            # Process tables.
+            yield defer.gatherResults(
+                [
+                    self.handle_table(table)
+                    for table in tables
+                    if table not in ["schema_version", "applied_schema_deltas"]
+                    and not table.startswith("sqlite_")
+                ],
+                consumeErrors=True,
+            )
+
+            self.progress.done()
+        except:
+            global end_error_exec_info
+            end_error_exec_info = sys.exc_info()
+            logger.exception("")
+        finally:
+            reactor.stop()
 
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
+    parser.add_argument("-v", action='store_true')
+    parser.add_argument("--curses", action='store_true')
     parser.add_argument("--sqlite-database")
     parser.add_argument(
         "--postgres-config", type=argparse.FileType('r'),
     )
 
+    parser.add_argument("--batch-size", type=int, default=1000)
+
     args = parser.parse_args()
-    logging.basicConfig(level=logging.INFO)
+
+
+    logging_config = {
+        "level": logging.DEBUG if args.v else logging.INFO,
+        "format": "%(asctime)s - %(name)s - %(lineno)d - %(levelname)s - %(message)s"
+    }
+
+    if args.curses:
+        logging_config["filename"] = "port-synapse.log"
+
+    logging.basicConfig(**logging_config)
 
     sqlite_config = {
         "name": "sqlite3",
@@ -364,10 +584,28 @@ if __name__ == "__main__":
 
     postgres_config = yaml.safe_load(args.postgres_config)
 
-    reactor.callWhenRunning(
-        main,
-        sqlite_config=sqlite_config,
-        postgres_config=postgres_config,
-    )
+    def start(stdscr=None):
+        if stdscr:
+            progress = CursesProgress(stdscr)
+        else:
+            progress = TerminalProgress()
+
+        porter = Porter(
+            sqlite_config=sqlite_config,
+            postgres_config=postgres_config,
+            progress=progress,
+            batch_size=args.batch_size,
+        )
+
+        reactor.callWhenRunning(porter.run)
+
+        reactor.run()
+
+    if args.curses:
+        curses.wrapper(start)
+    else:
+        start()
 
-    reactor.run()
+    if end_error_exec_info:
+        exc_type, exc_value, exc_traceback = end_error_exec_info
+        traceback.print_exception(exc_type, exc_value, exc_traceback)