From 4b4ce9e60ae22fadfc79d3e11df1bc3e682ca5da Mon Sep 17 00:00:00 2001
From: Alexey Masterov <alexeymasterov@neon.tech>
Date: Tue, 1 Apr 2025 15:12:27 +0200
Subject: [PATCH] Add support for BENCHMARTK_CONNSTR

---
 docker-compose/run-tests.sh | 29 +++++++++++++++++++++++------
 1 file changed, 23 insertions(+), 6 deletions(-)
 mode change 100644 => 100755 docker-compose/run-tests.sh

diff --git a/docker-compose/run-tests.sh b/docker-compose/run-tests.sh
old mode 100644
new mode 100755
index cb2de9770b..919fb66003
--- a/docker-compose/run-tests.sh
+++ b/docker-compose/run-tests.sh
@@ -1,12 +1,29 @@
 #!/bin/bash
 set -x
 
-if [[ -v BENCHMARK_CONNSTR ]] && [[ ${BENCHMARK_CONNSTR} =~ ^postgres(ql)?:\/\/([^:]+):([^@]+)@([^:/]+):?([0-9]*)\/(.+)$ ]]; then
-  export PGUSER="${BASH_REMATCH[2]}"
-  export PGPASSWORD="${BASH_REMATCH[3]}"
-  export PGHOST="${BASH_REMATCH[4]}"
-  export PGPORT="${BASH_REMATCH[5]:-5432}"
-  export PGDATABASE="${BASH_REMATCH[6]}"
+if [[ -v BENCHMARK_CONNSTR ]]; then
+  uri_no_proto="${BENCHMARK_CONNSTR#postgres://}"
+  uri_no_proto="${uri_no_proto#postgresql://}"
+  if [[ $uri_no_proto == *\?* ]]; then
+    base="${uri_no_proto%%\?*}"       # before '?'
+  else
+    base="$uri_no_proto"
+  fi
+  if [[ $base =~ ^([^:]+):([^@]+)@([^:/]+):?([0-9]*)/(.+)$ ]]; then
+    export PGUSER="${BASH_REMATCH[1]}"
+    export PGPASSWORD="${BASH_REMATCH[2]}"
+    export PGHOST="${BASH_REMATCH[3]}"
+    export PGPORT="${BASH_REMATCH[4]:-5432}"
+    export PGDATABASE="${BASH_REMATCH[5]}"
+    echo export PGUSER="${BASH_REMATCH[1]}"
+    echo export PGPASSWORD="${BASH_REMATCH[2]}"
+    echo export PGHOST="${BASH_REMATCH[3]}"
+    echo export PGPORT="${BASH_REMATCH[4]:-5432}"
+    echo export PGDATABASE="${BASH_REMATCH[5]}"
+  else
+    echo "Invalid PostgreSQL base URI"
+    exit 1
+  fi
 fi
 
 extdir=${1}