From 90702ac55e726837f62af7c3cfaa270af8f66afe Mon Sep 17 00:00:00 2001
From: Kevin Baensch <kb@ujamii.com>
Date: Wed, 18 Jan 2023 16:54:30 +0100
Subject: [PATCH] Add support for optional database names.

---
 script/myssh           | 54 ++++++++++++++++++++++++++++--------------
 script/sqlproxy_cli.sh |  2 +-
 templates/sqlproxy.sh  |  4 ++--
 3 files changed, 39 insertions(+), 21 deletions(-)

diff --git a/script/myssh b/script/myssh
index e56fdb4..d26b1bd 100755
--- a/script/myssh
+++ b/script/myssh
@@ -84,34 +84,50 @@ get_host() {
 # $4 - username (optional)
 # $5 - password (optional)
 set_host_env() {
-    if [ "$1" = 'mysql' ] || [ "$1" = 'postgres' ]
+    # Unquote argument
+    ARG="${1//\"/}"
+    if [ "${ARG}" = 'mysql' ] || [ "${ARG}" = 'postgres' ]
     then
-        TARGET_HOST_TYPE="$1"
+        TARGET_HOST_TYPE="${ARG}"
     else
-        printf 'Invalid Database type: "%s"\n' "$1"
+        printf 'Invalid Database type: "%s"\n' "$T"
         exit 1
     fi
-    if [[ "$2" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]
+
+    ARG="${2//\"/}"
+    if [[ "${ARG}" =~ ^[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}\.[0-9]{1,3}$ ]]
     then
-        TARGET_HOST_IP="$2"
+        TARGET_HOST_IP="${ARG}"
     else
-        printf 'Invalid Host IP "%s" given.\n' "$2"
+        printf 'Invalid Host IP "%s" given.\n' "${ARG}"
         exit 1
     fi
-    if [[ "$3" =~ ^[0-9]+$ ]]
+
+    ARG="${3//\"/}"
+    if [[ "${ARG}" =~ ^[0-9]+$ ]]
     then
-       TARGET_HOST_PORT="$3"
+       TARGET_HOST_PORT="${ARG}"
     else
-        printf 'Invalid Host Port "%s" given.\n' "$3"
+        printf 'Invalid Host Port "%s" given.\n' "${ARG}"
         exit 1
     fi
-    if [ -z "${TARGET_HOST_USERNAME}" ] && [ -n "$4" ]
+
+    ARG="${4//\"/}"
+    if [ -z "${TARGET_HOST_DBNAME}" ] && [ -n "${ARG}" ]
     then
-        TARGET_HOST_USERNAME="$4"
+        TARGET_HOST_DBNAME="${ARG}"
     fi
-    if [ -z "${TARGET_HOST_PASSWORD}" ] && [ -n "$5" ]
+
+    ARG="${5//\"/}"
+    if [ -z "${TARGET_HOST_USERNAME}" ] && [ -n "${ARG}" ]
     then
-        TARGET_HOST_PASSWORD="$5"
+        TARGET_HOST_USERNAME="${ARG}"
+    fi
+
+    ARG="${6//\"/}"
+    if [ -z "${TARGET_HOST_PASSWORD}" ] && [ -n "${ARG}" ]
+    then
+        TARGET_HOST_PASSWORD="${ARG}"
     fi
 
 }
@@ -121,6 +137,7 @@ set_host_env() {
 # $1 - sql type (mysql or psql)
 # $2 - username
 # $3 - password
+# $4 - dbname (optional)
 run_client() {
     if [ "$1" = 'mysql' ]
     then
@@ -129,7 +146,7 @@ run_client() {
             open "mysql://$2:$3@${SQL_PROXY_HOST}:6033" -a "Sequel Ace"
         elif which mysql >/dev/null 2>&1
         then
-            mysql --protocol=TCP -u "$2" -p"$3" -h "${SQL_PROXY_HOST}" -P 6033
+            mysql --protocol=TCP -u "$2" -p"$3" -h "${SQL_PROXY_HOST}" "${4:+-D$4}" -P 6033
         else
             SHOW_CLI_HELP=true
         fi
@@ -137,7 +154,7 @@ run_client() {
     then
         if which psql >/dev/null 2>&1
         then
-            psql "postgresql://$2:$3@${SQL_PROXY_HOST}:6033/postgres"
+            psql "postgresql://$2:$3@${SQL_PROXY_HOST}:6033/${4:-postgres}"
         else
             SHOW_CLI_HELP=true
         fi
@@ -175,8 +192,9 @@ case "${MAIN_OPTION}" in
         while getopts "u:p:" o
         do
             case "$o" in
-                u) TARGET_HOST_USERNAME="$OPTARG" ;;
-                p) TARGET_HOST_PASSWORD="$OPTARG" ;;
+                u) TARGET_HOST_USERNAME="${OPTARG}" ;;
+                p) TARGET_HOST_PASSWORD="${OPTARG}" ;;
+                d) TARGET_HOST_DBNAME="${OPTARG}" ;;
                 c) USE_CLI=true ;;
             esac
         done
@@ -186,7 +204,7 @@ case "${MAIN_OPTION}" in
         port_forward "${TARGET_HOST_IP}" "$TARGET_HOST_PORT"
         if [ -n "${TARGET_HOST_USERNAME}" ] && [ -n "${TARGET_HOST_PASSWORD}" ]
         then
-            run_client "${TARGET_HOST_TYPE}" "${TARGET_HOST_USERNAME}" "${TARGET_HOST_PASSWORD}"
+            run_client "${TARGET_HOST_TYPE}" "${TARGET_HOST_USERNAME}" "${TARGET_HOST_PASSWORD}" "${TARGET_HOST_DBNAME}"
         fi
         ;;
     disconnect)
diff --git a/script/sqlproxy_cli.sh b/script/sqlproxy_cli.sh
index f25b35d..74a8cd2 100755
--- a/script/sqlproxy_cli.sh
+++ b/script/sqlproxy_cli.sh
@@ -9,7 +9,7 @@ get_host() {
     export HOST=$(echo "${SSH_ORIGINAL_COMMAND}" | cut -d ' ' -f2)
     if [ "${HOST}" != 'get' ]
     then
-        jq -r 'first(.[] | select(.host == $ENV.HOST)) | [ .type, .ip, .port, .user, .password ] | join(" ")' < "${DB_DATA_FILE}"
+        jq -r 'first(.[] | select(.host == $ENV.HOST)) | [ .type, .ip, .port, .dbname, .user, .password ] | "\"" + join("\" \"") + "\""' < "${DB_DATA_FILE}"
     fi
 }
 
diff --git a/templates/sqlproxy.sh b/templates/sqlproxy.sh
index c6781ad..5b0db4d 100644
--- a/templates/sqlproxy.sh
+++ b/templates/sqlproxy.sh
@@ -4,9 +4,9 @@ set -e
 WRAP_START='[\n'
 if ${EXCLUDE_USERPASS}
 then
-    TEMPLATE='  { "ip": "${LOCAL_IP}", "type": "${LOCAL_DB_TYPE}", "host": "${LOCAL_DB_HOST}", "port": "${LOCAL_DB_PORT}" }'
+    TEMPLATE='  { "ip": "${LOCAL_IP}", "type": "${LOCAL_DB_TYPE}", "host": "${LOCAL_DB_HOST}", "port": "${LOCAL_DB_PORT}", "dbname": "${LOCAL_DB_DBNAME}" }'
 else
-    TEMPLATE='  { "ip": "${LOCAL_IP}", "type": "${LOCAL_DB_TYPE}", "user": "${LOCAL_DB_USER}", "password": "${LOCAL_DB_PASSWORD}", "host": "${LOCAL_DB_HOST}", "port": "${LOCAL_DB_PORT}" }'
+    TEMPLATE='  { "ip": "${LOCAL_IP}", "type": "${LOCAL_DB_TYPE}", "user": "${LOCAL_DB_USER}", "port": "${LOCAL_DB_PORT}", "dbname": "${LOCAL_DB_DBNAME}", "password": "${LOCAL_DB_PASSWORD}", "host": "${LOCAL_DB_HOST}" }'
 fi
 SEPARATOR=',\n'
 WRAP_END='\n]'