#!/usr/bin/env python3
# Reproduce stack overflow (SIGSEGV) in jsonb_plperl and jsonb_plpython3u.

import subprocess
import time

DEPTH = 100_000
PGLOG = "/home/eax/pginstall/data/logfile"


def psql(sql):
    subprocess.run(["psql", "-c", sql], check=True, capture_output=True)


def run_test(label, query):
    with open(PGLOG) as f:
        f.seek(0, 2)
        log_pos = f.tell()

    print(f"{label} (depth={DEPTH}): ", end="", flush=True)
    result = subprocess.run(["psql", "-c", query], capture_output=True)
    time.sleep(2)

    with open(PGLOG) as f:
        f.seek(log_pos)
        log = f.read()

    if "signal 11" in log:
        print("SIGSEGV")
        for line in log.splitlines():
            if "signal 11" in line or "DETAIL" in line:
                print(" ", line.strip())
    else:
        output = result.stdout.decode().strip()
        print(f"no crash:\n{output}")
        if result.stderr:
            print(result.stderr.decode().strip())


psql("CREATE EXTENSION IF NOT EXISTS jsonb_plpython3u CASCADE;")
psql("CREATE EXTENSION IF NOT EXISTS jsonb_plperl CASCADE;")
psql("""
    CREATE OR REPLACE FUNCTION py_deep(n int) RETURNS jsonb
    LANGUAGE plpython3u TRANSFORM FOR TYPE jsonb AS $$
    d = 1
    for i in range(n): d = {'x': d}
    return d
    $$;
""")
psql("""
    CREATE OR REPLACE FUNCTION perl_deep(n int) RETURNS jsonb
    LANGUAGE plperl TRANSFORM FOR TYPE jsonb AS $$
    my $h = 1;
    for my $i (1..$_[0]) { $h = {x => $h}; }
    return $h;
    $$;
""")

run_test("plpython3u", f"SELECT py_deep({DEPTH});")

# Wait for postmaster to restart the backend after the crash.
for _ in range(5):
    result = subprocess.run(["psql", "-c", "SELECT 1;"], capture_output=True)
    if result.returncode == 0:
        break
    time.sleep(2)

run_test("plperl    ", f"SELECT perl_deep({DEPTH});")