Add `--run-background-updates` option to `update_database` script. (#10954)

Signed-off-by: Nick Barrett <nick@beeper.com>

1 files changed, 116 insertions, 0 deletions

diff --git a/scripts/update_synapse_database b/scripts/update_synapse_database
new file mode 100755
index 0000000000..26b29b0b45
--- /dev/null
+++ b/scripts/update_synapse_database
@@ -0,0 +1,116 @@
+#!/usr/bin/env python
+# 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 logging
+import sys
+
+import yaml
+
+from twisted.internet import defer, reactor
+
+import synapse
+from synapse.config.homeserver import HomeServerConfig
+from synapse.metrics.background_process_metrics import run_as_background_process
+from synapse.server import HomeServer
+from synapse.storage import DataStore
+from synapse.util.versionstring import get_version_string
+
+logger = logging.getLogger("update_database")
+
+
+class MockHomeserver(HomeServer):
+    DATASTORE_CLASS = DataStore
+
+    def __init__(self, config, **kwargs):
+        super(MockHomeserver, self).__init__(
+            config.server_name, reactor=reactor, config=config, **kwargs
+        )
+
+        self.version_string = "Synapse/" + get_version_string(synapse)
+
+
+def run_background_updates(hs):
+    store = hs.get_datastore()
+
+    async def run_background_updates():
+        await store.db_pool.updates.run_background_updates(sleep=False)
+        # Stop the reactor to exit the script once every background update is run.
+        reactor.stop()
+
+    def run():
+        # Apply all background updates on the database.
+        defer.ensureDeferred(
+            run_as_background_process("background_updates", run_background_updates)
+        )
+
+    reactor.callWhenRunning(run)
+
+    reactor.run()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description=(
+            "Updates a synapse database to the latest schema and optionally runs background updates"
+            " on it."
+        )
+    )
+    parser.add_argument("-v", action="store_true")
+    parser.add_argument(
+        "--database-config",
+        type=argparse.FileType("r"),
+        required=True,
+        help="Synapse configuration file, giving the details of the database to be updated",
+    )
+    parser.add_argument(
+        "--run-background-updates",
+        action="store_true",
+        required=False,
+        help="run background updates after upgrading the database schema",
+    )
+
+    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",
+    }
+
+    logging.basicConfig(**logging_config)
+
+    # Load, process and sanity-check the config.
+    hs_config = yaml.safe_load(args.database_config)
+
+    if "database" not in hs_config:
+        sys.stderr.write("The configuration file must have a 'database' section.\n")
+        sys.exit(4)
+
+    config = HomeServerConfig()
+    config.parse_config_dict(hs_config, "", "")
+
+    # Instantiate and initialise the homeserver object.
+    hs = MockHomeserver(config)
+
+    # Setup instantiates the store within the homeserver object and updates the
+    # DB.
+    hs.setup()
+
+    if args.run_background_updates:
+        run_background_updates(hs)
+
+
+if __name__ == "__main__":
+    main()