synapse_port_db 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992
  1. #!/usr/bin/env python
  2. # -*- coding: utf-8 -*-
  3. # Copyright 2015, 2016 OpenMarket Ltd
  4. #
  5. # Licensed under the Apache License, Version 2.0 (the "License");
  6. # you may not use this file except in compliance with the License.
  7. # You may obtain a copy of the License at
  8. #
  9. # http://www.apache.org/licenses/LICENSE-2.0
  10. #
  11. # Unless required by applicable law or agreed to in writing, software
  12. # distributed under the License is distributed on an "AS IS" BASIS,
  13. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14. # See the License for the specific language governing permissions and
  15. # limitations under the License.
  16. from twisted.internet import defer, reactor
  17. from twisted.enterprise import adbapi
  18. from synapse.storage._base import LoggingTransaction, SQLBaseStore
  19. from synapse.storage.engines import create_engine
  20. from synapse.storage.prepare_database import prepare_database
  21. import argparse
  22. import curses
  23. import logging
  24. import sys
  25. import time
  26. import traceback
  27. import yaml
  28. logger = logging.getLogger("synapse_port_db")
  29. BOOLEAN_COLUMNS = {
  30. "events": ["processed", "outlier", "contains_url"],
  31. "rooms": ["is_public"],
  32. "event_edges": ["is_state"],
  33. "presence_list": ["accepted"],
  34. "presence_stream": ["currently_active"],
  35. "public_room_list_stream": ["visibility"],
  36. "device_lists_outbound_pokes": ["sent"],
  37. "users_who_share_rooms": ["share_private"],
  38. "groups": ["is_public"],
  39. "group_rooms": ["is_public"],
  40. "group_users": ["is_public", "is_admin"],
  41. "group_summary_rooms": ["is_public"],
  42. "group_room_categories": ["is_public"],
  43. "group_summary_users": ["is_public"],
  44. "group_roles": ["is_public"],
  45. "local_group_membership": ["is_publicised", "is_admin"],
  46. }
  47. APPEND_ONLY_TABLES = [
  48. "event_content_hashes",
  49. "event_reference_hashes",
  50. "event_signatures",
  51. "event_edge_hashes",
  52. "events",
  53. "event_json",
  54. "state_events",
  55. "room_memberships",
  56. "feedback",
  57. "topics",
  58. "room_names",
  59. "rooms",
  60. "local_media_repository",
  61. "local_media_repository_thumbnails",
  62. "remote_media_cache",
  63. "remote_media_cache_thumbnails",
  64. "redactions",
  65. "event_edges",
  66. "event_auth",
  67. "received_transactions",
  68. "sent_transactions",
  69. "transaction_id_to_pdu",
  70. "users",
  71. "state_groups",
  72. "state_groups_state",
  73. "event_to_state_groups",
  74. "rejections",
  75. "event_search",
  76. "presence_stream",
  77. "push_rules_stream",
  78. "current_state_resets",
  79. "ex_outlier_stream",
  80. "cache_invalidation_stream",
  81. "public_room_list_stream",
  82. "state_group_edges",
  83. "stream_ordering_to_exterm",
  84. ]
  85. end_error_exec_info = None
  86. class Store(object):
  87. """This object is used to pull out some of the convenience API from the
  88. Storage layer.
  89. *All* database interactions should go through this object.
  90. """
  91. def __init__(self, db_pool, engine):
  92. self.db_pool = db_pool
  93. self.database_engine = engine
  94. _simple_insert_txn = SQLBaseStore.__dict__["_simple_insert_txn"]
  95. _simple_insert = SQLBaseStore.__dict__["_simple_insert"]
  96. _simple_select_onecol_txn = SQLBaseStore.__dict__["_simple_select_onecol_txn"]
  97. _simple_select_onecol = SQLBaseStore.__dict__["_simple_select_onecol"]
  98. _simple_select_one = SQLBaseStore.__dict__["_simple_select_one"]
  99. _simple_select_one_txn = SQLBaseStore.__dict__["_simple_select_one_txn"]
  100. _simple_select_one_onecol = SQLBaseStore.__dict__["_simple_select_one_onecol"]
  101. _simple_select_one_onecol_txn = SQLBaseStore.__dict__[
  102. "_simple_select_one_onecol_txn"
  103. ]
  104. _simple_update_one = SQLBaseStore.__dict__["_simple_update_one"]
  105. _simple_update_one_txn = SQLBaseStore.__dict__["_simple_update_one_txn"]
  106. _simple_update_txn = SQLBaseStore.__dict__["_simple_update_txn"]
  107. def runInteraction(self, desc, func, *args, **kwargs):
  108. def r(conn):
  109. try:
  110. i = 0
  111. N = 5
  112. while True:
  113. try:
  114. txn = conn.cursor()
  115. return func(
  116. LoggingTransaction(txn, desc, self.database_engine, [], []),
  117. *args, **kwargs
  118. )
  119. except self.database_engine.module.DatabaseError as e:
  120. if self.database_engine.is_deadlock(e):
  121. logger.warn("[TXN DEADLOCK] {%s} %d/%d", desc, i, N)
  122. if i < N:
  123. i += 1
  124. conn.rollback()
  125. continue
  126. raise
  127. except Exception as e:
  128. logger.debug("[TXN FAIL] {%s} %s", desc, e)
  129. raise
  130. return self.db_pool.runWithConnection(r)
  131. def execute(self, f, *args, **kwargs):
  132. return self.runInteraction(f.__name__, f, *args, **kwargs)
  133. def execute_sql(self, sql, *args):
  134. def r(txn):
  135. txn.execute(sql, args)
  136. return txn.fetchall()
  137. return self.runInteraction("execute_sql", r)
  138. def insert_many_txn(self, txn, table, headers, rows):
  139. sql = "INSERT INTO %s (%s) VALUES (%s)" % (
  140. table,
  141. ", ".join(k for k in headers),
  142. ", ".join("%s" for _ in headers)
  143. )
  144. try:
  145. txn.executemany(sql, rows)
  146. except:
  147. logger.exception(
  148. "Failed to insert: %s",
  149. table,
  150. )
  151. raise
  152. class Porter(object):
  153. def __init__(self, **kwargs):
  154. self.__dict__.update(kwargs)
  155. @defer.inlineCallbacks
  156. def setup_table(self, table):
  157. if table in APPEND_ONLY_TABLES:
  158. # It's safe to just carry on inserting.
  159. row = yield self.postgres_store._simple_select_one(
  160. table="port_from_sqlite3",
  161. keyvalues={"table_name": table},
  162. retcols=("forward_rowid", "backward_rowid"),
  163. allow_none=True,
  164. )
  165. total_to_port = None
  166. if row is None:
  167. if table == "sent_transactions":
  168. forward_chunk, already_ported, total_to_port = (
  169. yield self._setup_sent_transactions()
  170. )
  171. backward_chunk = 0
  172. else:
  173. yield self.postgres_store._simple_insert(
  174. table="port_from_sqlite3",
  175. values={
  176. "table_name": table,
  177. "forward_rowid": 1,
  178. "backward_rowid": 0,
  179. }
  180. )
  181. forward_chunk = 1
  182. backward_chunk = 0
  183. already_ported = 0
  184. else:
  185. forward_chunk = row["forward_rowid"]
  186. backward_chunk = row["backward_rowid"]
  187. if total_to_port is None:
  188. already_ported, total_to_port = yield self._get_total_count_to_port(
  189. table, forward_chunk, backward_chunk
  190. )
  191. else:
  192. def delete_all(txn):
  193. txn.execute(
  194. "DELETE FROM port_from_sqlite3 WHERE table_name = %s",
  195. (table,)
  196. )
  197. txn.execute("TRUNCATE %s CASCADE" % (table,))
  198. yield self.postgres_store.execute(delete_all)
  199. yield self.postgres_store._simple_insert(
  200. table="port_from_sqlite3",
  201. values={
  202. "table_name": table,
  203. "forward_rowid": 1,
  204. "backward_rowid": 0,
  205. }
  206. )
  207. forward_chunk = 1
  208. backward_chunk = 0
  209. already_ported, total_to_port = yield self._get_total_count_to_port(
  210. table, forward_chunk, backward_chunk
  211. )
  212. defer.returnValue(
  213. (table, already_ported, total_to_port, forward_chunk, backward_chunk)
  214. )
  215. @defer.inlineCallbacks
  216. def handle_table(self, table, postgres_size, table_size, forward_chunk,
  217. backward_chunk):
  218. if not table_size:
  219. return
  220. self.progress.add_table(table, postgres_size, table_size)
  221. if table == "event_search":
  222. yield self.handle_search_table(
  223. postgres_size, table_size, forward_chunk, backward_chunk
  224. )
  225. return
  226. if table in (
  227. "user_directory", "user_directory_search", "users_who_share_rooms",
  228. "users_in_pubic_room",
  229. ):
  230. # We don't port these tables, as they're a faff and we can regenreate
  231. # them anyway.
  232. self.progress.update(table, table_size) # Mark table as done
  233. return
  234. if table == "user_directory_stream_pos":
  235. # We need to make sure there is a single row, `(X, null), as that is
  236. # what synapse expects to be there.
  237. yield self.postgres_store._simple_insert(
  238. table=table,
  239. values={"stream_id": None},
  240. )
  241. self.progress.update(table, table_size) # Mark table as done
  242. return
  243. forward_select = (
  244. "SELECT rowid, * FROM %s WHERE rowid >= ? ORDER BY rowid LIMIT ?"
  245. % (table,)
  246. )
  247. backward_select = (
  248. "SELECT rowid, * FROM %s WHERE rowid <= ? ORDER BY rowid LIMIT ?"
  249. % (table,)
  250. )
  251. do_forward = [True]
  252. do_backward = [True]
  253. while True:
  254. def r(txn):
  255. forward_rows = []
  256. backward_rows = []
  257. if do_forward[0]:
  258. txn.execute(forward_select, (forward_chunk, self.batch_size,))
  259. forward_rows = txn.fetchall()
  260. if not forward_rows:
  261. do_forward[0] = False
  262. if do_backward[0]:
  263. txn.execute(backward_select, (backward_chunk, self.batch_size,))
  264. backward_rows = txn.fetchall()
  265. if not backward_rows:
  266. do_backward[0] = False
  267. if forward_rows or backward_rows:
  268. headers = [column[0] for column in txn.description]
  269. else:
  270. headers = None
  271. return headers, forward_rows, backward_rows
  272. headers, frows, brows = yield self.sqlite_store.runInteraction(
  273. "select", r
  274. )
  275. if frows or brows:
  276. if frows:
  277. forward_chunk = max(row[0] for row in frows) + 1
  278. if brows:
  279. backward_chunk = min(row[0] for row in brows) - 1
  280. rows = frows + brows
  281. rows = self._convert_rows(table, headers, rows)
  282. def insert(txn):
  283. self.postgres_store.insert_many_txn(
  284. txn, table, headers[1:], rows
  285. )
  286. self.postgres_store._simple_update_one_txn(
  287. txn,
  288. table="port_from_sqlite3",
  289. keyvalues={"table_name": table},
  290. updatevalues={
  291. "forward_rowid": forward_chunk,
  292. "backward_rowid": backward_chunk,
  293. },
  294. )
  295. yield self.postgres_store.execute(insert)
  296. postgres_size += len(rows)
  297. self.progress.update(table, postgres_size)
  298. else:
  299. return
  300. @defer.inlineCallbacks
  301. def handle_search_table(self, postgres_size, table_size, forward_chunk,
  302. backward_chunk):
  303. select = (
  304. "SELECT es.rowid, es.*, e.origin_server_ts, e.stream_ordering"
  305. " FROM event_search as es"
  306. " INNER JOIN events AS e USING (event_id, room_id)"
  307. " WHERE es.rowid >= ?"
  308. " ORDER BY es.rowid LIMIT ?"
  309. )
  310. while True:
  311. def r(txn):
  312. txn.execute(select, (forward_chunk, self.batch_size,))
  313. rows = txn.fetchall()
  314. headers = [column[0] for column in txn.description]
  315. return headers, rows
  316. headers, rows = yield self.sqlite_store.runInteraction("select", r)
  317. if rows:
  318. forward_chunk = rows[-1][0] + 1
  319. # We have to treat event_search differently since it has a
  320. # different structure in the two different databases.
  321. def insert(txn):
  322. sql = (
  323. "INSERT INTO event_search (event_id, room_id, key,"
  324. " sender, vector, origin_server_ts, stream_ordering)"
  325. " VALUES (?,?,?,?,to_tsvector('english', ?),?,?)"
  326. )
  327. rows_dict = []
  328. for row in rows:
  329. d = dict(zip(headers, row))
  330. if "\0" in d['value']:
  331. logger.warn('dropping search row %s', d)
  332. else:
  333. rows_dict.append(d)
  334. txn.executemany(sql, [
  335. (
  336. row["event_id"],
  337. row["room_id"],
  338. row["key"],
  339. row["sender"],
  340. row["value"],
  341. row["origin_server_ts"],
  342. row["stream_ordering"],
  343. )
  344. for row in rows_dict
  345. ])
  346. self.postgres_store._simple_update_one_txn(
  347. txn,
  348. table="port_from_sqlite3",
  349. keyvalues={"table_name": "event_search"},
  350. updatevalues={
  351. "forward_rowid": forward_chunk,
  352. "backward_rowid": backward_chunk,
  353. },
  354. )
  355. yield self.postgres_store.execute(insert)
  356. postgres_size += len(rows)
  357. self.progress.update("event_search", postgres_size)
  358. else:
  359. return
  360. def handle_state_group_seq(self):
  361. """Updates postgres handle_state_group_seq with last state_group ID
  362. """
  363. def handle_state_group_seq_txn(txn):
  364. txn.execute("""
  365. SELECT setval('state_group_id_seq',
  366. SELECT COALESCE(MAX(id), 0) FROM state_groups
  367. )
  368. """)
  369. return self.postgres_store.runInteraction(
  370. "handle_state_group_seq", handle_state_group_seq_txn,
  371. )
  372. def setup_db(self, db_config, database_engine):
  373. db_conn = database_engine.module.connect(
  374. **{
  375. k: v for k, v in db_config.get("args", {}).items()
  376. if not k.startswith("cp_")
  377. }
  378. )
  379. prepare_database(db_conn, database_engine, config=None)
  380. db_conn.commit()
  381. @defer.inlineCallbacks
  382. def run(self):
  383. try:
  384. sqlite_db_pool = adbapi.ConnectionPool(
  385. self.sqlite_config["name"],
  386. **self.sqlite_config["args"]
  387. )
  388. postgres_db_pool = adbapi.ConnectionPool(
  389. self.postgres_config["name"],
  390. **self.postgres_config["args"]
  391. )
  392. sqlite_engine = create_engine(sqlite_config)
  393. postgres_engine = create_engine(postgres_config)
  394. self.sqlite_store = Store(sqlite_db_pool, sqlite_engine)
  395. self.postgres_store = Store(postgres_db_pool, postgres_engine)
  396. yield self.postgres_store.execute(
  397. postgres_engine.check_database
  398. )
  399. # Step 1. Set up databases.
  400. self.progress.set_state("Preparing SQLite3")
  401. self.setup_db(sqlite_config, sqlite_engine)
  402. self.progress.set_state("Preparing PostgreSQL")
  403. self.setup_db(postgres_config, postgres_engine)
  404. # Step 2. Get tables.
  405. self.progress.set_state("Fetching tables")
  406. sqlite_tables = yield self.sqlite_store._simple_select_onecol(
  407. table="sqlite_master",
  408. keyvalues={
  409. "type": "table",
  410. },
  411. retcol="name",
  412. )
  413. postgres_tables = yield self.postgres_store._simple_select_onecol(
  414. table="information_schema.tables",
  415. keyvalues={},
  416. retcol="distinct table_name",
  417. )
  418. tables = set(sqlite_tables) & set(postgres_tables)
  419. self.progress.set_state("Creating tables")
  420. logger.info("Found %d tables", len(tables))
  421. def create_port_table(txn):
  422. txn.execute(
  423. "CREATE TABLE port_from_sqlite3 ("
  424. " table_name varchar(100) NOT NULL UNIQUE,"
  425. " forward_rowid bigint NOT NULL,"
  426. " backward_rowid bigint NOT NULL"
  427. ")"
  428. )
  429. # The old port script created a table with just a "rowid" column.
  430. # We want people to be able to rerun this script from an old port
  431. # so that they can pick up any missing events that were not
  432. # ported across.
  433. def alter_table(txn):
  434. txn.execute(
  435. "ALTER TABLE IF EXISTS port_from_sqlite3"
  436. " RENAME rowid TO forward_rowid"
  437. )
  438. txn.execute(
  439. "ALTER TABLE IF EXISTS port_from_sqlite3"
  440. " ADD backward_rowid bigint NOT NULL DEFAULT 0"
  441. )
  442. try:
  443. yield self.postgres_store.runInteraction(
  444. "alter_table", alter_table
  445. )
  446. except Exception as e:
  447. logger.info("Failed to create port table: %s", e)
  448. try:
  449. yield self.postgres_store.runInteraction(
  450. "create_port_table", create_port_table
  451. )
  452. except Exception as e:
  453. logger.info("Failed to create port table: %s", e)
  454. self.progress.set_state("Setting up")
  455. # Set up tables.
  456. setup_res = yield defer.gatherResults(
  457. [
  458. self.setup_table(table)
  459. for table in tables
  460. if table not in ["schema_version", "applied_schema_deltas"]
  461. and not table.startswith("sqlite_")
  462. ],
  463. consumeErrors=True,
  464. )
  465. # Process tables.
  466. yield defer.gatherResults(
  467. [
  468. self.handle_table(*res)
  469. for res in setup_res
  470. ],
  471. consumeErrors=True,
  472. )
  473. yield self.handle_state_group_seq()
  474. self.progress.done()
  475. except:
  476. global end_error_exec_info
  477. end_error_exec_info = sys.exc_info()
  478. logger.exception("")
  479. finally:
  480. reactor.stop()
  481. def _convert_rows(self, table, headers, rows):
  482. bool_col_names = BOOLEAN_COLUMNS.get(table, [])
  483. bool_cols = [
  484. i for i, h in enumerate(headers) if h in bool_col_names
  485. ]
  486. class BadValueException(Exception):
  487. pass
  488. def conv(j, col):
  489. if j in bool_cols:
  490. return bool(col)
  491. elif isinstance(col, basestring) and "\0" in col:
  492. logger.warn("DROPPING ROW: NUL value in table %s col %s: %r", table, headers[j], col)
  493. raise BadValueException();
  494. return col
  495. outrows = []
  496. for i, row in enumerate(rows):
  497. try:
  498. outrows.append(tuple(
  499. conv(j, col)
  500. for j, col in enumerate(row)
  501. if j > 0
  502. ))
  503. except BadValueException:
  504. pass
  505. return outrows
  506. @defer.inlineCallbacks
  507. def _setup_sent_transactions(self):
  508. # Only save things from the last day
  509. yesterday = int(time.time() * 1000) - 86400000
  510. # And save the max transaction id from each destination
  511. select = (
  512. "SELECT rowid, * FROM sent_transactions WHERE rowid IN ("
  513. "SELECT max(rowid) FROM sent_transactions"
  514. " GROUP BY destination"
  515. ")"
  516. )
  517. def r(txn):
  518. txn.execute(select)
  519. rows = txn.fetchall()
  520. headers = [column[0] for column in txn.description]
  521. ts_ind = headers.index('ts')
  522. return headers, [r for r in rows if r[ts_ind] < yesterday]
  523. headers, rows = yield self.sqlite_store.runInteraction(
  524. "select", r,
  525. )
  526. rows = self._convert_rows("sent_transactions", headers, rows)
  527. inserted_rows = len(rows)
  528. if inserted_rows:
  529. max_inserted_rowid = max(r[0] for r in rows)
  530. def insert(txn):
  531. self.postgres_store.insert_many_txn(
  532. txn, "sent_transactions", headers[1:], rows
  533. )
  534. yield self.postgres_store.execute(insert)
  535. else:
  536. max_inserted_rowid = 0
  537. def get_start_id(txn):
  538. txn.execute(
  539. "SELECT rowid FROM sent_transactions WHERE ts >= ?"
  540. " ORDER BY rowid ASC LIMIT 1",
  541. (yesterday,)
  542. )
  543. rows = txn.fetchall()
  544. if rows:
  545. return rows[0][0]
  546. else:
  547. return 1
  548. next_chunk = yield self.sqlite_store.execute(get_start_id)
  549. next_chunk = max(max_inserted_rowid + 1, next_chunk)
  550. yield self.postgres_store._simple_insert(
  551. table="port_from_sqlite3",
  552. values={
  553. "table_name": "sent_transactions",
  554. "forward_rowid": next_chunk,
  555. "backward_rowid": 0,
  556. }
  557. )
  558. def get_sent_table_size(txn):
  559. txn.execute(
  560. "SELECT count(*) FROM sent_transactions"
  561. " WHERE ts >= ?",
  562. (yesterday,)
  563. )
  564. size, = txn.fetchone()
  565. return int(size)
  566. remaining_count = yield self.sqlite_store.execute(
  567. get_sent_table_size
  568. )
  569. total_count = remaining_count + inserted_rows
  570. defer.returnValue((next_chunk, inserted_rows, total_count))
  571. @defer.inlineCallbacks
  572. def _get_remaining_count_to_port(self, table, forward_chunk, backward_chunk):
  573. frows = yield self.sqlite_store.execute_sql(
  574. "SELECT count(*) FROM %s WHERE rowid >= ?" % (table,),
  575. forward_chunk,
  576. )
  577. brows = yield self.sqlite_store.execute_sql(
  578. "SELECT count(*) FROM %s WHERE rowid <= ?" % (table,),
  579. backward_chunk,
  580. )
  581. defer.returnValue(frows[0][0] + brows[0][0])
  582. @defer.inlineCallbacks
  583. def _get_already_ported_count(self, table):
  584. rows = yield self.postgres_store.execute_sql(
  585. "SELECT count(*) FROM %s" % (table,),
  586. )
  587. defer.returnValue(rows[0][0])
  588. @defer.inlineCallbacks
  589. def _get_total_count_to_port(self, table, forward_chunk, backward_chunk):
  590. remaining, done = yield defer.gatherResults(
  591. [
  592. self._get_remaining_count_to_port(table, forward_chunk, backward_chunk),
  593. self._get_already_ported_count(table),
  594. ],
  595. consumeErrors=True,
  596. )
  597. remaining = int(remaining) if remaining else 0
  598. done = int(done) if done else 0
  599. defer.returnValue((done, remaining + done))
  600. ##############################################
  601. ###### The following is simply UI stuff ######
  602. ##############################################
  603. class Progress(object):
  604. """Used to report progress of the port
  605. """
  606. def __init__(self):
  607. self.tables = {}
  608. self.start_time = int(time.time())
  609. def add_table(self, table, cur, size):
  610. self.tables[table] = {
  611. "start": cur,
  612. "num_done": cur,
  613. "total": size,
  614. "perc": int(cur * 100 / size),
  615. }
  616. def update(self, table, num_done):
  617. data = self.tables[table]
  618. data["num_done"] = num_done
  619. data["perc"] = int(num_done * 100 / data["total"])
  620. def done(self):
  621. pass
  622. class CursesProgress(Progress):
  623. """Reports progress to a curses window
  624. """
  625. def __init__(self, stdscr):
  626. self.stdscr = stdscr
  627. curses.use_default_colors()
  628. curses.curs_set(0)
  629. curses.init_pair(1, curses.COLOR_RED, -1)
  630. curses.init_pair(2, curses.COLOR_GREEN, -1)
  631. self.last_update = 0
  632. self.finished = False
  633. self.total_processed = 0
  634. self.total_remaining = 0
  635. super(CursesProgress, self).__init__()
  636. def update(self, table, num_done):
  637. super(CursesProgress, self).update(table, num_done)
  638. self.total_processed = 0
  639. self.total_remaining = 0
  640. for table, data in self.tables.items():
  641. self.total_processed += data["num_done"] - data["start"]
  642. self.total_remaining += data["total"] - data["num_done"]
  643. self.render()
  644. def render(self, force=False):
  645. now = time.time()
  646. if not force and now - self.last_update < 0.2:
  647. # reactor.callLater(1, self.render)
  648. return
  649. self.stdscr.clear()
  650. rows, cols = self.stdscr.getmaxyx()
  651. duration = int(now) - int(self.start_time)
  652. minutes, seconds = divmod(duration, 60)
  653. duration_str = '%02dm %02ds' % (minutes, seconds,)
  654. if self.finished:
  655. status = "Time spent: %s (Done!)" % (duration_str,)
  656. else:
  657. if self.total_processed > 0:
  658. left = float(self.total_remaining) / self.total_processed
  659. est_remaining = (int(now) - self.start_time) * left
  660. est_remaining_str = '%02dm %02ds remaining' % divmod(est_remaining, 60)
  661. else:
  662. est_remaining_str = "Unknown"
  663. status = (
  664. "Time spent: %s (est. remaining: %s)"
  665. % (duration_str, est_remaining_str,)
  666. )
  667. self.stdscr.addstr(
  668. 0, 0,
  669. status,
  670. curses.A_BOLD,
  671. )
  672. max_len = max([len(t) for t in self.tables.keys()])
  673. left_margin = 5
  674. middle_space = 1
  675. items = self.tables.items()
  676. items.sort(
  677. key=lambda i: (i[1]["perc"], i[0]),
  678. )
  679. for i, (table, data) in enumerate(items):
  680. if i + 2 >= rows:
  681. break
  682. perc = data["perc"]
  683. color = curses.color_pair(2) if perc == 100 else curses.color_pair(1)
  684. self.stdscr.addstr(
  685. i + 2, left_margin + max_len - len(table),
  686. table,
  687. curses.A_BOLD | color,
  688. )
  689. size = 20
  690. progress = "[%s%s]" % (
  691. "#" * int(perc * size / 100),
  692. " " * (size - int(perc * size / 100)),
  693. )
  694. self.stdscr.addstr(
  695. i + 2, left_margin + max_len + middle_space,
  696. "%s %3d%% (%d/%d)" % (progress, perc, data["num_done"], data["total"]),
  697. )
  698. if self.finished:
  699. self.stdscr.addstr(
  700. rows - 1, 0,
  701. "Press any key to exit...",
  702. )
  703. self.stdscr.refresh()
  704. self.last_update = time.time()
  705. def done(self):
  706. self.finished = True
  707. self.render(True)
  708. self.stdscr.getch()
  709. def set_state(self, state):
  710. self.stdscr.clear()
  711. self.stdscr.addstr(
  712. 0, 0,
  713. state + "...",
  714. curses.A_BOLD,
  715. )
  716. self.stdscr.refresh()
  717. class TerminalProgress(Progress):
  718. """Just prints progress to the terminal
  719. """
  720. def update(self, table, num_done):
  721. super(TerminalProgress, self).update(table, num_done)
  722. data = self.tables[table]
  723. print "%s: %d%% (%d/%d)" % (
  724. table, data["perc"],
  725. data["num_done"], data["total"],
  726. )
  727. def set_state(self, state):
  728. print state + "..."
  729. ##############################################
  730. ##############################################
  731. if __name__ == "__main__":
  732. parser = argparse.ArgumentParser(
  733. description="A script to port an existing synapse SQLite database to"
  734. " a new PostgreSQL database."
  735. )
  736. parser.add_argument("-v", action='store_true')
  737. parser.add_argument(
  738. "--sqlite-database", required=True,
  739. help="The snapshot of the SQLite database file. This must not be"
  740. " currently used by a running synapse server"
  741. )
  742. parser.add_argument(
  743. "--postgres-config", type=argparse.FileType('r'), required=True,
  744. help="The database config file for the PostgreSQL database"
  745. )
  746. parser.add_argument(
  747. "--curses", action='store_true',
  748. help="display a curses based progress UI"
  749. )
  750. parser.add_argument(
  751. "--batch-size", type=int, default=1000,
  752. help="The number of rows to select from the SQLite table each"
  753. " iteration [default=1000]",
  754. )
  755. args = parser.parse_args()
  756. logging_config = {
  757. "level": logging.DEBUG if args.v else logging.INFO,
  758. "format": "%(asctime)s - %(name)s - %(lineno)d - %(levelname)s - %(message)s"
  759. }
  760. if args.curses:
  761. logging_config["filename"] = "port-synapse.log"
  762. logging.basicConfig(**logging_config)
  763. sqlite_config = {
  764. "name": "sqlite3",
  765. "args": {
  766. "database": args.sqlite_database,
  767. "cp_min": 1,
  768. "cp_max": 1,
  769. "check_same_thread": False,
  770. },
  771. }
  772. postgres_config = yaml.safe_load(args.postgres_config)
  773. if "database" in postgres_config:
  774. postgres_config = postgres_config["database"]
  775. if "name" not in postgres_config:
  776. sys.stderr.write("Malformed database config: no 'name'")
  777. sys.exit(2)
  778. if postgres_config["name"] != "psycopg2":
  779. sys.stderr.write("Database must use 'psycopg2' connector.")
  780. sys.exit(3)
  781. def start(stdscr=None):
  782. if stdscr:
  783. progress = CursesProgress(stdscr)
  784. else:
  785. progress = TerminalProgress()
  786. porter = Porter(
  787. sqlite_config=sqlite_config,
  788. postgres_config=postgres_config,
  789. progress=progress,
  790. batch_size=args.batch_size,
  791. )
  792. reactor.callWhenRunning(porter.run)
  793. reactor.run()
  794. if args.curses:
  795. curses.wrapper(start)
  796. else:
  797. start()
  798. if end_error_exec_info:
  799. exc_type, exc_value, exc_traceback = end_error_exec_info
  800. traceback.print_exception(exc_type, exc_value, exc_traceback)