From d245bbd600811762a74544aeff6086a6c94dd7f3 Mon Sep 17 00:00:00 2001
From: Chris Lovering <chris.lovering.95@gmail.com>
Date: Tue, 7 Sep 2021 09:59:18 +0100
Subject: Create and populate metricity in a single connection

---
 manage.py         | 35 ++++++++++-------------------------
 postgres/init.sql | 21 +++++++++++++++++++++
 2 files changed, 31 insertions(+), 25 deletions(-)

diff --git a/manage.py b/manage.py
index cb4eb567..308f917d 100755
--- a/manage.py
+++ b/manage.py
@@ -152,36 +152,21 @@ class SiteManager:
         from psycopg2.extensions import ISOLATION_LEVEL_AUTOCOMMIT
 
         db_url_parts = SiteManager.parse_db_url(os.environ["DATABASE_URL"])
-        db_connection_kwargs = {
-            "host": db_url_parts.hostname,
-            "port": db_url_parts.port,
-            "user": db_url_parts.username,
-            "password": db_url_parts.password,
-        }
-        # Connect to pysite first to create metricity db
         conn = psycopg2.connect(
-            database=db_url_parts.path[1:],
-            **db_connection_kwargs
+            host=db_url_parts.hostname,
+            port=db_url_parts.port,
+            user=db_url_parts.username,
+            password=db_url_parts.password,
+            database=db_url_parts.path[1:]
         )
+        # Required to create a db from `cursor.execute()`
         conn.set_isolation_level(ISOLATION_LEVEL_AUTOCOMMIT)
 
-        with conn.cursor() as cursor:
-            cursor.execute("SELECT 1 FROM pg_catalog.pg_database WHERE datname = 'metricity'")
-            exists = cursor.fetchone()
-            if exists:
-                # Assume metricity is already populated if it exists
-                return
-            print("Creating metricity relations and populating with some data.")
-            cursor.execute("CREATE DATABASE metricity")
-        conn.close()
-
-        # Switch connection to metricity and initialise some data
-        conn = psycopg2.connect(
-            database="metricity",
-            **db_connection_kwargs
-        )
         with conn.cursor() as cursor, open("postgres/init.sql", encoding="utf-8") as f:
-            cursor.execute(f.read())
+            cursor.execute(
+                f.read(),
+                ("metricity", db_url_parts.username, db_url_parts.password)
+            )
         conn.close()
 
     def prepare_server(self) -> None:
diff --git a/postgres/init.sql b/postgres/init.sql
index ea748480..ee1acd8d 100644
--- a/postgres/init.sql
+++ b/postgres/init.sql
@@ -1,3 +1,24 @@
+DO
+$do$
+DECLARE
+  _db TEXT := %s;
+  _user TEXT := %s;
+  _password TEXT := %s;
+BEGIN
+  CREATE EXTENSION IF NOT EXISTS dblink;
+  IF EXISTS (SELECT 1 FROM pg_database WHERE datname = _db) THEN
+    RAISE NOTICE 'Database already exists';
+  ELSE
+    PERFORM dblink_connect(
+        'host=localhost user=' || _user ||
+        ' password=' || _password ||
+        ' dbname=' || current_database()
+    );
+    PERFORM dblink_exec('CREATE DATABASE ' || _db);
+  END IF;
+END
+$do$;
+
 CREATE TABLE IF NOT EXISTS users (
     id varchar,
     joined_at timestamp,
-- 
cgit v1.2.3