#!/usr/bin/env python3
"""
Galera #613 check: does wsrep_sync_wait still enforce causal reads after COM_CHANGE_USER?

Writer on node2 inserts a row, reader on node1 immediately reads it back.
The global status counter wsrep_causal_reads on node1 is sampled from a separate
monitor connection (with wsrep_sync_wait=0 so it does not count itself).

If wsrep_sync_wait=1 is honoured, the counter grows by one per reader SELECT.
If it stays flat after COM_CHANGE_USER, the causal wait is skipped (bug reproduced).

Requirements, either driver:
  pip install mysql-connector-python   (prebuilt wheel, no compiler needed)
  pip install mariadb                  (needs gcc, python3-devel, MariaDB-devel)
Run it on an otherwise idle lab cluster, so no other traffic moves the counter.
"""
import argparse
import statistics
import sys
import time

try:
    import mariadb as driver
    DRIVER = "mariadb"
except ImportError:
    import mysql.connector as driver
    DRIVER = "mysql-connector-python"


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--node1", required=True, help="reader node (counter is sampled here)")
    p.add_argument("--node2", required=True, help="writer node")
    p.add_argument("--port", type=int, default=3306)
    p.add_argument("--user", default="root")
    p.add_argument("--password", default="")
    p.add_argument("--db", default="test")
    p.add_argument("-n", type=int, default=200, help="reads per phase")
    p.add_argument("--keep", action="store_true", help="do not drop the test table")
    a = p.parse_args()

    def connect(host):
        return driver.connect(host=host, port=a.port, user=a.user, password=a.password,
                              database=a.db, autocommit=True)

    def do_change_user(conn):
        if DRIVER == "mariadb":
            conn.change_user(a.user, a.password, a.db)
        else:
            conn.cmd_change_user(username=a.user, password=a.password, database=a.db)

    def one(cur, sql, params=()):
        cur.execute(sql, params)
        return cur.fetchone()

    mon = connect(a.node1)
    mc = mon.cursor()
    mc.execute("SET SESSION wsrep_sync_wait=0")

    def causal():
        return int(one(mc, "SHOW GLOBAL STATUS LIKE 'wsrep_causal_reads'")[1])

    print("driver                 :", DRIVER)
    print("node1 version          :", one(mc, "SELECT @@version")[0])
    print("node1 global sync_wait :", one(mc, "SELECT @@global.wsrep_sync_wait")[0])
    print("node1 provider version :", one(mc, "SHOW GLOBAL STATUS LIKE 'wsrep_provider_version'")[1])
    print("node1 cluster size     :", one(mc, "SHOW GLOBAL STATUS LIKE 'wsrep_cluster_size'")[1])
    print()

    w = connect(a.node2)
    wc = w.cursor()
    wc.execute("CREATE TABLE IF NOT EXISTS g613 (id BIGINT AUTO_INCREMENT PRIMARY KEY, v INT) ENGINE=InnoDB")
    time.sleep(2)

    def change_user(conn, cur):
        do_change_user(conn)

    def change_user_proxysql(conn, cur):
        # What ProxySQL does on handout: COM_CHANGE_USER, then restores session vars
        do_change_user(conn)
        cur.execute("SET NAMES utf8mb4 COLLATE 'utf8mb4_unicode_ci'")

    def change_user_set_sync(conn, cur):
        do_change_user(conn)
        cur.execute("SET SESSION wsrep_sync_wait=1")

    def sync_off(conn, cur):
        cur.execute("SET SESSION wsrep_sync_wait=0")

    def sync_on(conn, cur):
        cur.execute("SET SESSION wsrep_sync_wait=1")

    def run(name, prep=None, per_read=None):
        r = connect(a.node1)
        rc = r.cursor()
        if prep:
            prep(r, rc)
        sw = one(rc, "SELECT @@session.wsrep_sync_wait")[0]
        c0 = causal()
        stale = 0
        lat = []
        for i in range(a.n):
            if per_read:
                per_read(r, rc)
            wc.execute("INSERT INTO g613 (v) VALUES (%s)", (i,))
            nid = wc.lastrowid
            t = time.perf_counter()
            cnt = one(rc, "SELECT COUNT(*) FROM g613 WHERE id=%s", (nid,))[0]
            lat.append((time.perf_counter() - t) * 1000)
            if cnt == 0:
                stale += 1
        delta = causal() - c0
        r.close()
        print(f"{name:<62} sync_wait={sw}  causal_reads={delta:>5}/{a.n:<5} "
              f"stale={stale:>4}  avg_read_ms={statistics.mean(lat):.3f}")
        return delta, stale

    gsw = int(one(mc, "SELECT @@global.wsrep_sync_wait")[0])
    if not gsw & 1:
        print("WARNING: global wsrep_sync_wait on node1 is not 1 (the scenario under test uses 1).")
        print("         Phases 1, 3, 4, 5 rely on the global value. Run: SET GLOBAL wsrep_sync_wait=1;")
        print()

    res = {}
    res[0] = run("0 fresh, SET SESSION sync_wait=1       (control: causal=N, stale=0)", prep=sync_on)
    res[1] = run("1 fresh connection, global value       (expect causal=N)")
    res[2] = run("2 fresh, SET SESSION sync_wait=0       (control: causal=0)", prep=sync_off)
    res[3] = run("3 after one COM_CHANGE_USER            (bug: causal=0)", prep=change_user)
    res[4] = run("4 COM_CHANGE_USER + SET NAMES          (ProxySQL handout)", prep=change_user_proxysql)
    res[5] = run("5 COM_CHANGE_USER before every read    (ProxySQL per request)", per_read=change_user_proxysql)
    res[6] = run("6 COM_CHANGE_USER + SET sync_wait=1    (app-side workaround?)", prep=change_user_set_sync)

    if not a.keep:
        wc.execute("DROP TABLE g613")

    print()
    n = a.n
    on_ok = res[0][0] >= n * 0.9 and res[0][1] == 0
    off_ok = res[2][0] <= n * 0.1
    if not (on_ok and off_ok):
        print("INCONCLUSIVE: control phases 0/2 did not behave as expected "
              "(other traffic on node1, or wsrep_causal_reads not counting). Check the setup.")
        return 2
    broken = [k for k in (3, 4, 5, 6) if res[k][0] < n * 0.1 or res[k][1] > 0]
    if not gsw & 1:
        broken = [k for k in broken if k == 6]
    if broken:
        print("REPRODUCED: wsrep_sync_wait not enforced after COM_CHANGE_USER in phase(s) "
              f"{', '.join(map(str, broken))} (galera #613).")
        return 1
    print("NOT REPRODUCED: causal reads are still performed after COM_CHANGE_USER.")
    return 0

if __name__ == "__main__":
    sys.exit(main())