diff options
Diffstat (limited to 'synapse/_scripts/synapse_port_db.py')
| -rwxr-xr-x | synapse/_scripts/synapse_port_db.py | 1257 |
1 files changed, 1257 insertions, 0 deletions
diff --git a/synapse/_scripts/synapse_port_db.py b/synapse/_scripts/synapse_port_db.py new file mode 100755 index 0000000000..c38666da18 --- /dev/null +++ b/synapse/_scripts/synapse_port_db.py @@ -0,0 +1,1257 @@ +#!/usr/bin/env python +# Copyright 2015, 2016 OpenMarket Ltd +# Copyright 2018 New Vector Ltd +# Copyright 2019 The Matrix.org Foundation C.I.C. +# +# 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. + +import argparse +import curses +import logging +import sys +import time +import traceback +from typing import Dict, Iterable, Optional, Set + +import yaml +from matrix_common.versionstring import get_distribution_version_string + +from twisted.internet import defer, reactor + +from synapse.config.database import DatabaseConnectionConfig +from synapse.config.homeserver import HomeServerConfig +from synapse.logging.context import ( + LoggingContext, + make_deferred_yieldable, + run_in_background, +) +from synapse.storage.database import DatabasePool, make_conn +from synapse.storage.databases.main import PushRuleStore +from synapse.storage.databases.main.account_data import AccountDataWorkerStore +from synapse.storage.databases.main.client_ips import ClientIpBackgroundUpdateStore +from synapse.storage.databases.main.deviceinbox import DeviceInboxBackgroundUpdateStore +from synapse.storage.databases.main.devices import DeviceBackgroundUpdateStore +from synapse.storage.databases.main.end_to_end_keys import EndToEndKeyBackgroundStore +from synapse.storage.databases.main.events_bg_updates import ( + EventsBackgroundUpdatesStore, +) +from synapse.storage.databases.main.group_server import GroupServerWorkerStore +from synapse.storage.databases.main.media_repository import ( + MediaRepositoryBackgroundUpdateStore, +) +from synapse.storage.databases.main.presence import PresenceBackgroundUpdateStore +from synapse.storage.databases.main.pusher import PusherWorkerStore +from synapse.storage.databases.main.registration import ( + RegistrationBackgroundUpdateStore, + find_max_generated_user_id_localpart, +) +from synapse.storage.databases.main.room import RoomBackgroundUpdateStore +from synapse.storage.databases.main.roommember import RoomMemberBackgroundUpdateStore +from synapse.storage.databases.main.search import SearchBackgroundUpdateStore +from synapse.storage.databases.main.state import MainStateBackgroundUpdateStore +from synapse.storage.databases.main.stats import StatsStore +from synapse.storage.databases.main.user_directory import ( + UserDirectoryBackgroundUpdateStore, +) +from synapse.storage.databases.state.bg_updates import StateBackgroundUpdateStore +from synapse.storage.engines import create_engine +from synapse.storage.prepare_database import prepare_database +from synapse.util import Clock + +logger = logging.getLogger("synapse_port_db") + + +BOOLEAN_COLUMNS = { + "events": ["processed", "outlier", "contains_url"], + "rooms": ["is_public", "has_auth_chain_index"], + "event_edges": ["is_state"], + "presence_list": ["accepted"], + "presence_stream": ["currently_active"], + "public_room_list_stream": ["visibility"], + "devices": ["hidden"], + "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"], + "e2e_room_keys": ["is_verified"], + "account_validity": ["email_sent"], + "redactions": ["have_censored"], + "room_stats_state": ["is_federatable"], + "local_media_repository": ["safe_from_quarantine"], + "users": ["shadow_banned"], + "e2e_fallback_keys_json": ["used"], + "access_tokens": ["used"], +} + + +APPEND_ONLY_TABLES = [ + "event_reference_hashes", + "events", + "event_json", + "state_events", + "room_memberships", + "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", + "ex_outlier_stream", + "cache_invalidation_stream_by_instance", + "public_room_list_stream", + "state_group_edges", + "stream_ordering_to_exterm", +] + + +IGNORED_TABLES = { + # We don't port these tables, as they're a faff and we can regenerate + # them anyway. + "user_directory", + "user_directory_search", + "user_directory_search_content", + "user_directory_search_docsize", + "user_directory_search_segdir", + "user_directory_search_segments", + "user_directory_search_stat", + "user_directory_search_pos", + "users_who_share_private_rooms", + "users_in_public_room", + # UI auth sessions have foreign keys so additional care needs to be taken, + # the sessions are transient anyway, so ignore them. + "ui_auth_sessions", + "ui_auth_sessions_credentials", + "ui_auth_sessions_ips", +} + + +# Error returned by the run function. Used at the top-level part of the script to +# handle errors and return codes. +end_error = None # type: Optional[str] +# The exec_info for the error, if any. If error is defined but not exec_info the script +# will show only the error message without the stacktrace, if exec_info is defined but +# not the error then the script will show nothing outside of what's printed in the run +# function. If both are defined, the script will print both the error and the stacktrace. +end_error_exec_info = None + + +class Store( + ClientIpBackgroundUpdateStore, + DeviceInboxBackgroundUpdateStore, + DeviceBackgroundUpdateStore, + EventsBackgroundUpdatesStore, + MediaRepositoryBackgroundUpdateStore, + RegistrationBackgroundUpdateStore, + RoomBackgroundUpdateStore, + RoomMemberBackgroundUpdateStore, + SearchBackgroundUpdateStore, + StateBackgroundUpdateStore, + MainStateBackgroundUpdateStore, + UserDirectoryBackgroundUpdateStore, + EndToEndKeyBackgroundStore, + StatsStore, + AccountDataWorkerStore, + PushRuleStore, + PusherWorkerStore, + PresenceBackgroundUpdateStore, + GroupServerWorkerStore, +): + def execute(self, f, *args, **kwargs): + return self.db_pool.runInteraction(f.__name__, f, *args, **kwargs) + + def execute_sql(self, sql, *args): + def r(txn): + txn.execute(sql, args) + return txn.fetchall() + + return self.db_pool.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 Exception: + logger.exception("Failed to insert: %s", table) + raise + + def set_room_is_public(self, room_id, is_public): + raise Exception( + "Attempt to set room_is_public during port_db: database not empty?" + ) + + +class MockHomeserver: + def __init__(self, config): + self.clock = Clock(reactor) + self.config = config + self.hostname = config.server.server_name + self.version_string = "Synapse/" + get_distribution_version_string( + "matrix-synapse" + ) + + def get_clock(self): + return self.clock + + def get_reactor(self): + return reactor + + def get_instance_name(self): + return "master" + + +class Porter(object): + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + async def setup_table(self, table): + if table in APPEND_ONLY_TABLES: + # It's safe to just carry on inserting. + row = await self.postgres_store.db_pool.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, + ) = await self._setup_sent_transactions() + backward_chunk = 0 + else: + await self.postgres_store.db_pool.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 = await 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,)) + + await self.postgres_store.execute(delete_all) + + await self.postgres_store.db_pool.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 = await self._get_total_count_to_port( + table, forward_chunk, backward_chunk + ) + + return table, already_ported, total_to_port, forward_chunk, backward_chunk + + async def get_table_constraints(self) -> Dict[str, Set[str]]: + """Returns a map of tables that have foreign key constraints to tables they depend on.""" + + def _get_constraints(txn): + # We can pull the information about foreign key constraints out from + # the postgres schema tables. + sql = """ + SELECT DISTINCT + tc.table_name, + ccu.table_name AS foreign_table_name + FROM + information_schema.table_constraints AS tc + INNER JOIN information_schema.constraint_column_usage AS ccu + USING (table_schema, constraint_name) + WHERE tc.constraint_type = 'FOREIGN KEY' + AND tc.table_name != ccu.table_name; + """ + txn.execute(sql) + + results = {} + for table, foreign_table in txn: + results.setdefault(table, set()).add(foreign_table) + return results + + return await self.postgres_store.db_pool.runInteraction( + "get_table_constraints", _get_constraints + ) + + async 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": + await self.handle_search_table( + postgres_size, table_size, forward_chunk, backward_chunk + ) + return + + if table in IGNORED_TABLES: + 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. + await self.postgres_store.db_pool.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 = await self.sqlite_store.db_pool.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.db_pool.simple_update_one_txn( + txn, + table="port_from_sqlite3", + keyvalues={"table_name": table}, + updatevalues={ + "forward_rowid": forward_chunk, + "backward_rowid": backward_chunk, + }, + ) + + await self.postgres_store.execute(insert) + + postgres_size += len(rows) + + self.progress.update(table, postgres_size) + else: + return + + async 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 = await self.sqlite_store.db_pool.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.warning("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.db_pool.simple_update_one_txn( + txn, + table="port_from_sqlite3", + keyvalues={"table_name": "event_search"}, + updatevalues={ + "forward_rowid": forward_chunk, + "backward_rowid": backward_chunk, + }, + ) + + await self.postgres_store.execute(insert) + + postgres_size += len(rows) + + self.progress.update("event_search", postgres_size) + + else: + return + + def build_db_store( + self, + db_config: DatabaseConnectionConfig, + allow_outdated_version: bool = False, + ): + """Builds and returns a database store using the provided configuration. + + Args: + db_config: The database configuration + allow_outdated_version: True to suppress errors about the database server + version being too old to run a complete synapse + + Returns: + The built Store object. + """ + self.progress.set_state("Preparing %s" % db_config.config["name"]) + + engine = create_engine(db_config.config) + + hs = MockHomeserver(self.hs_config) + + with make_conn(db_config, engine, "portdb") as db_conn: + engine.check_database( + db_conn, allow_outdated_version=allow_outdated_version + ) + prepare_database(db_conn, engine, config=self.hs_config) + store = Store(DatabasePool(hs, db_config, engine), db_conn, hs) + db_conn.commit() + + return store + + async def run_background_updates_on_postgres(self): + # Manually apply all background updates on the PostgreSQL database. + postgres_ready = ( + await self.postgres_store.db_pool.updates.has_completed_background_updates() + ) + + if not postgres_ready: + # Only say that we're running background updates when there are background + # updates to run. + self.progress.set_state("Running background updates on PostgreSQL") + + while not postgres_ready: + await self.postgres_store.db_pool.updates.do_next_background_update(100) + postgres_ready = await ( + self.postgres_store.db_pool.updates.has_completed_background_updates() + ) + + async def run(self): + """Ports the SQLite database to a PostgreSQL database. + + When a fatal error is met, its message is assigned to the global "end_error" + variable. When this error comes with a stacktrace, its exec_info is assigned to + the global "end_error_exec_info" variable. + """ + global end_error + + try: + # we allow people to port away from outdated versions of sqlite. + self.sqlite_store = self.build_db_store( + DatabaseConnectionConfig("master-sqlite", self.sqlite_config), + allow_outdated_version=True, + ) + + # Check if all background updates are done, abort if not. + updates_complete = ( + await self.sqlite_store.db_pool.updates.has_completed_background_updates() + ) + if not updates_complete: + end_error = ( + "Pending background updates exist in the SQLite3 database." + " Please start Synapse again and wait until every update has finished" + " before running this script.\n" + ) + return + + self.postgres_store = self.build_db_store( + self.hs_config.database.get_single_database() + ) + + await self.run_background_updates_on_postgres() + + 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: + await self.postgres_store.db_pool.runInteraction( + "alter_table", alter_table + ) + except Exception: + # On Error Resume Next + pass + + await self.postgres_store.db_pool.runInteraction( + "create_port_table", create_port_table + ) + + # Step 2. Set up sequences + # + # We do this before porting the tables so that event if we fail half + # way through the postgres DB always have sequences that are greater + # than their respective tables. If we don't then creating the + # `DataStore` object will fail due to the inconsistency. + self.progress.set_state("Setting up sequence generators") + await self._setup_state_group_id_seq() + await self._setup_user_id_seq() + await self._setup_events_stream_seqs() + await self._setup_sequence( + "device_inbox_sequence", ("device_inbox", "device_federation_outbox") + ) + await self._setup_sequence( + "account_data_sequence", + ("room_account_data", "room_tags_revisions", "account_data"), + ) + await self._setup_sequence("receipts_sequence", ("receipts_linearized",)) + await self._setup_sequence("presence_stream_sequence", ("presence_stream",)) + await self._setup_auth_chain_sequence() + + # Step 3. Get tables. + self.progress.set_state("Fetching tables") + sqlite_tables = await self.sqlite_store.db_pool.simple_select_onecol( + table="sqlite_master", keyvalues={"type": "table"}, retcol="name" + ) + + postgres_tables = await self.postgres_store.db_pool.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 4. Figure out what still needs copying + self.progress.set_state("Checking on port progress") + setup_res = await make_deferred_yieldable( + defer.gatherResults( + [ + run_in_background(self.setup_table, table) + for table in tables + if table not in ["schema_version", "applied_schema_deltas"] + and not table.startswith("sqlite_") + ], + consumeErrors=True, + ) + ) + # Map from table name to args passed to `handle_table`, i.e. a tuple + # of: `postgres_size`, `table_size`, `forward_chunk`, `backward_chunk`. + tables_to_port_info_map = {r[0]: r[1:] for r in setup_res} + + # Step 5. Do the copying. + # + # This is slightly convoluted as we need to ensure tables are ported + # in the correct order due to foreign key constraints. + self.progress.set_state("Copying to postgres") + + constraints = await self.get_table_constraints() + tables_ported = set() # type: Set[str] + + while tables_to_port_info_map: + # Pulls out all tables that are still to be ported and which + # only depend on tables that are already ported (if any). + tables_to_port = [ + table + for table in tables_to_port_info_map + if not constraints.get(table, set()) - tables_ported + ] + + await make_deferred_yieldable( + defer.gatherResults( + [ + run_in_background( + self.handle_table, + table, + *tables_to_port_info_map.pop(table), + ) + for table in tables_to_port + ], + consumeErrors=True, + ) + ) + + tables_ported.update(tables_to_port) + + self.progress.done() + except Exception as e: + global end_error_exec_info + end_error = str(e) + 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) + if isinstance(col, bytes): + return bytearray(col) + elif isinstance(col, str) and "\0" in col: + logger.warning( + "DROPPING ROW: NUL value in table %s col %s: %r", + table, + headers[j], + col, + ) + raise BadValueException() + return col + + outrows = [] + for row in rows: + try: + outrows.append( + tuple(conv(j, col) for j, col in enumerate(row) if j > 0) + ) + except BadValueException: + pass + + return outrows + + async 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 = await self.sqlite_store.db_pool.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 + ) + + await 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 = await self.sqlite_store.execute(get_start_id) + next_chunk = max(max_inserted_rowid + 1, next_chunk) + + await self.postgres_store.db_pool.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 = await self.sqlite_store.execute(get_sent_table_size) + + total_count = remaining_count + inserted_rows + + return next_chunk, inserted_rows, total_count + + async def _get_remaining_count_to_port(self, table, forward_chunk, backward_chunk): + frows = await self.sqlite_store.execute_sql( + "SELECT count(*) FROM %s WHERE rowid >= ?" % (table,), forward_chunk + ) + + brows = await self.sqlite_store.execute_sql( + "SELECT count(*) FROM %s WHERE rowid <= ?" % (table,), backward_chunk + ) + + return frows[0][0] + brows[0][0] + + async def _get_already_ported_count(self, table): + rows = await self.postgres_store.execute_sql( + "SELECT count(*) FROM %s" % (table,) + ) + + return rows[0][0] + + async def _get_total_count_to_port(self, table, forward_chunk, backward_chunk): + remaining, done = await make_deferred_yieldable( + defer.gatherResults( + [ + run_in_background( + self._get_remaining_count_to_port, + table, + forward_chunk, + backward_chunk, + ), + run_in_background(self._get_already_ported_count, table), + ], + ) + ) + + remaining = int(remaining) if remaining else 0 + done = int(done) if done else 0 + + return done, remaining + done + + async def _setup_state_group_id_seq(self) -> None: + curr_id = await self.sqlite_store.db_pool.simple_select_one_onecol( + table="state_groups", keyvalues={}, retcol="MAX(id)", allow_none=True + ) + + if not curr_id: + return + + def r(txn): + next_id = curr_id + 1 + txn.execute("ALTER SEQUENCE state_group_id_seq RESTART WITH %s", (next_id,)) + + await self.postgres_store.db_pool.runInteraction("setup_state_group_id_seq", r) + + async def _setup_user_id_seq(self) -> None: + curr_id = await self.sqlite_store.db_pool.runInteraction( + "setup_user_id_seq", find_max_generated_user_id_localpart + ) + + def r(txn): + next_id = curr_id + 1 + txn.execute("ALTER SEQUENCE user_id_seq RESTART WITH %s", (next_id,)) + + await self.postgres_store.db_pool.runInteraction("setup_user_id_seq", r) + + async def _setup_events_stream_seqs(self) -> None: + """Set the event stream sequences to the correct values.""" + + # We get called before we've ported the events table, so we need to + # fetch the current positions from the SQLite store. + curr_forward_id = await self.sqlite_store.db_pool.simple_select_one_onecol( + table="events", keyvalues={}, retcol="MAX(stream_ordering)", allow_none=True + ) + + curr_backward_id = await self.sqlite_store.db_pool.simple_select_one_onecol( + table="events", + keyvalues={}, + retcol="MAX(-MIN(stream_ordering), 1)", + allow_none=True, + ) + + def _setup_events_stream_seqs_set_pos(txn): + if curr_forward_id: + txn.execute( + "ALTER SEQUENCE events_stream_seq RESTART WITH %s", + (curr_forward_id + 1,), + ) + + if curr_backward_id: + txn.execute( + "ALTER SEQUENCE events_backfill_stream_seq RESTART WITH %s", + (curr_backward_id + 1,), + ) + + await self.postgres_store.db_pool.runInteraction( + "_setup_events_stream_seqs", + _setup_events_stream_seqs_set_pos, + ) + + async def _setup_sequence( + self, sequence_name: str, stream_id_tables: Iterable[str] + ) -> None: + """Set a sequence to the correct value.""" + current_stream_ids = [] + for stream_id_table in stream_id_tables: + max_stream_id = await self.sqlite_store.db_pool.simple_select_one_onecol( + table=stream_id_table, + keyvalues={}, + retcol="COALESCE(MAX(stream_id), 1)", + allow_none=True, + ) + current_stream_ids.append(max_stream_id) + + next_id = max(current_stream_ids) + 1 + + def r(txn): + sql = "ALTER SEQUENCE %s RESTART WITH" % (sequence_name,) + txn.execute(sql + " %s", (next_id,)) + + await self.postgres_store.db_pool.runInteraction( + "_setup_%s" % (sequence_name,), r + ) + + async def _setup_auth_chain_sequence(self) -> None: + curr_chain_id = await self.sqlite_store.db_pool.simple_select_one_onecol( + table="event_auth_chains", + keyvalues={}, + retcol="MAX(chain_id)", + allow_none=True, + ) + + def r(txn): + txn.execute( + "ALTER SEQUENCE event_auth_chain_id RESTART WITH %s", + (curr_chain_id + 1,), + ) + + if curr_chain_id is not None: + await self.postgres_store.db_pool.runInteraction( + "_setup_event_auth_chain_id", + 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 data in self.tables.values(): + 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 = sorted(items, 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 + "...") + + +############################################## +############################################## + + +def 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, + }, + } + + hs_config = yaml.safe_load(args.postgres_config) + + if "database" not in hs_config: + sys.stderr.write("The configuration file must have a 'database' section.\n") + sys.exit(4) + + postgres_config = hs_config["database"] + + if "name" not in postgres_config: + sys.stderr.write("Malformed database config: no 'name'\n") + sys.exit(2) + if postgres_config["name"] != "psycopg2": + sys.stderr.write("Database must use the 'psycopg2' connector.\n") + sys.exit(3) + + config = HomeServerConfig() + config.parse_config_dict(hs_config, "", "") + + def start(stdscr=None): + if stdscr: + progress = CursesProgress(stdscr) + else: + progress = TerminalProgress() + + porter = Porter( + sqlite_config=sqlite_config, + progress=progress, + batch_size=args.batch_size, + hs_config=config, + ) + + @defer.inlineCallbacks + def run(): + with LoggingContext("synapse_port_db_run"): + yield defer.ensureDeferred(porter.run()) + + reactor.callWhenRunning(run) + + reactor.run() + + if args.curses: + curses.wrapper(start) + else: + start() + + if end_error: + if end_error_exec_info: + exc_type, exc_value, exc_traceback = end_error_exec_info + traceback.print_exception(exc_type, exc_value, exc_traceback) + + sys.stderr.write(end_error) + + sys.exit(5) + + +if __name__ == "__main__": + main() |