From b63c24acbd19f035e34eae4a4b52cf1607447a95 Mon Sep 17 00:00:00 2001
From: Martin Grund <martin.grund@databricks.com>
Date: Tue, 13 Feb 2024 04:59:27 -0800
Subject: [PATCH] [MINOR][CONNECT] Allow Spark Connect Server Script to wait

---
 sbin/start-connect-server.sh | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/sbin/start-connect-server.sh b/sbin/start-connect-server.sh
index a347f43db8b16..fecda717eb341 100755
--- a/sbin/start-connect-server.sh
+++ b/sbin/start-connect-server.sh
@@ -38,4 +38,10 @@ fi
 
 . "${SPARK_HOME}/bin/load-spark-env.sh"
 
-exec "${SPARK_HOME}"/sbin/spark-daemon.sh submit $CLASS 1 --name "Spark Connect server" "$@"
+if [ "$1" == "--wait" ]; then
+  shift
+  exec "${SPARK_HOME}"/bin/spark-submit --class $CLASS 1 --name "Spark Connect Server" "$@"
+else
+  exec "${SPARK_HOME}"/sbin/spark-daemon.sh submit $CLASS 1 --name "Spark Connect server" "$@"
+fi
+