Move database preparing code out of homserver.py into storage where it belongs

1 files changed, 61 insertions, 0 deletions

diff --git a/synapse/storage/__init__.py b/synapse/storage/__init__.py
index ad2a484c16..2543fb12b7 100644
--- a/synapse/storage/__init__.py
+++ b/synapse/storage/__init__.py
@@ -43,10 +43,28 @@ from .keys import KeyStore
 import json
 import logging
 import os
+import sqlite3
 
 
 logger = logging.getLogger(__name__)
 
+
+SCHEMAS = [
+    "transactions",
+    "pdu",
+    "users",
+    "profiles",
+    "presence",
+    "im",
+    "room_aliases",
+]
+
+
+# Remember to update this number every time an incompatible change is made to
+# database schema files, so the users will be informed on server restarts.
+SCHEMA_VERSION = 3
+
+
 class _RollbackButIsFineException(Exception):
     """ This exception is used to rollback a transaction without implying
     something went wrong.
@@ -350,3 +368,46 @@ def read_schema(schema):
     """
     with open(schema_path(schema)) as schema_file:
         return schema_file.read()
+
+
+def prepare_database(db_name):
+    """ Set up all the dbs. Since all the *.sql have IF NOT EXISTS, so we
+    don't have to worry about overwriting existing content.
+    """
+    logging.info("Preparing database: %s...", db_name)
+
+    with sqlite3.connect(db_name) as db_conn:
+        c = db_conn.cursor()
+        c.execute("PRAGMA user_version")
+        row = c.fetchone()
+
+        if row and row[0]:
+            user_version = row[0]
+
+            if user_version > SCHEMA_VERSION:
+                raise ValueError("Cannot use this database as it is too " +
+                    "new for the server to understand"
+                )
+            elif user_version < SCHEMA_VERSION:
+                logging.info("Upgrading database from version %d",
+                    user_version
+                )
+
+                # Run every version since after the current version.
+                for v in range(user_version + 1, SCHEMA_VERSION + 1):
+                    sql_script = read_schema("delta/v%d" % (v))
+                    c.executescript(sql_script)
+
+                db_conn.commit()
+
+        else:
+            for sql_loc in SCHEMAS:
+                sql_script = read_schema(sql_loc)
+
+                c.executescript(sql_script)
+            db_conn.commit()
+            c.execute("PRAGMA user_version = %d" % SCHEMA_VERSION)
+
+        c.close()
+
+    logging.info("Database prepared in %s.", db_name)