#!/usr/bin/bash
set -euo pipefail

CURRENT_USER=$(whoami)

LOG="/tmp/pg_connection_reaper_${CURRENT_USER}.log"

if [[ "${CURRENT_USER}" == "postgres" ]]; then
        prefix=""
else
        prefix="sudo -u postgres"
fi

echo "$(date '+%y-%m-%d %H:%M:%S.%N') - BEGIN pg connection reaper" | tee -a $LOG

threshold=$(($(nproc) * 2))

active_connection_count=$($prefix psql -Atc "SELECT count(*) FROM pg_stat_activity WHERE state = 'active'")

WHERE="state = 'active' AND usename NOT IN ('postgres','replica','rewind') AND usename NOT LIKE ('pitr%') AND application_name NOT IN ('pg_receivewal', 'pg_dump')"

if [[ ${active_connection_count} -gt $threshold ]]; then
        CSV=/srv/log/main/pg_connection_reaper_$(date "+%Y%m%d_%H%M%S").csv
        echo "$(date '+%y-%m-%d %H:%M:%S.%N') - current number of active connections is ${active_connection_count}, which is more than the double of threads ($threshold)" | tee -a $LOG
        echo "$(date '+%y-%m-%d %H:%M:%S.%N') - dumping current active queries to $CSV" | tee -a $LOG
        $prefix psql -Atc "COPY(SELECT * from pg_stat_activity WHERE $WHERE) to stdout delimiter ';' csv" | $prefix tee $CSV
        TERMINATE="SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE $WHERE"
        echo "$(date '+%y-%m-%d %H:%M:%S.%N') - $TERMINATE" | tee -a $LOG
        $prefix psql -Atc "$TERMINATE"
else
        echo "$(date '+%y-%m-%d %H:%M:%S.%N') - current number of connection is ${active_connection_count}, which is inferior to $threshold, doing nothing" | tee -a $LOG
fi

echo "$(date '+%y-%m-%d %H:%M:%S.%N') - END pg connection reaper" | tee -a $LOG