# /// script
# requires-python = ">=3.10"
# dependencies = ["psycopg[binary]"]
# ///
# Copyright (c) 2026 Aayush Sharma · SPDX-License-Identifier: MIT
# https://thedataengineering.com/lab-code/LICENSE.txt
"""Read your writes with a replica.

Update a row on a PostgreSQL primary and read it straight back from a
streaming replica that applies changes with a delay. Then carry a token
(the write's log position) and make the replica catch up to it before
reading; then lose the token on the way, as a request routed elsewhere would.
Needs the primary and replica from the compose file (ports 55470, 55471).
"""
import argparse, json, random, time
import psycopg

p = argparse.ArgumentParser(description=__doc__.splitlines()[0])
p.add_argument("--primary", default="postgresql://postgres:lab@localhost:55470/postgres", help="Primary connection string.")
p.add_argument("--replica", default="postgresql://postgres:lab@localhost:55471/postgres", help="Replica connection string.")
p.add_argument("--delay-ms", type=int, default=200, help="recovery_min_apply_delay on the replica.")
p.add_argument("--reads", type=int, default=1000, help="Write-then-read pairs without the token.")
p.add_argument("--token-reads", type=int, default=200, help="Write-then-read pairs with the token (each may wait for the replica).")
p.add_argument("--timeout", type=float, default=2.0, help="Seconds to wait for the replica before reading the primary instead.")
p.add_argument("--seed", type=int, default=0, help="Random seed.")
p.add_argument("--json", help="Also write the results to this JSON file.")
args = p.parse_args()
random.seed(args.seed)

pri = psycopg.connect(args.primary, autocommit=True)
rep = psycopg.connect(args.replica, autocommit=True)
rep.execute(f"ALTER SYSTEM SET recovery_min_apply_delay = '{args.delay_ms}ms'")
rep.execute("SELECT pg_reload_conf()")
pri.execute("DROP TABLE IF EXISTS profile")
pri.execute("CREATE TABLE profile (id int PRIMARY KEY, name text)")
pri.execute("INSERT INTO profile VALUES (1, 'v0')")
time.sleep(args.delay_ms / 1000 + 1)

def write(i):
    pri.execute("UPDATE profile SET name = %s WHERE id = 1", (f"v{i}",))
    return pri.execute("SELECT pg_current_wal_lsn()").fetchone()[0]          # the token

def read_replica():
    return rep.execute("SELECT name FROM profile WHERE id = 1").fetchone()[0]

def wait_for(token):
    t = time.perf_counter()
    while time.perf_counter() - t < args.timeout:
        if rep.execute("SELECT pg_last_wal_replay_lsn() >= %s::pg_lsn", (token,)).fetchone()[0]:
            return time.perf_counter() - t, "replica"
        time.sleep(0.005)
    return time.perf_counter() - t, "primary"

# 1. No token: write on the primary, read straight back from the replica
stale = 0
for i in range(1, args.reads + 1):
    write(i)
    stale += read_replica() != f"v{i}"

# 2. The token: wait until the replica has replayed the write
stale_tok, waits, fallbacks, n = 0, [], 0, args.reads
for i in range(n + 1, n + args.token_reads + 1):
    token = write(i)
    waited, src = wait_for(token)
    waits.append(waited); fallbacks += src == "primary"
    value = read_replica() if src == "replica" else pri.execute("SELECT name FROM profile WHERE id = 1").fetchone()[0]
    stale_tok += value != f"v{i}"

# 3. The token lost: half the reads go through a path that drops it (another endpoint, a new session)
stale_lost, n = 0, n + args.token_reads
for i in range(n + 1, n + args.token_reads + 1):
    token = write(i)
    if random.random() < 0.5: wait_for(token)       # this request still carries the token
    stale_lost += read_replica() != f"v{i}"

print(f"replica applies with a {args.delay_ms} ms delay")
print(f"no token:          {stale:>5} of {args.reads} reads stale ({stale / args.reads:.1%})")
print(f"with the token:    {stale_tok:>5} of {args.token_reads} stale; mean wait {sum(waits) / len(waits) * 1000:.0f} ms, "
      f"{fallbacks} fell back to the primary")
print(f"token lost on half the requests: {stale_lost:>3} of {args.token_reads} stale ({stale_lost / args.token_reads:.1%})")

if args.json:
    with open(args.json, "w") as f:
        json.dump({"tables": {"variants": [
            {"variant": "no token", "reads": args.reads, "stale": stale, "stale_pct": stale / args.reads * 100},
            {"variant": "token", "reads": args.token_reads, "stale": stale_tok, "stale_pct": stale_tok / args.token_reads * 100},
            {"variant": "token lost on half the requests", "reads": args.token_reads, "stale": stale_lost,
             "stale_pct": stale_lost / args.token_reads * 100}],
            "summary": [{"delay_ms": args.delay_ms, "stale_pct": stale / args.reads * 100, "token_stale": stale_tok,
                         "token_wait_ms": sum(waits) / len(waits) * 1000, "lost_pct": stale_lost / args.token_reads * 100}]}}, f)