From b21487300c54242cf18b106e681d1b11a61b53ba Mon Sep 17 00:00:00 2001
From: Olivier 'reivilibre <oliverw@matrix.org>
Date: Mon, 9 Sep 2024 16:15:52 +0100
Subject: [PATCH] Add `application_name`s to the postgres connections based on
 the worker names

---
 docker/conf-workers/worker.yaml.j2    | 15 +++++++++++++++
 docker/conf/homeserver.yaml           |  1 +
 docker/configure_workers_and_start.py |  6 ++++++
 3 files changed, 22 insertions(+)

diff --git a/docker/conf-workers/worker.yaml.j2 b/docker/conf-workers/worker.yaml.j2
index a750eec8a8..9aebc87cca 100644
--- a/docker/conf-workers/worker.yaml.j2
+++ b/docker/conf-workers/worker.yaml.j2
@@ -27,4 +27,19 @@ worker_listeners:
 
 worker_log_config: {{ worker_log_config_filepath }}
 
+database:
+  name: "psycopg2"
+  args:
+    user: "{{ POSTGRES_USER or "synapse" }}"
+    password: "{{ POSTGRES_PASSWORD }}"
+    database: "{{ POSTGRES_DB or "synapse" }}"
+{% if not SYNAPSE_USE_UNIX_SOCKET %}
+{# Synapse will use a default unix socket for Postgres when host/port is not specified (behavior from `psycopg2`). #}
+    host: "{{ POSTGRES_HOST or "db" }}"
+    port: "{{ POSTGRES_PORT or "5432" }}"
+{% endif %}
+    cp_min: {{ POSTGRES_CP_MIN or 5 }}
+    cp_max: {{ POSTGRES_CP_MAX or 10 }}
+    application_name: "{{ name }}"
+
 {{ worker_extra_conf }}
diff --git a/docker/conf/homeserver.yaml b/docker/conf/homeserver.yaml
index 76b0762df6..be45f4ffd0 100644
--- a/docker/conf/homeserver.yaml
+++ b/docker/conf/homeserver.yaml
@@ -69,6 +69,7 @@ database:
 {% endif %}
     cp_min: {{ POSTGRES_CP_MIN or 5 }}
     cp_max: {{ POSTGRES_CP_MAX or 10 }}
+    application_name: synapse
 {% else %}
 database:
   name: "sqlite3"
diff --git a/docker/configure_workers_and_start.py b/docker/configure_workers_and_start.py
index c52b175bad..86e9c44f0f 100755
--- a/docker/configure_workers_and_start.py
+++ b/docker/configure_workers_and_start.py
@@ -962,11 +962,17 @@ def generate_worker_files(
             environ, worker_name, workers_config_dir, template_dir, data_dir
         )
 
+        extra_env: Dict[str, str] = {}
+        for extra_env_var in ("POSTGRES_HOST", "POSTGRES_PORT", "POSTGRES_USER", "POSTGRES_PASSWORD", "POSTGRES_SOCKET", "POSTGRES_DB", "SYNAPSE_USE_UNIX_SOCKET", "POSTGRES_CP_MIN", "POSTGRES_CP_MAX"):
+            if extra_env_var in os.environ:
+                extra_env[extra_env_var] = os.environ[extra_env_var]
+
         # Then a worker config file
         convert(
             os.path.join(template_dir, "worker.yaml.j2"),
             os.path.join(workers_config_dir, f"{worker_name}.yaml"),
             **worker_config,
+            **extra_env,
             worker_log_config_filepath=log_config_filepath,
             using_unix_sockets=using_unix_sockets,
         )