1 files changed, 50 insertions, 14 deletions

diff --git a/synapse/app/homeserver.py b/synapse/app/homeserver.py
index 27e53a9e56..93500dd791 100755
--- a/synapse/app/homeserver.py
+++ b/synapse/app/homeserver.py
@@ -17,9 +17,8 @@
 import sys
 sys.dont_write_bytecode = True
 
-from synapse.storage import (
-    prepare_database, prepare_sqlite3_database, UpgradeDatabaseException,
-)
+from synapse.storage import UpgradeDatabaseException
+from synapse.storage.engines import create_engine
 
 from synapse.server import HomeServer
 
@@ -57,9 +56,10 @@ import os
 import re
 import resource
 import subprocess
-import sqlite3
+import yaml
+
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger("synapse.app.homeserver")
 
 
 class SynapseHomeServer(HomeServer):
@@ -103,13 +103,11 @@ class SynapseHomeServer(HomeServer):
             return None
 
     def build_db_pool(self):
+        name = self.db_config["name"]
+
         return adbapi.ConnectionPool(
-            "sqlite3", self.get_db_name(),
-            check_same_thread=False,
-            cp_min=1,
-            cp_max=1,
-            cp_openfun=prepare_database,  # Prepare the database for each conn
-                                          # so that :memory: sqlite works
+            name,
+            **self.db_config.get("args", {})
         )
 
     def create_resource_tree(self, redirect_root_to_web_client):
@@ -352,15 +350,46 @@ def setup(config_options):
 
     tls_context_factory = context_factory.ServerContextFactory(config)
 
+    if config.database_config:
+        with open(config.database_config, 'r') as f:
+            db_config = yaml.safe_load(f)
+    else:
+        db_config = {
+            "name": "sqlite3",
+            "args": {
+                "database": config.database_path,
+            },
+        }
+
+    db_config = {
+        k: v for k, v in db_config.items()
+    }
+
+    name = db_config.get("name", None)
+    if name == "psycopg2":
+        pass
+    elif name == "sqlite3":
+        db_config.setdefault("args", {}).update({
+            "cp_min": 1,
+            "cp_max": 1,
+        })
+    else:
+        raise RuntimeError("Unsupported database type '%s'" % (name,))
+
+    database_engine = create_engine(name)
+    db_config["args"]["cp_openfun"] = database_engine.on_new_connection
+
     hs = SynapseHomeServer(
         config.server_name,
         domain_with_port=domain_with_port,
         upload_dir=os.path.abspath("uploads"),
         db_name=config.database_path,
+        db_config=db_config,
         tls_context_factory=tls_context_factory,
         config=config,
         content_addr=config.content_addr,
         version_string=version_string,
+        database_engine=database_engine,
     )
 
     hs.create_resource_tree(
@@ -372,9 +401,16 @@ def setup(config_options):
     logger.info("Preparing database: %s...", db_name)
 
     try:
-        with sqlite3.connect(db_name) as db_conn:
-            prepare_sqlite3_database(db_conn)
-            prepare_database(db_conn)
+        db_conn = database_engine.module.connect(
+            **{
+                k: v for k, v in db_config.get("args", {}).items()
+                if not k.startswith("cp_")
+            }
+        )
+
+        database_engine.prepare_database(db_conn)
+
+        db_conn.commit()
     except UpgradeDatabaseException:
         sys.stderr.write(
             "\nFailed to upgrade database.\n"