From 6dfb1de5c6af0b4d737b574cda9612c965edb62c Mon Sep 17 00:00:00 2001
From: Revant Nandgaonkar <revant.one@gmail.com>
Date: Wed, 1 Jul 2020 20:01:47 +0530
Subject: [PATCH] feat: ping additional services using doctor

doctor --ping-service postgres:5432 -p mariadb:3306
---
 build/common/commands/check_connection.py | 26 +++++++++++++----
 build/common/commands/doctor.py           | 34 +++++++++++++++++++++--
 build/common/worker/docker-entrypoint.sh  |  2 +-
 3 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/build/common/commands/check_connection.py b/build/common/commands/check_connection.py
index cd87e44b..2b47c753 100644
--- a/build/common/commands/check_connection.py
+++ b/build/common/commands/check_connection.py
@@ -53,18 +53,32 @@ def get_config():
     return config
 
 
-# Check mariadb
-def check_mariadb(retry=10, delay=3, print_attempt=True):
+# Check service
+def check_service(
+    retry=10,
+    delay=3,
+    print_attempt=True,
+    service_name=None,
+    service_port=None):
+
     config = get_config()
+    if not service_name:
+        service_name = config.get(DB_HOST_KEY, 'mariadb')
+    if not service_port:
+        service_port = config.get(DB_PORT_KEY, DB_PORT)
+
     is_db_connected = False
     is_db_connected = check_host(
-        config.get(DB_HOST_KEY, 'mariadb'),
-        config.get(DB_PORT_KEY, DB_PORT),
+        service_name,
+        service_port,
         retry,
         delay,
         print_attempt)
     if not is_db_connected:
-        print("Connection to MariaDB timed out")
+        print("Connection to {service_name}:{service_port} timed out".format(
+            service_name=service_name,
+            service_port=service_port,
+        ))
         exit(1)
 
 
@@ -128,7 +142,7 @@ def get_site_config(site_name):
 
 
 def main():
-    check_mariadb()
+    check_service()
     check_redis_queue()
     check_redis_cache()
     check_redis_socketio()
diff --git a/build/common/commands/doctor.py b/build/common/commands/doctor.py
index c317f291..5354e476 100644
--- a/build/common/commands/doctor.py
+++ b/build/common/commands/doctor.py
@@ -1,20 +1,48 @@
+import argparse
+
 from check_connection import (
-    check_mariadb,
+    check_service,
     check_redis_cache,
     check_redis_queue,
     check_redis_socketio,
 )
 
 
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '-p',
+        '--ping-service',
+        dest='ping_services',
+        action='append',
+        type=str,
+        help='list of services to ping, e.g. doctor -p "postgres:5432" --ping-service "mariadb:3306"',
+    )
+    args = parser.parse_args()
+    return args
+
+
 def main():
-    check_mariadb(retry=1, delay=0, print_attempt=False)
-    print("MariaDB Connected")
+    args = parse_args()
+    check_service(retry=1, delay=0, print_attempt=False)
+    print("Bench database Connected")
     check_redis_cache(retry=1, delay=0, print_attempt=False)
     print("Redis Cache Connected")
     check_redis_queue(retry=1, delay=0, print_attempt=False)
     print("Redis Queue Connected")
     check_redis_socketio(retry=1, delay=0, print_attempt=False)
     print("Redis SocketIO Connected")
+    if(args.ping_services):
+        for service in args.ping_services:
+            service_name, service_port = service.split(':')
+            check_service(
+                retry=1,
+                delay=0,
+                print_attempt=False,
+                service_name=service_name,
+                service_port=service_port,
+            )
+            print("{0}:{1} Connected".format(service_name, service_port))
 
     print("Health check successful")
     exit(0)
diff --git a/build/common/worker/docker-entrypoint.sh b/build/common/worker/docker-entrypoint.sh
index 0a8fc0f6..5b68eec4 100755
--- a/build/common/worker/docker-entrypoint.sh
+++ b/build/common/worker/docker-entrypoint.sh
@@ -154,7 +154,7 @@ elif [ "$1" = 'migrate' ]; then
 elif [ "$1" = 'doctor' ]; then
 
   su frappe -c ". /home/frappe/frappe-bench/env/bin/activate \
-    && python /home/frappe/frappe-bench/commands/doctor.py"
+    && python /home/frappe/frappe-bench/commands/doctor.py ${@:2}"
   exit
 
 elif [ "$1" = 'backup' ]; then