From bc0943c5698bda5863daf0e866863b10a54b4dbd Mon Sep 17 00:00:00 2001
From: Jesse Wierzbinski <contact@cpluspatch.com>
Date: Mon, 26 Aug 2024 18:04:22 +0200
Subject: [PATCH] feat(database): :sparkles: Implement read replicas for
 database

---
 config/config.example.toml             |  8 +++
 config/config.schema.json              | 33 +++++++++++++
 drizzle/db.ts                          | 68 +++++++++++++++++++-------
 packages/config-manager/config.type.ts | 16 ++++++
 4 files changed, 106 insertions(+), 19 deletions(-)

diff --git a/config/config.example.toml b/config/config.example.toml
index 71594784..1fc34e42 100644
--- a/config/config.example.toml
+++ b/config/config.example.toml
@@ -10,6 +10,14 @@ username = "versia"
 password = "mycoolpassword"
 database = "versia"
 
+# Add any eventual read-only database replicas here
+# [[database.replicas]]
+# host = "other-host"
+# port = 5432
+# username = "versia"
+# password = "mycoolpassword2"
+# database = "replica1"
+
 [redis.queue]
 # Redis instance for storing the federation queue
 # Required for federation
diff --git a/config/config.schema.json b/config/config.schema.json
index ffafadef..2f094f18 100644
--- a/config/config.schema.json
+++ b/config/config.schema.json
@@ -27,6 +27,39 @@
                     "type": "string",
                     "minLength": 1,
                     "default": "versia"
+                },
+                "replicas": {
+                    "type": "array",
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "host": {
+                                "type": "string",
+                                "minLength": 1
+                            },
+                            "port": {
+                                "type": "integer",
+                                "minimum": 1,
+                                "maximum": 65535,
+                                "default": 5432
+                            },
+                            "username": {
+                                "type": "string",
+                                "minLength": 1
+                            },
+                            "password": {
+                                "type": "string",
+                                "default": ""
+                            },
+                            "database": {
+                                "type": "string",
+                                "minLength": 1,
+                                "default": "versia"
+                            }
+                        },
+                        "required": ["host", "username"],
+                        "additionalProperties": false
+                    }
                 }
             },
             "required": ["username"],
diff --git a/drizzle/db.ts b/drizzle/db.ts
index 8f4ae087..8c029a0b 100644
--- a/drizzle/db.ts
+++ b/drizzle/db.ts
@@ -1,11 +1,13 @@
 import { getLogger } from "@logtape/logtape";
-import { drizzle } from "drizzle-orm/node-postgres";
+import chalk from "chalk";
+import { type NodePgDatabase, drizzle } from "drizzle-orm/node-postgres";
+import { withReplicas } from "drizzle-orm/pg-core";
 import { migrate } from "drizzle-orm/postgres-js/migrator";
-import { Client } from "pg";
+import { Pool } from "pg";
 import { config } from "~/packages/config-manager";
 import * as schema from "./schema";
 
-export const client = new Client({
+const primaryDb = new Pool({
     host: config.database.host,
     port: Number(config.database.port),
     user: config.database.username,
@@ -13,24 +15,54 @@ export const client = new Client({
     database: config.database.database,
 });
 
+const replicas =
+    config.database.replicas?.map(
+        (replica) =>
+            new Pool({
+                host: replica.host,
+                port: Number(replica.port),
+                user: replica.username,
+                password: replica.password,
+                database: replica.database,
+            }),
+    ) ?? [];
+
+export const db =
+    (replicas.length ?? 0) > 0
+        ? withReplicas(
+              drizzle(primaryDb, { schema }),
+              replicas.map((r) => drizzle(r, { schema })) as [
+                  NodePgDatabase<typeof schema>,
+                  ...NodePgDatabase<typeof schema>[],
+              ],
+          )
+        : drizzle(primaryDb, { schema });
+
 export const setupDatabase = async (info = true) => {
     const logger = getLogger("database");
 
-    try {
-        await client.connect();
-    } catch (e) {
-        if (
-            (e as Error).message ===
-            "Client has already been connected. You cannot reuse a client."
-        ) {
-            return;
+    for (const dbPool of [primaryDb, ...replicas]) {
+        try {
+            await dbPool.connect();
+        } catch (e) {
+            if (
+                (e as Error).message ===
+                "Client has already been connected. You cannot reuse a client."
+            ) {
+                return;
+            }
+
+            logger.fatal`${e}`;
+            logger.fatal`Failed to connect to database ${chalk.bold(
+                // Index of the database in the array
+                replicas.indexOf(dbPool) === -1
+                    ? "primary"
+                    : `replica-${replicas.indexOf(dbPool)}`,
+            )}. Please check your configuration.`;
+
+            // Hang until Ctrl+C is pressed
+            await Bun.sleep(Number.POSITIVE_INFINITY);
         }
-
-        logger.fatal`${e}`;
-        logger.fatal`Failed to connect to database. Please check your configuration.`;
-
-        // Hang until Ctrl+C is pressed
-        await Bun.sleep(Number.POSITIVE_INFINITY);
     }
 
     // Migrate the database
@@ -50,5 +82,3 @@ export const setupDatabase = async (info = true) => {
 
     info && logger.info`Database migrated`;
 };
-
-export const db = drizzle(client, { schema });
diff --git a/packages/config-manager/config.type.ts b/packages/config-manager/config.type.ts
index e1c9cd00..005bcced 100644
--- a/packages/config-manager/config.type.ts
+++ b/packages/config-manager/config.type.ts
@@ -33,6 +33,22 @@ export const configValidator = z.object({
         username: z.string().min(1),
         password: z.string().default(""),
         database: z.string().min(1).default("versia"),
+        replicas: z
+            .array(
+                z.object({
+                    host: z.string().min(1),
+                    port: z
+                        .number()
+                        .int()
+                        .min(1)
+                        .max(2 ** 16 - 1)
+                        .default(5432),
+                    username: z.string().min(1),
+                    password: z.string().default(""),
+                    database: z.string().min(1).default("versia"),
+                }),
+            )
+            .optional(),
     }),
     redis: z.object({
         queue: z