EXPLAIN a sql query in a kubernetes pod and open the result on explain.dalibo.com.

dalibo-explain

#!/usr/bin/env -S uv run --script
"""EXPLAIN a sql query in a kubernetes pod and open the result on explain.dalibo.com."""

# /// script
# requires-python = ">=3.9"
# dependencies = [
#     "requests",
# ]
# ///

import subprocess
import sys
import webbrowser
from urllib.parse import urljoin

import requests

DALIBO_URL = "https://explain.dalibo.com"


def _input(prompt: str, default: str | None = None, multiline: bool = False) -> str:
    full_prompt = prompt
    if default:
        full_prompt = f"{full_prompt} ({default})"
    while True:
        if not multiline:
            value = input(f"{full_prompt}: ")
        else:
            print(
                full_prompt,
                "[multiple lines allowed, press ENTER followed by CTRL-D to confirm]:",
            )
            value = sys.stdin.read()
        value = value.strip()
        if not value and default:
            value = default
        if value:
            break
        print("Please enter a value!")
    return value

def _get_default_kube_context() -> str:
    return subprocess.check_output(
	["kubectl", "config", "current-context"],
	text=True,
    ).strip()

def _get_default_kube_namespace() -> str:
    return subprocess.check_output(
	["kubectl", "config", "view", "-o", "jsonpath={..namespace}", "--minify"],
	text=True,
    ).strip()

def _get_default_postgres_pod() -> str:
    return "postgres-1"

def _get_default_database() -> str:
    return "odoo"

def _make_explain_query(query: str) -> str:
    preamble = (
	"\\set ON_ERROR_STOP true\n"
	"EXPLAIN (ANALYZE, COSTS, VERBOSE, BUFFERS, FORMAT JSON)\n"
    )
    return f"{preamble} {query}"

def _make_remote_command() -> list[str]:
    kube_context = _input("kube context", _get_default_kube_context())
    kube_namespace = _input("kube namespace", _get_default_kube_namespace())
    postgres_pod = _input("postgres pod", _get_default_postgres_pod())
    return [
	"kubectl",
	"exec",
	"--context",
	kube_context,
	"--namespace",
	kube_namespace,
	"--stdin",
	postgres_pod,
	"--",
    ]

def _make_psql_command() -> list[str]:
    database = _input("database", _get_default_database())
    return [
	"psql",
	"-XqAt",
	"-d",
	database,
    ]

# User input
remote_command = _make_remote_command()
psql_command = _make_psql_command()
query = _input("sql query, without EXPLAIN", multiline=True)
explain_query = _make_explain_query(query)

# Run the sql query
proc = subprocess.Popen(
    remote_command + psql_command,
    text=True,
    stdout=subprocess.PIPE,
    stdin=subprocess.PIPE,
)
plan, _ = proc.communicate(input=explain_query)
if proc.returncode != 0:
    sys.exit(1)

# Post the plan for visualization
response = requests.post(
    urljoin(DALIBO_URL, "/new"),
    data={
        "query": explain_query,
        "plan": plan,
    },
    allow_redirects=False,
)
if response.status_code == 302:
    webbrowser.open(urljoin(DALIBO_URL, response.headers["location"]))
else:
    print(response.text)