From 7dda7aef23a8caae3bb917dabba3894fcf639af3 Mon Sep 17 00:00:00 2001 From: Charles Santos Date: Tue, 19 May 2026 14:42:37 -0300 Subject: [PATCH] feat: Add standalone script for measuring datasource connection times --- probe_connections.py | 179 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 179 insertions(+) create mode 100644 probe_connections.py diff --git a/probe_connections.py b/probe_connections.py new file mode 100644 index 0000000..b7ebfec --- /dev/null +++ b/probe_connections.py @@ -0,0 +1,179 @@ +""" +Standalone script to measure connection time for each datasource: +impala, redis, hbase, and rabbit. + +Reads the same environment variables used by the app. +Run with: python probe_connections.py +""" + +import os +import subprocess +import time +import ssl +from datetime import datetime + + +def check_klist(): + """Parse and display the current Kerberos TGT status from the ticket cache.""" + cache = os.environ.get("KRB5CCNAME", "") + cmd = ["klist"] + (["-c", cache] if cache else []) + try: + out = subprocess.run(cmd, capture_output=True, text=True, timeout=5) + lines = out.stdout.strip().splitlines() + except Exception as exc: + print(f"[klist] ERROR: {exc}\n") + return + + principal, valid_from, expires, renew_until = "-", "-", "-", "-" + for line in lines: + if line.startswith("Default principal:"): + principal = line.split(":", 1)[1].strip() + # ticket lines: "MM/DD/YY HH:MM:SS MM/DD/YY HH:MM:SS service" + if "krbtgt/" in line: + parts = line.split() + if len(parts) >= 5: + valid_from = f"{parts[0]} {parts[1]}" + expires = f"{parts[2]} {parts[3]}" + if "renew until" in line: + parts = line.split() + if len(parts) >= 3: + renew_until = f"{parts[2]} {parts[3]}" + + now = datetime.now() + try: + exp_dt = datetime.strptime(expires, "%m/%d/%y %H:%M:%S") + time_left = exp_dt - now + hours, rem = divmod(int(time_left.total_seconds()), 3600) + minutes = rem // 60 + ttl_str = f"{hours}h {minutes}m" if time_left.total_seconds() > 0 else "EXPIRED" + except ValueError: + ttl_str = "unknown" + + print("[Kerberos TGT]") + print(f" Principal : {principal}") + print(f" Valid from : {valid_from}") + print(f" Expires : {expires} (TTL: {ttl_str})") + print(f" Renew until: {renew_until}") + print() + + +def probe_impala(): + from impala.dbapi import connect + import kerberos + + host = os.environ["IMPALA_HOST"] + port = int(os.environ["IMPALA_PORT"]) + + # Pre-fetch the Kerberos service ticket so we can time the KDC round-trip + # separately from the actual TCP/SSL connection. + # After authGSSClientStep the ticket is cached, so connect() reuses it. + impala_service = f"impala/{host}@MPT.INTRA" + t0 = time.perf_counter() + __, krb_context = kerberos.authGSSClientInit(impala_service) + kerberos.authGSSClientStep(krb_context, "") + t1 = time.perf_counter() + krb_elapsed = t1 - t0 + + conn = connect( + host=host, + database="spai", + port=port, + kerberos_service_name="impala", + auth_mechanism="GSSAPI", + use_ssl=True, + ) + conn_elapsed = time.perf_counter() - t1 + conn.close() + return {"impala (krb)": krb_elapsed, "impala (conn)": conn_elapsed - krb_elapsed} + + +def probe_redis(): + import redis + + host = os.environ["REDIS_HOST"] + port = int(os.environ["REDIS_PORT"]) + db = int(os.environ["REDIS_DB"]) + + start = time.perf_counter() + r = redis.Redis(host=host, port=port, db=db) + r.ping() + elapsed = time.perf_counter() - start + return elapsed + + +def probe_hbase(): + import kerberos + from thrift.transport import THttpClient + from thrift.protocol import TBinaryProtocol + from hbase import Hbase + + host = os.environ["HBASE_HOST"] + port = os.environ["HBASE_PORT"] + + hbase_service = f"HTTP/{host}@MPT.INTRA" + ssl._create_default_https_context = ssl._create_unverified_context + + t0 = time.perf_counter() + __, krb_context = kerberos.authGSSClientInit(hbase_service) + kerberos.authGSSClientStep(krb_context, "") + negotiate_details = kerberos.authGSSClientResponse(krb_context) + t1 = time.perf_counter() + krb_elapsed = t1 - t0 + + headers = { + "Authorization": "Negotiate " + negotiate_details, + "Content-Type": "application/binary", + } + http_client = THttpClient.THttpClient(f"https://{host}:{port}/") + http_client.setCustomHeaders(headers=headers) + protocol = TBinaryProtocol.TBinaryProtocol(http_client) + http_client.open() + Hbase.Client(protocol) + http_client.close() + conn_elapsed = time.perf_counter() - t1 + + return {"hbase (krb)": krb_elapsed, "hbase (conn)": conn_elapsed} + + +def probe_rabbit(): + import pika + + host = os.environ["RABBIT_HOST"] + port = int(os.environ["RABBIT_PORT"]) + user = os.environ["RABBIT_USER"] + password = os.environ["RABBIT_PASSWORD"] + env = os.environ.get("RABBIT_ENV", "stg") + vhost = f"/suetonio/{env}" + + credentials = pika.PlainCredentials(user, password) + parameters = pika.ConnectionParameters(host, port, vhost, credentials) + + start = time.perf_counter() + conn = pika.BlockingConnection(parameters) + conn.close() + elapsed = time.perf_counter() - start + return elapsed + + +PROBES = { + "impala": probe_impala, + "redis": probe_redis, + "hbase": probe_hbase, + "rabbit": probe_rabbit, +} + +if __name__ == "__main__": + check_klist() + print(f"{'Service':<14} {'Status':<8} {'Time (s)':>10} {'Error'}") + print("-" * 70) + + for name, fn in PROBES.items(): + try: + result = fn() + if isinstance(result, dict): + for sub_name, elapsed in result.items(): + print(f"{sub_name:<14} {'OK':<8} {elapsed:>10.3f}") + else: + print(f"{name:<14} {'OK':<8} {result:>10.3f}") + except Exception as exc: + print(f"{name:<14} {'FAILED':<8} {'N/A':>10} {exc}")