From 617375fb46b4166f2558859f08793da64b8848ca Mon Sep 17 00:00:00 2001
From: Dave Arter <davea@mysociety.org>
Date: Thu, 23 Nov 2023 19:55:43 +0000
Subject: [PATCH] Add support for read-only database replica

Lifted from https://github.com/mysociety/mapit.mysociety.org/pull/90
---
 labour_project/middleware.py | 22 ++++++++++++++++++
 labour_project/multidb.py    | 16 +++++++++++++
 labour_project/settings.py   | 45 ++++++++++++++++++++++++++++++++++++
 3 files changed, 83 insertions(+)
 create mode 100644 labour_project/middleware.py
 create mode 100644 labour_project/multidb.py

diff --git a/labour_project/middleware.py b/labour_project/middleware.py
new file mode 100644
index 0000000..1675646
--- /dev/null
+++ b/labour_project/middleware.py
@@ -0,0 +1,22 @@
+from .multidb import use_primary
+
+READ_METHODS = frozenset(['GET', 'HEAD', 'OPTIONS', 'TRACE'])
+
+
+def force_primary_middleware(get_response):
+    """On a non-read request (e.g. POST), always use the primary database, and
+    set a short-lived cookie. Then on any request made including that cookie,
+    also always use the primary database. This is so we don't need to worry
+    about any replication lag for reads made immediately after writes, because
+    we force those reads to go to the primary."""
+    def middleware(request):
+        use_primary('primary_forced' in request.COOKIES or request.method not in READ_METHODS)
+
+        response = get_response(request)
+
+        if request.method not in READ_METHODS:
+            response.set_cookie('primary_forced', value='y', max_age=10)
+
+        return response
+
+    return middleware
diff --git a/labour_project/multidb.py b/labour_project/multidb.py
new file mode 100644
index 0000000..8c7aeec
--- /dev/null
+++ b/labour_project/multidb.py
@@ -0,0 +1,16 @@
+"""
+A getter/setter for a thread-local variable for using the primary database.
+"""
+
+import threading
+
+_locals = threading.local()
+
+
+def use_primary(val=None):
+    """If passed no argument, return the current value (or False if unset);
+    if passed an argument, set the variable to that."""
+    if val is None:
+        return getattr(_locals, 'primary', False)
+    else:
+        _locals.primary = val
diff --git a/labour_project/settings.py b/labour_project/settings.py
index 71d1a63..883736c 100644
--- a/labour_project/settings.py
+++ b/labour_project/settings.py
@@ -1,4 +1,5 @@
 import os
+import sys
 import yaml
 from .utils import skip_unreadable_post
 
@@ -97,6 +98,47 @@
     }
 }
 
+if config.get('MAPIT_DB_RO_HOST', '') and 'test' not in sys.argv:
+    DATABASES['replica'] = {
+        'ENGINE': 'django.contrib.gis.db.backends.postgis',
+        'NAME': config.get('MAPIT_DB_NAME', 'mapit'),
+        'USER': config.get('MAPIT_DB_USER', 'mapit'),
+        'PASSWORD': config.get('MAPIT_DB_PASS', ''),
+        'HOST': config['MAPIT_DB_RO_HOST'],
+        'PORT': config['MAPIT_DB_RO_PORT'],
+        # Should work, but does not appear to (hence sys.argv test above); see
+        # https://stackoverflow.com/questions/33941139/test-mirror-default-database-but-no-data
+        # 'TEST': {
+        #     'MIRROR': 'default',
+        # },
+    }
+
+    import random
+    from .multidb import use_primary
+
+    class PrimaryReplicaRouter(object):
+        """A basic primary/replica database router."""
+        def db_for_read(self, model, **hints):
+            """Randomly pick between default and replica databases, unless the
+            request (via middleware) demands we use the primary."""
+            if use_primary():
+                return 'default'
+            return random.choice(['default', 'replica'])
+
+        def db_for_write(self, model, **hints):
+            """Always write to the primary database."""
+            return 'default'
+
+        def allow_relation(self, obj1, obj2, **hints):
+            """Any relation between objects is allowed, as same data."""
+            return True
+
+        def allow_migrate(self, db, app_label, model_name=None, **hints):
+            """migrate is only ever called on the default database."""
+            return True
+
+    DATABASE_ROUTERS = ['labour_project.settings.PrimaryReplicaRouter']
+
 # Make this unique, and don't share it with anybody.
 SECRET_KEY = config.get('DJANGO_SECRET_KEY', '')
 
@@ -171,6 +213,9 @@
     'mapit.middleware.ViewExceptionMiddleware',
 ]
 
+if config.get('MAPIT_DB_RO_HOST', '') and 'test' not in sys.argv:
+    MIDDLEWARE.insert(0, 'labour_project.middleware.force_primary_middleware')
+
 if DEBUG:
     MIDDLEWARE.insert(0, 'debug_toolbar.middleware.DebugToolbarMiddleware')