feat: add RangerConnectionChecker (pinterest#59)

* feat: add RangerConnectionChecker * chore: clean imports * fix: combine ranger check
lilyli9 · Nov 18, 2022 · 5d63eb7 · 5d63eb7
1 parent 875c906
commit 5d63eb7
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 7 deletions.
diff --git a/plugins/engine_status_checker_plugin/trino_connection_checker.py b/plugins/engine_status_checker_plugin/trino_connection_checker.py
@@ -4,6 +4,7 @@
 import requests
 import os
 from pathlib import Path
+from env import QuerybookSettings
 
 from lib.engine_status_checker.base_checker import BaseEngineStatusChecker, EngineStatus
 from const.query_execution import QueryEngineStatus
@@ -30,17 +31,32 @@ def check_connection(
 
     utc_now_str = datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S") 
 
+    ranger_response_code_list = []
     try:
         with Timeout(20, "Connection took too long"):
-            conn_str = Path(client_settings["connection_string"])
-            url = os.path.join('https://', conn_str.parts[1], "v1/info/state")
-            response = requests.get(url=url, timeout=10)
-            data = response.json()
-            if data == "ACTIVE":
-                result["messages"].append(f"Connection ACTIVE at {utc_now_str} UTC")
+            # Check response from trino API
+            trino_conn_str = Path(client_settings["connection_string"])
+            trino_url = os.path.join('https://', trino_conn_str.parts[1], "v1/info/state")
+            trino_response = requests.get(url=trino_url, timeout=10)
+            trino_response_data = trino_response.json()
+
+            if trino_response_data == "ACTIVE":
+                result["messages"].append(f"Trino connection ACTIVE at {utc_now_str} UTC")
+            else:
+                result["status"] = QueryEngineStatus.ERROR.value
+                result["messages"].append(f"Trino connection returned {trino_response_data} at {utc_now_str} UTC")
+
+            # Check ranger instance for response
+            ranger_response = requests.get(url=QuerybookSettings.RANGER_URL, timeout=10)
+            ranger_response_code = ranger_response.status_code
+            ranger_response_code_list.append(ranger_response_code)
+
+            if all(x == 200 for x in ranger_response_code_list):
+                result["messages"].append(f"Ranger connection ACTIVE at {utc_now_str} UTC")
             else:
                 result["status"] = QueryEngineStatus.ERROR.value
-                result["messages"].append(f"Connection returned {data} at {utc_now_str} UTC")
+                result["messages"].append(f"Ranger connection INACTIVE at {utc_now_str} UTC")
+
     except Exception as e:
         result["status"] = QueryEngineStatus.ERROR.value
         result["messages"].append(f"Connection INACTIVE at {utc_now_str} UTC: {str(e)}")

diff --git a/querybook/server/env.py b/querybook/server/env.py
@@ -136,3 +136,6 @@ class QuerybookSettings(object):
 
     # Event Logging
     EVENT_LOGGER_NAME = get_env_config("EVENT_LOGGER_NAME") or "null"
+
+    # Ranger
+    RANGER_URL = get_env_config("RANGER_URL")