From 4a34ee1f1e2ec6c6d6ddbd0adf9fd6e6bb0eac9a Mon Sep 17 00:00:00 2001
From: Rick Elrod <rick@elrod.me>
Date: Wed, 5 Jul 2023 13:41:47 -0500
Subject: [PATCH] Add optional pgbouncer to dev environment (#14083)

Signed-off-by: Rick Elrod <rick@elrod.me>
---
 Makefile                                      |  3 +++
 .../ansible/roles/sources/defaults/main.yml   |  5 +++++
 .../roles/sources/templates/database.py.j2    |  5 +++++
 .../sources/templates/docker-compose.yml.j2   | 19 +++++++++++++++++++
 4 files changed, 32 insertions(+)

diff --git a/Makefile b/Makefile
index ca3b134728..3e18e5e3e4 100644
--- a/Makefile
+++ b/Makefile
@@ -27,6 +27,8 @@ COLLECTION_TEMPLATE_VERSION ?= false
 # NOTE: This defaults the container image version to the branch that's active
 COMPOSE_TAG ?= $(GIT_BRANCH)
 MAIN_NODE_TYPE ?= hybrid
+# If set to true docker-compose will also start a pgbouncer instance and use it
+PGBOUNCER ?= false
 # If set to true docker-compose will also start a keycloak instance
 KEYCLOAK ?= false
 # If set to true docker-compose will also start an ldap instance
@@ -522,6 +524,7 @@ docker-compose-sources: .git/hooks/pre-commit
 	    -e control_plane_node_count=$(CONTROL_PLANE_NODE_COUNT) \
 	    -e execution_node_count=$(EXECUTION_NODE_COUNT) \
 	    -e minikube_container_group=$(MINIKUBE_CONTAINER_GROUP) \
+	    -e enable_pgbouncer=$(PGBOUNCER) \
 	    -e enable_keycloak=$(KEYCLOAK) \
 	    -e enable_ldap=$(LDAP) \
 	    -e enable_splunk=$(SPLUNK) \
diff --git a/tools/docker-compose/ansible/roles/sources/defaults/main.yml b/tools/docker-compose/ansible/roles/sources/defaults/main.yml
index 03e7d05bd3..b918104c6a 100644
--- a/tools/docker-compose/ansible/roles/sources/defaults/main.yml
+++ b/tools/docker-compose/ansible/roles/sources/defaults/main.yml
@@ -35,3 +35,8 @@ enable_splunk: false
 enable_grafana: false
 enable_prometheus: false
 scrape_interval: '5s'
+
+# pgbouncer
+enable_pgbouncer: false
+pgbouncer_port: 6432
+pgbouncer_max_pool_size: 70
diff --git a/tools/docker-compose/ansible/roles/sources/templates/database.py.j2 b/tools/docker-compose/ansible/roles/sources/templates/database.py.j2
index c1cd6fd992..76bebff159 100644
--- a/tools/docker-compose/ansible/roles/sources/templates/database.py.j2
+++ b/tools/docker-compose/ansible/roles/sources/templates/database.py.j2
@@ -5,7 +5,12 @@ DATABASES = {
         'NAME': "{{ pg_database }}",
         'USER': "{{ pg_username }}",
         'PASSWORD': "{{ pg_password }}",
+{% if enable_pgbouncer|bool %}
+        'HOST': "pgbouncer",
+        'PORT': "{{ pgbouncer_port }}",
+{% else %}
         'HOST': "{{ pg_hostname | default('postgres') }}",
         'PORT': "{{ pg_port }}",
+{% endif %}
     }
 }
diff --git a/tools/docker-compose/ansible/roles/sources/templates/docker-compose.yml.j2 b/tools/docker-compose/ansible/roles/sources/templates/docker-compose.yml.j2
index c605a1a8a3..e1942d0d37 100644
--- a/tools/docker-compose/ansible/roles/sources/templates/docker-compose.yml.j2
+++ b/tools/docker-compose/ansible/roles/sources/templates/docker-compose.yml.j2
@@ -201,6 +201,25 @@ services:
       POSTGRES_PASSWORD: {{ pg_password }}
     volumes:
       - "awx_db:/var/lib/postgresql/data"
+{% if enable_pgbouncer|bool %}
+  pgbouncer:
+    image: bitnami/pgbouncer:latest
+    container_name: tools_pgbouncer_1
+    hostname: pgbouncer
+    environment:
+      POSTGRESQL_USERNAME: {{ pg_username }}
+      POSTGRESQL_DATABASE: {{ pg_database }}
+      PGBOUNCER_DATABASE: {{ pg_database }}
+      POSTGRESQL_PASSWORD: {{ pg_password }}
+      POSTGRESQL_HOST: {{ pg_hostname | default('postgres') }}
+      POSTGRESQL_PORT: {{ pg_port }}
+      PGBOUNCER_AUTH_TYPE: trust
+      PGBOUNCER_PORT: {{ pgbouncer_port }}
+      PGBOUNCER_DEFAULT_POOL_SIZE: {{ pgbouncer_max_pool_size }}
+      # This is the default, but we're being explicit here because it's important:
+      # pg_notify will NOT work in transaction mode.
+      PGBOUNCER_POOL_MODE: session
+{% endif %}
 {% if execution_node_count|int > 0 %}
   receptor-hop:
     image: {{ receptor_image }}