From b69fd32605bf088b47bc72013742147396c500e0 Mon Sep 17 00:00:00 2001
From: Gopal Patel <gopal@qualified.com>
Date: Fri, 26 Jan 2024 15:44:49 -0800
Subject: [PATCH] Support connect_query via PGBOUNCER_CONNECT_QUERY

Squashed version of #1 and
https://github.com/heroku/heroku-buildpack-pgbouncer/pull/95 to make
patches easier to carry forward.
---
 README.md                 | 1 +
 bin/gen-pgbouncer-conf.sh | 2 +-
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 15bfe6d..5e8f5b0 100644
--- a/README.md
+++ b/README.md
@@ -146,6 +146,7 @@ Some settings are configurable through app config vars at runtime. Refer to the
 - `PGBOUNCER_SERVER_RESET_QUERY` Default is empty when pool mode is transaction, and "DISCARD ALL;" when session.
 - `PGBOUNCER_IGNORE_STARTUP_PARAMETERS` Adds parameters to ignore when pgbouncer is starting. Some postgres libraries, like Go's pq, append this parameter, making it impossible to use this buildpack. Default is empty and the most common ignored parameter is `extra_float_digits`. Multiple parameters can be seperated via commas. Example: `PGBOUNCER_IGNORE_STARTUP_PARAMETERS="extra_float_digits, some_other_param"`
 - `PGBOUNCER_QUERY_WAIT_TIMEOUT` Default is 120 seconds, helps when the server is down or the database rejects connections for any reason. If this is disabled, clients will be queued infinitely.
+- `PGBOUNCER_CONNECT_QUERY` Default is empty. The query to be executed on newly created connections. This is useful for setting things like `statement_timeout`. e.g. `PGBOUNCER_CONNECT_QUERY="set statement_timeout to 10000"`
 
 For more info, see [CONTRIBUTING.md](CONTRIBUTING.md)
 
diff --git a/bin/gen-pgbouncer-conf.sh b/bin/gen-pgbouncer-conf.sh
index 0dad387..2c26bb7 100644
--- a/bin/gen-pgbouncer-conf.sh
+++ b/bin/gen-pgbouncer-conf.sh
@@ -90,7 +90,7 @@ do
 EOFEOF
 
   cat >> "$CONFIG_DIR/pgbouncer.ini" << EOFEOF
-$CLIENT_DB_NAME= host=$DB_HOST dbname=$DB_NAME port=$DB_PORT
+$CLIENT_DB_NAME= host=$DB_HOST dbname=$DB_NAME port=$DB_PORT ${PGBOUNCER_CONNECT_QUERY:+connect_query=\'${PGBOUNCER_CONNECT_QUERY//\'/\'\'}\'}
 EOFEOF
 
   (( n += 1 ))