Quick and dirty check that Postgres database didn't lose acknowledged data

pg-test.py

#!/usr/bin/env python3

"""
CREATE TABLE t1 (id SERIAL PRIMARY KEY, val INT, date TIMESTAMPTZ DEFAULT NOW());
CREATE INDEX t1_val ON t1 (val);
"""

import logging
import sys
import time

import psycopg2
from psycopg2.errors import DatabaseError

LOG = logging.getLogger(__name__)

conn = psycopg2.connect("")
cur = conn.cursor()


def insert_loop():
    global cur, conn
    seq = 0

    cur.execute('SELECT COALESCE(MAX(val), 0) FROM t1')
    start, = cur.fetchone()

    while True:
        try:
            cur.execute('INSERT INTO t1 (val) VALUES (%s) RETURNING *', (seq + start,))
            row = cur.fetchone()
            conn.commit()
            seq += 1
        except DatabaseError as e:
            print(e)
            try:
                conn.close()
            except Exception:
                pass
            conn = psycopg2.connect("")
            cur = conn.cursor()
            time.sleep(0.5)
        except KeyboardInterrupt:
            print(f'Inserted {seq} acknowledged values')
            return


def verify():
    cur.execute('SELECT id, val, date FROM t1 ORDER BY val')
    last = cur.fetchone()
    vals = 1
    for row in cur:
        vals += 1
        if row[1] != last[1] + 1:
            print(f'! {last[1]} -> {row[1]}  @{last[2]} diff {row[2] - last[2]}')
        last = row

    print(f'Verified {vals} values')


if __name__ == "__main__":
    if sys.argv[1] == "insert":
        insert_loop()
    elif sys.argv[1] == "verify":
        verify()
    else:
        print(f"Usage: {sys.argv[0]} insert|verify")

Useful for testing distributed systems (Citus, Patroni, Cockroach, Yugabyte) under stress - severing connection, removing servers etc.

The insert mode inserts values in a loop as fast as it can. Occasionally, it may insert one value twice, if it doesn't receive a server reply from the commit command.

The verify mode goes over the table and looks for gaps.