From ecda060dff7ed372525716b39c98454bc8f7a5c5 Mon Sep 17 00:00:00 2001
From: Denys Tsomenko <denys.tsomenko@caspiandb.com>
Date: Wed, 19 Feb 2025 22:34:39 +0200
Subject: [PATCH] Update query runner with session id

---
 utils/query_runner.sh | 64 +++++++++++++++++++++++++++----------------
 1 file changed, 40 insertions(+), 24 deletions(-)

diff --git a/utils/query_runner.sh b/utils/query_runner.sh
index 7ef8c94f4..6e2e00f00 100755
--- a/utils/query_runner.sh
+++ b/utils/query_runner.sh
@@ -1,39 +1,55 @@
 #!/bin/bash
 
-if [ "$#" -lt 2 ]; then
-  echo "Usage: $0 <location_or_cluster_id> <sql_query>"
+if [ "$#" -lt 1 ]; then
+  echo "Usage: $0 <sql_query>"
   exit 1
 fi
 
-LOCATION_OR_CLUSTER=$1
-
-shift
-SQL_QUERY="$@"
-
+SQL_QUERY=$1
 SQL_QUERY=$(echo "$SQL_QUERY" | tr '\n' ' ')
 
-# TODO: This should be change. Use env variable for cluster location or provide seperate params.
-if [[ $LOCATION_OR_CLUSTER == *\/* ]]; then
-  LOCATION=$LOCATION_OR_CLUSTER
-else
-  DEFAULT_BASE_URL="0.0.0.0:3000/ui/warehouses"
-  LOCATION="$DEFAULT_BASE_URL/$LOCATION_OR_CLUSTER/databases/datasets/tables/table_upload_tests/query"
-fi
+# Set default location if not provided
+LOCATION="${QUERY_LOCATION:-0.0.0.0:3000/ui/query}"
 
+SESSION_COOKIE=""
+if [ -n "$SESSION_ID" ]; then
+    SESSION_COOKIE="id=${SESSION_ID}"
+fi
 
-JSON_PAYLOAD=$(jq -n --arg query "$SQL_QUERY" '{"query": $query}')
+JSON_PAYLOAD=$(jq -n --arg query "$SQL_QUERY" '{query: $query}')
 
 RESPONSE=$(curl --location "$LOCATION" \
+  --include \
   --header "Content-Type: application/json" \
+  --cookie "$SESSION_COOKIE" \
   --data "$JSON_PAYLOAD")
 
-if echo "$RESPONSE" | jq empty > /dev/null 2>&1; then
-# TODO: Find library or write recursive traversal of nested json instead of checking for result here.
-    if echo "$RESPONSE" | jq -e '.result' > /dev/null 2>&1; then
-        echo "$RESPONSE" | jq -r '.result' | jq
-    else
-        echo "$RESPONSE" | jq
-    fi
-else
+HEADERS=$(echo "$RESPONSE" | sed -n '/^HTTP/{:a;n;/^\r$/!ba;p}')
+
+HTTP_STATUS=$(echo "$HEADERS" | grep -oP '(?<=HTTP/1.1 )\d{3}')
+
+# Check for errors (status codes 4xx or 5xx)
+if [[ "$HTTP_STATUS" =~ ^4|5 ]]; then
+  echo "Error detected. Full response:"
   echo "$RESPONSE"
-fi
\ No newline at end of file
+fi
+
+BODY=$(echo "$RESPONSE" | sed -n '/^\r$/,$p' | sed '1d')
+
+if echo "$BODY" | jq empty > /dev/null 2>&1; then
+  # TODO: Find library or write recursive traversal of nested json instead of checking for result here.
+  if echo "$BODY" | jq -e '.result' > /dev/null 2>&1; then
+    echo "$BODY" | jq -r '.result' | jq
+  else
+    echo "$BODY" | jq
+  fi
+else
+  echo "$BODY"
+fi
+
+SET_COOKIE=$(echo "$RESPONSE" | grep -oP '(?<=set-cookie: id=)[^;]+')
+
+if [ -n "$SET_COOKIE" ]; then
+  echo "Set-Cookie:"
+  echo "$SET_COOKIE"
+fi