123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991 |
- #!/usr/bin/env python
- # -*- coding: utf-8 -*-
- # Copyright 2015, 2016 OpenMarket Ltd
- # Copyright 2018 New Vector Ltd
- #
- # Licensed under the Apache License, Version 2.0 (the "License");
- # you may not use this file except in compliance with the License.
- # You may obtain a copy of the License at
- #
- # http://www.apache.org/licenses/LICENSE-2.0
- #
- # Unless required by applicable law or agreed to in writing, software
- # distributed under the License is distributed on an "AS IS" BASIS,
- # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- # See the License for the specific language governing permissions and
- # limitations under the License.
- from twisted.internet import defer, reactor
- from twisted.enterprise import adbapi
- from synapse.storage._base import LoggingTransaction, SQLBaseStore
- from synapse.storage.engines import create_engine
- from synapse.storage.prepare_database import prepare_database
- import argparse
- import curses
- import logging
- import sys
- import time
- import traceback
- import yaml
- logger = logging.getLogger("synapse_port_db")
- BOOLEAN_COLUMNS = {
- "events": ["processed", "outlier", "contains_url"],
- "rooms": ["is_public"],
- "event_edges": ["is_state"],
- "presence_list": ["accepted"],
- "presence_stream": ["currently_active"],
- "public_room_list_stream": ["visibility"],
- "device_lists_outbound_pokes": ["sent"],
- "users_who_share_rooms": ["share_private"],
- "groups": ["is_public"],
- "group_rooms": ["is_public"],
- "group_users": ["is_public", "is_admin"],
- "group_summary_rooms": ["is_public"],
- "group_room_categories": ["is_public"],
- "group_summary_users": ["is_public"],
- "group_roles": ["is_public"],
- "local_group_membership": ["is_publicised", "is_admin"],
- }
- APPEND_ONLY_TABLES = [
- "event_content_hashes",
- "event_reference_hashes",
- "event_signatures",
- "event_edge_hashes",
- "events",
- "event_json",
- "state_events",
- "room_memberships",
- "feedback",
- "topics",
- "room_names",
- "rooms",
- "local_media_repository",
- "local_media_repository_thumbnails",
- "remote_media_cache",
- "remote_media_cache_thumbnails",
- "redactions",
- "event_edges",
- "event_auth",
- "received_transactions",
- "sent_transactions",
- "transaction_id_to_pdu",
- "users",
- "state_groups",
- "state_groups_state",
- "event_to_state_groups",
- "rejections",
- "event_search",
- "presence_stream",
- "push_rules_stream",
- "current_state_resets",
- "ex_outlier_stream",
- "cache_invalidation_stream",
- "public_room_list_stream",
- "state_group_edges",
- "stream_ordering_to_exterm",
- ]
- 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
- _simple_insert_txn = SQLBaseStore.__dict__["_simple_insert_txn"]
- _simple_insert = SQLBaseStore.__dict__["_simple_insert"]
- _simple_select_onecol_txn = SQLBaseStore.__dict__["_simple_select_onecol_txn"]
- _simple_select_onecol = SQLBaseStore.__dict__["_simple_select_onecol"]
- _simple_select_one = SQLBaseStore.__dict__["_simple_select_one"]
- _simple_select_one_txn = SQLBaseStore.__dict__["_simple_select_one_txn"]
- _simple_select_one_onecol = SQLBaseStore.__dict__["_simple_select_one_onecol"]
- _simple_select_one_onecol_txn = SQLBaseStore.__dict__[
- "_simple_select_one_onecol_txn"
- ]
- _simple_update_one = SQLBaseStore.__dict__["_simple_update_one"]
- _simple_update_one_txn = SQLBaseStore.__dict__["_simple_update_one_txn"]
- _simple_update_txn = SQLBaseStore.__dict__["_simple_update_txn"]
- def runInteraction(self, desc, func, *args, **kwargs):
- def r(conn):
- try:
- i = 0
- N = 5
- while True:
- try:
- txn = conn.cursor()
- return func(
- LoggingTransaction(txn, desc, self.database_engine, [], []),
- *args, **kwargs
- )
- except self.database_engine.module.DatabaseError as e:
- if self.database_engine.is_deadlock(e):
- logger.warn("[TXN DEADLOCK] {%s} %d/%d", desc, i, N)
- if i < N:
- i += 1
- conn.rollback()
- continue
- raise
- except Exception as e:
- logger.debug("[TXN FAIL] {%s} %s", desc, e)
- raise
- return self.db_pool.runWithConnection(r)
- def execute(self, f, *args, **kwargs):
- return self.runInteraction(f.__name__, f, *args, **kwargs)
- def execute_sql(self, sql, *args):
- def r(txn):
- txn.execute(sql, args)
- return txn.fetchall()
- return self.runInteraction("execute_sql", r)
- def insert_many_txn(self, txn, table, headers, rows):
- sql = "INSERT INTO %s (%s) VALUES (%s)" % (
- table,
- ", ".join(k for k in headers),
- ", ".join("%s" for _ in headers)
- )
- try:
- txn.executemany(sql, rows)
- except:
- logger.exception(
- "Failed to insert: %s",
- table,
- )
- raise
- class Porter(object):
- def __init__(self, **kwargs):
- self.__dict__.update(kwargs)
- @defer.inlineCallbacks
- def setup_table(self, table):
- if table in APPEND_ONLY_TABLES:
- # It's safe to just carry on inserting.
- row = yield self.postgres_store._simple_select_one(
- table="port_from_sqlite3",
- keyvalues={"table_name": table},
- retcols=("forward_rowid", "backward_rowid"),
- allow_none=True,
- )
- total_to_port = None
- if row is None:
- if table == "sent_transactions":
- forward_chunk, already_ported, total_to_port = (
- yield self._setup_sent_transactions()
- )
- backward_chunk = 0
- else:
- yield self.postgres_store._simple_insert(
- table="port_from_sqlite3",
- values={
- "table_name": table,
- "forward_rowid": 1,
- "backward_rowid": 0,
- }
- )
- forward_chunk = 1
- backward_chunk = 0
- already_ported = 0
- else:
- forward_chunk = row["forward_rowid"]
- backward_chunk = row["backward_rowid"]
- if total_to_port is None:
- already_ported, total_to_port = yield self._get_total_count_to_port(
- table, forward_chunk, backward_chunk
- )
- else:
- def delete_all(txn):
- txn.execute(
- "DELETE FROM port_from_sqlite3 WHERE table_name = %s",
- (table,)
- )
- txn.execute("TRUNCATE %s CASCADE" % (table,))
- yield self.postgres_store.execute(delete_all)
- yield self.postgres_store._simple_insert(
- table="port_from_sqlite3",
- values={
- "table_name": table,
- "forward_rowid": 1,
- "backward_rowid": 0,
- }
- )
- forward_chunk = 1
- backward_chunk = 0
- already_ported, total_to_port = yield self._get_total_count_to_port(
- table, forward_chunk, backward_chunk
- )
- defer.returnValue(
- (table, already_ported, total_to_port, forward_chunk, backward_chunk)
- )
- @defer.inlineCallbacks
- def handle_table(self, table, postgres_size, table_size, forward_chunk,
- backward_chunk):
- logger.info(
- "Table %s: %i/%i (rows %i-%i) already ported",
- table, postgres_size, table_size,
- backward_chunk+1, forward_chunk-1,
- )
- if not table_size:
- return
- self.progress.add_table(table, postgres_size, table_size)
- if table == "event_search":
- yield self.handle_search_table(
- postgres_size, table_size, forward_chunk, backward_chunk
- )
- return
- if table in (
- "user_directory", "user_directory_search", "users_who_share_rooms",
- "users_in_pubic_room",
- ):
- # We don't port these tables, as they're a faff and we can regenreate
- # them anyway.
- self.progress.update(table, table_size) # Mark table as done
- return
- if table == "user_directory_stream_pos":
- # We need to make sure there is a single row, `(X, null), as that is
- # what synapse expects to be there.
- yield self.postgres_store._simple_insert(
- table=table,
- values={"stream_id": None},
- )
- self.progress.update(table, table_size) # Mark table as done
- return
- forward_select = (
- "SELECT rowid, * FROM %s WHERE rowid >= ? ORDER BY rowid LIMIT ?"
- % (table,)
- )
- backward_select = (
- "SELECT rowid, * FROM %s WHERE rowid <= ? ORDER BY rowid LIMIT ?"
- % (table,)
- )
- do_forward = [True]
- do_backward = [True]
- while True:
- def r(txn):
- forward_rows = []
- backward_rows = []
- if do_forward[0]:
- txn.execute(forward_select, (forward_chunk, self.batch_size,))
- forward_rows = txn.fetchall()
- if not forward_rows:
- do_forward[0] = False
- if do_backward[0]:
- txn.execute(backward_select, (backward_chunk, self.batch_size,))
- backward_rows = txn.fetchall()
- if not backward_rows:
- do_backward[0] = False
- if forward_rows or backward_rows:
- headers = [column[0] for column in txn.description]
- else:
- headers = None
- return headers, forward_rows, backward_rows
- headers, frows, brows = yield self.sqlite_store.runInteraction(
- "select", r
- )
- if frows or brows:
- if frows:
- forward_chunk = max(row[0] for row in frows) + 1
- if brows:
- backward_chunk = min(row[0] for row in brows) - 1
- rows = frows + brows
- rows = self._convert_rows(table, headers, rows)
- def insert(txn):
- self.postgres_store.insert_many_txn(
- txn, table, headers[1:], rows
- )
- self.postgres_store._simple_update_one_txn(
- txn,
- table="port_from_sqlite3",
- keyvalues={"table_name": table},
- updatevalues={
- "forward_rowid": forward_chunk,
- "backward_rowid": backward_chunk,
- },
- )
- yield self.postgres_store.execute(insert)
- postgres_size += len(rows)
- self.progress.update(table, postgres_size)
- else:
- return
- @defer.inlineCallbacks
- def handle_search_table(self, postgres_size, table_size, forward_chunk,
- backward_chunk):
- select = (
- "SELECT es.rowid, es.*, e.origin_server_ts, e.stream_ordering"
- " FROM event_search as es"
- " INNER JOIN events AS e USING (event_id, room_id)"
- " WHERE es.rowid >= ?"
- " ORDER BY es.rowid LIMIT ?"
- )
- while True:
- def r(txn):
- txn.execute(select, (forward_chunk, self.batch_size,))
- rows = txn.fetchall()
- headers = [column[0] for column in txn.description]
- return headers, rows
- headers, rows = yield self.sqlite_store.runInteraction("select", r)
- if rows:
- forward_chunk = rows[-1][0] + 1
- # We have to treat event_search differently since it has a
- # different structure in the two different databases.
- def insert(txn):
- sql = (
- "INSERT INTO event_search (event_id, room_id, key,"
- " sender, vector, origin_server_ts, stream_ordering)"
- " VALUES (?,?,?,?,to_tsvector('english', ?),?,?)"
- )
- rows_dict = []
- for row in rows:
- d = dict(zip(headers, row))
- if "\0" in d['value']:
- logger.warn('dropping search row %s', d)
- else:
- rows_dict.append(d)
- txn.executemany(sql, [
- (
- row["event_id"],
- row["room_id"],
- row["key"],
- row["sender"],
- row["value"],
- row["origin_server_ts"],
- row["stream_ordering"],
- )
- for row in rows_dict
- ])
- self.postgres_store._simple_update_one_txn(
- txn,
- table="port_from_sqlite3",
- keyvalues={"table_name": "event_search"},
- updatevalues={
- "forward_rowid": forward_chunk,
- "backward_rowid": backward_chunk,
- },
- )
- yield self.postgres_store.execute(insert)
- postgres_size += len(rows)
- self.progress.update("event_search", 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_")
- }
- )
- prepare_database(db_conn, database_engine, config=None)
- db_conn.commit()
- @defer.inlineCallbacks
- def run(self):
- try:
- sqlite_db_pool = adbapi.ConnectionPool(
- self.sqlite_config["name"],
- **self.sqlite_config["args"]
- )
- postgres_db_pool = adbapi.ConnectionPool(
- self.postgres_config["name"],
- **self.postgres_config["args"]
- )
- sqlite_engine = create_engine(sqlite_config)
- postgres_engine = create_engine(postgres_config)
- self.sqlite_store = Store(sqlite_db_pool, sqlite_engine)
- self.postgres_store = Store(postgres_db_pool, postgres_engine)
- yield self.postgres_store.execute(
- postgres_engine.check_database
- )
- # Step 1. Set up databases.
- self.progress.set_state("Preparing SQLite3")
- self.setup_db(sqlite_config, sqlite_engine)
- self.progress.set_state("Preparing PostgreSQL")
- self.setup_db(postgres_config, postgres_engine)
- self.progress.set_state("Creating port tables")
- def create_port_table(txn):
- txn.execute(
- "CREATE TABLE IF NOT EXISTS port_from_sqlite3 ("
- " table_name varchar(100) NOT NULL UNIQUE,"
- " forward_rowid bigint NOT NULL,"
- " backward_rowid bigint NOT NULL"
- ")"
- )
- # The old port script created a table with just a "rowid" column.
- # We want people to be able to rerun this script from an old port
- # so that they can pick up any missing events that were not
- # ported across.
- def alter_table(txn):
- txn.execute(
- "ALTER TABLE IF EXISTS port_from_sqlite3"
- " RENAME rowid TO forward_rowid"
- )
- txn.execute(
- "ALTER TABLE IF EXISTS port_from_sqlite3"
- " ADD backward_rowid bigint NOT NULL DEFAULT 0"
- )
- try:
- yield self.postgres_store.runInteraction(
- "alter_table", alter_table
- )
- except Exception as e:
- pass
- yield self.postgres_store.runInteraction(
- "create_port_table", create_port_table
- )
- # Step 2. Get tables.
- self.progress.set_state("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={},
- retcol="distinct table_name",
- )
- tables = set(sqlite_tables) & set(postgres_tables)
- logger.info("Found %d tables", len(tables))
- # Step 3. Figure out what still needs copying
- self.progress.set_state("Checking on port progress")
- setup_res = yield defer.gatherResults(
- [
- self.setup_table(table)
- for table in tables
- if table not in ["schema_version", "applied_schema_deltas"]
- and not table.startswith("sqlite_")
- ],
- consumeErrors=True,
- )
- # Step 4. Do the copying.
- self.progress.set_state("Copying to postgres")
- yield defer.gatherResults(
- [
- self.handle_table(*res)
- for res in setup_res
- ],
- consumeErrors=True,
- )
- # Step 5. Do final post-processing
- yield self._setup_state_group_id_seq()
- self.progress.done()
- except:
- global end_error_exec_info
- end_error_exec_info = sys.exc_info()
- logger.exception("")
- finally:
- reactor.stop()
- def _convert_rows(self, table, headers, rows):
- bool_col_names = BOOLEAN_COLUMNS.get(table, [])
- bool_cols = [
- i for i, h in enumerate(headers) if h in bool_col_names
- ]
- class BadValueException(Exception):
- pass
- def conv(j, col):
- if j in bool_cols:
- return bool(col)
- elif isinstance(col, basestring) and "\0" in col:
- logger.warn("DROPPING ROW: NUL value in table %s col %s: %r", table, headers[j], col)
- raise BadValueException();
- return col
- outrows = []
- for i, row in enumerate(rows):
- try:
- outrows.append(tuple(
- conv(j, col)
- for j, col in enumerate(row)
- if j > 0
- ))
- except BadValueException:
- pass
- return outrows
- @defer.inlineCallbacks
- def _setup_sent_transactions(self):
- # Only save things from the last day
- yesterday = int(time.time() * 1000) - 86400000
- # And save the max transaction id from each destination
- select = (
- "SELECT rowid, * FROM sent_transactions WHERE rowid IN ("
- "SELECT max(rowid) FROM sent_transactions"
- " GROUP BY destination"
- ")"
- )
- def r(txn):
- txn.execute(select)
- rows = txn.fetchall()
- headers = [column[0] for column in txn.description]
- ts_ind = headers.index('ts')
- return headers, [r for r in rows if r[ts_ind] < yesterday]
- headers, rows = yield self.sqlite_store.runInteraction(
- "select", r,
- )
- rows = self._convert_rows("sent_transactions", headers, rows)
- inserted_rows = len(rows)
- if inserted_rows:
- max_inserted_rowid = max(r[0] for r in rows)
- def insert(txn):
- self.postgres_store.insert_many_txn(
- txn, "sent_transactions", headers[1:], rows
- )
- yield self.postgres_store.execute(insert)
- else:
- max_inserted_rowid = 0
- def get_start_id(txn):
- txn.execute(
- "SELECT rowid FROM sent_transactions WHERE ts >= ?"
- " ORDER BY rowid ASC LIMIT 1",
- (yesterday,)
- )
- rows = txn.fetchall()
- if rows:
- return rows[0][0]
- else:
- return 1
- next_chunk = yield self.sqlite_store.execute(get_start_id)
- next_chunk = max(max_inserted_rowid + 1, next_chunk)
- yield self.postgres_store._simple_insert(
- table="port_from_sqlite3",
- values={
- "table_name": "sent_transactions",
- "forward_rowid": next_chunk,
- "backward_rowid": 0,
- }
- )
- def get_sent_table_size(txn):
- txn.execute(
- "SELECT count(*) FROM sent_transactions"
- " WHERE ts >= ?",
- (yesterday,)
- )
- size, = txn.fetchone()
- return int(size)
- remaining_count = yield self.sqlite_store.execute(
- get_sent_table_size
- )
- total_count = remaining_count + inserted_rows
- defer.returnValue((next_chunk, inserted_rows, total_count))
- @defer.inlineCallbacks
- def _get_remaining_count_to_port(self, table, forward_chunk, backward_chunk):
- frows = yield self.sqlite_store.execute_sql(
- "SELECT count(*) FROM %s WHERE rowid >= ?" % (table,),
- forward_chunk,
- )
- brows = yield self.sqlite_store.execute_sql(
- "SELECT count(*) FROM %s WHERE rowid <= ?" % (table,),
- backward_chunk,
- )
- defer.returnValue(frows[0][0] + brows[0][0])
- @defer.inlineCallbacks
- def _get_already_ported_count(self, table):
- rows = yield self.postgres_store.execute_sql(
- "SELECT count(*) FROM %s" % (table,),
- )
- defer.returnValue(rows[0][0])
- @defer.inlineCallbacks
- def _get_total_count_to_port(self, table, forward_chunk, backward_chunk):
- remaining, done = yield defer.gatherResults(
- [
- self._get_remaining_count_to_port(table, forward_chunk, backward_chunk),
- self._get_already_ported_count(table),
- ],
- consumeErrors=True,
- )
- remaining = int(remaining) if remaining else 0
- done = int(done) if done else 0
- defer.returnValue((done, remaining + done))
- def _setup_state_group_id_seq(self):
- def r(txn):
- txn.execute("SELECT MAX(id) FROM state_groups")
- next_id = txn.fetchone()[0]+1
- txn.execute(
- "ALTER SEQUENCE state_group_id_seq RESTART WITH %s",
- (next_id,),
- )
- return self.postgres_store.runInteraction("setup_state_group_id_seq", r)
- ##############################################
- ###### The following is simply UI stuff ######
- ##############################################
- 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
- class CursesProgress(Progress):
- """Reports progress to a curses window
- """
- def __init__(self, stdscr):
- self.stdscr = stdscr
- curses.use_default_colors()
- curses.curs_set(0)
- curses.init_pair(1, curses.COLOR_RED, -1)
- curses.init_pair(2, curses.COLOR_GREEN, -1)
- self.last_update = 0
- self.finished = False
- self.total_processed = 0
- self.total_remaining = 0
- super(CursesProgress, self).__init__()
- def update(self, table, num_done):
- super(CursesProgress, self).update(table, num_done)
- self.total_processed = 0
- self.total_remaining = 0
- for table, data in self.tables.items():
- self.total_processed += data["num_done"] - data["start"]
- self.total_remaining += data["total"] - data["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:
- if self.total_processed > 0:
- left = float(self.total_remaining) / self.total_processed
- est_remaining = (int(now) - self.start_time) * left
- 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,
- )
- 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,
- )
- size = 20
- progress = "[%s%s]" % (
- "#" * int(perc * size / 100),
- " " * (size - int(perc * size / 100)),
- )
- self.stdscr.addstr(
- i + 2, left_margin + max_len + middle_space,
- "%s %3d%% (%d/%d)" % (progress, perc, data["num_done"], data["total"]),
- )
- if self.finished:
- self.stdscr.addstr(
- rows - 1, 0,
- "Press any key to exit...",
- )
- self.stdscr.refresh()
- self.last_update = time.time()
- def done(self):
- self.finished = True
- self.render(True)
- self.stdscr.getch()
- def set_state(self, state):
- self.stdscr.clear()
- self.stdscr.addstr(
- 0, 0,
- state + "...",
- curses.A_BOLD,
- )
- self.stdscr.refresh()
- class TerminalProgress(Progress):
- """Just prints progress to the terminal
- """
- def update(self, table, num_done):
- super(TerminalProgress, self).update(table, num_done)
- data = self.tables[table]
- print "%s: %d%% (%d/%d)" % (
- table, data["perc"],
- data["num_done"], data["total"],
- )
- def set_state(self, state):
- print state + "..."
- ##############################################
- ##############################################
- if __name__ == "__main__":
- parser = argparse.ArgumentParser(
- description="A script to port an existing synapse SQLite database to"
- " a new PostgreSQL database."
- )
- parser.add_argument("-v", action='store_true')
- parser.add_argument(
- "--sqlite-database", required=True,
- help="The snapshot of the SQLite database file. This must not be"
- " currently used by a running synapse server"
- )
- parser.add_argument(
- "--postgres-config", type=argparse.FileType('r'), required=True,
- help="The database config file for the PostgreSQL database"
- )
- parser.add_argument(
- "--curses", action='store_true',
- help="display a curses based progress UI"
- )
- parser.add_argument(
- "--batch-size", type=int, default=1000,
- help="The number of rows to select from the SQLite table each"
- " iteration [default=1000]",
- )
- args = parser.parse_args()
- 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",
- "args": {
- "database": args.sqlite_database,
- "cp_min": 1,
- "cp_max": 1,
- "check_same_thread": False,
- },
- }
- postgres_config = yaml.safe_load(args.postgres_config)
- if "database" in postgres_config:
- postgres_config = postgres_config["database"]
- if "name" not in postgres_config:
- sys.stderr.write("Malformed database config: no 'name'")
- sys.exit(2)
- if postgres_config["name"] != "psycopg2":
- sys.stderr.write("Database must use 'psycopg2' connector.")
- sys.exit(3)
- 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()
- if end_error_exec_info:
- exc_type, exc_value, exc_traceback = end_error_exec_info
- traceback.print_exception(exc_type, exc_value, exc_traceback)
|