Add an argument to set the statement timeout for the API user

openhive-network · Feb 6, 2025 · c432399 · c432399
1 parent 1690a78
commit c432399
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 10 deletions.
diff --git a/docker/docker_entrypoint.sh b/docker/docker_entrypoint.sh
@@ -28,9 +28,9 @@ INSTALL_APP=0
 DO_SCHEMA_UPGRADE=0
 WITH_REPTRACKER=0
 REPTRACKER_SCHEMA=reptracker_app
+STATEMENT_TIMEOUT=""
 reptracker_dir="$SCRIPT_DIR/app/reputation_tracker"
 
-
 while [ $# -gt 0 ]; do
   case "$1" in
     --database-url=*)
@@ -44,6 +44,9 @@ while [ $# -gt 0 ]; do
         export POSTGRES_URL="${1#*=}"
         export POSTGRES_ADMIN_URL="${1#*=}"
         ;;
+    --statement-timeout=*)
+        STATEMENT_TIMEOUT="${1#*=}"
+        ;;
     --add-mocks=*)
         ADD_MOCKS="${1#*=}"
         ;;
@@ -65,7 +68,8 @@ while [ $# -gt 0 ]; do
         ;;
     *)
         arg=$1
-        [[ -n "${arg}" ]] && HIVEMIND_ARGS+=("${arg}")  
+        [[ -n "${arg}" ]] && HIVEMIND_ARGS+=("${arg}")
+        ;;
   esac
   shift
 done
@@ -113,7 +117,12 @@ run_hive() {
 setup() {
   log "setup" "Setting up the database..."
   cd /home/hivemind/app
-  ./setup_postgres.sh --postgres-url="${POSTGRES_ADMIN_URL}"
+  # If STATEMENT_TIMEOUT was provided, pass it to setup_postgres.sh
+  if [[ -n "${STATEMENT_TIMEOUT}" ]]; then
+      ./setup_postgres.sh --postgres-url="${POSTGRES_ADMIN_URL}" --statement-timeout="${STATEMENT_TIMEOUT}"
+  else
+      ./setup_postgres.sh --postgres-url="${POSTGRES_ADMIN_URL}"
+  fi
 
   if [ "${WITH_REPTRACKER}" -eq 1 ]; then
     # if we force to install rep tracker then we setup it as non-forking app
@@ -124,7 +133,7 @@ setup() {
   fi
 
   ./install_app.sh --reptracker-schema-name="${REPTRACKER_SCHEMA}" --postgres-url="${POSTGRES_ADMIN_URL}"
-  
+
   if [[ "$ADD_MOCKS" == "true" ]]; then
     log "setup" "Adding mocks to database..."
     # shellcheck source=/dev/null

diff --git a/haf b/haf
diff --git a/scripts/setup_postgres.sh b/scripts/setup_postgres.sh
@@ -11,7 +11,7 @@ source "$SCRIPTPATH/common.sh"
 
 log_exec_params "$@"
 
-# Script reponsible for setup of specified postgres instance.
+# Script responsible for setup of specified postgres instance.
 #
 # - creates all builtin hivemind roles on pointed PostgreSQL server instance
 
@@ -20,10 +20,11 @@ print_help () {
     echo
     echo "Allows to setup a database already filled by HAF instance, to work with hivemind application."
     echo "OPTIONS:"
-    echo "  --host=VALUE         Allows to specify a PostgreSQL host location (defaults to /var/run/postgresql)"
-    echo "  --port=NUMBER        Allows to specify a PostgreSQL operating port (defaults to 5432)"
-    echo "  --postgres-url=URL   Allows to specify a PostgreSQL URL (in opposite to separate --host and --port options)"
-    echo "  --help               Display this help screen and exit"
+    echo "  --host=VALUE               Allows to specify a PostgreSQL host location (defaults to /var/run/postgresql)"
+    echo "  --port=NUMBER              Allows to specify a PostgreSQL operating port (defaults to 5432)"
+    echo "  --postgres-url=URL         Allows to specify a PostgreSQL URL (in opposite to separate --host and --port options)"
+    echo "  --statement-timeout=VALUE  Set the statement_timeout for the hivemind_user role (e.g., 10s)"
+    echo "  --help                   Display this help screen and exit"
     echo
 }
 
@@ -37,6 +38,7 @@ supplement_builtin_roles() {
 POSTGRES_HOST="/var/run/postgresql"
 POSTGRES_PORT=5432
 POSTGRES_URL=""
+STATEMENT_TIMEOUT=""
 
 while [ $# -gt 0 ]; do
   case "$1" in
@@ -49,6 +51,9 @@ while [ $# -gt 0 ]; do
     --postgres-url=*)
         export POSTGRES_URL="${1#*=}"
         ;;
+    --statement-timeout=*)
+        STATEMENT_TIMEOUT="${1#*=}"
+        ;;
     --help)
         print_help
         exit 0
@@ -81,3 +86,8 @@ fi
 #psql "$POSTGRES_ACCESS" -c "GRANT SET ON PARAMETER log_min_messages TO hivemind;"
 
 supplement_builtin_roles "$POSTGRES_ACCESS"
+
+if [ -n "$STATEMENT_TIMEOUT" ]; then
+    echo "Setting statement timeout for hivemind_user role to ${STATEMENT_TIMEOUT}..."
+    psql "$POSTGRES_ACCESS" -v ON_ERROR_STOP=on -c "ALTER ROLE hivemind_user SET statement_timeout TO '${STATEMENT_TIMEOUT}';"
+fi