"""Операции, которые задевают больше одного шарда: уникальность и фиксация.

ЗАЧЕМ ЭТОТ СКРИПТ. Пока запрос идёт в один шард, шардированная таблица ведёт
себя как обычная. Два обещания обычной таблицы при этом молча перестают
действовать, стоит операции задеть второй шард: уникальность по столбцу, не
входящему в ключ шардирования, и атомарность транзакции. Здесь проверяется,
что именно PostgreSQL говорит в каждом из двух случаев.

ЧТО ИМЕННО ПРОВЕРЯЕТСЯ.

  1. Уникальный email на таблице, шардированной по user_id: на таблице с
     внешними секциями (шарды) и на обычной секционированной. Печатается
     ответ сервера дословно.
  2. Транзакция, вставившая по строке в два шарда, когда один из шардов
     отказывает при фиксации. Отказ устроен отложенным триггером, который
     срабатывает ровно на COMMIT, — то есть ПОСЛЕ того, как обе вставки
     прошли. Проверяется, что осталось в каждом шарде, при отказе каждого
     из двух и с parallel_commit и без него.
  3. Двухфазная фиксация на шарде: разрешена ли она по умолчанию.

ПОЧЕМУ ОТКАЗ ИМЕННО НА ФИКСАЦИИ. Ошибка во время вставки откатывает всё —
это неинтересно и честно. Опасный случай — когда всё прошло, и отказал один
участник на последнем шаге: postgres_fdw фиксирует удалённые транзакции по
очереди (или одновременно, с parallel_commit), и к моменту отказа второго
первый уже мог зафиксироваться.

ЗАПУСК:
    DE_BENCH_DSN="postgres://postgres@localhost:5433/postgres" \\
      python3 bench/sharding/crossshard.py
Вывод: runs/crossshard.txt
"""

from __future__ import annotations

import psycopg

from stand import Stand, version


def show(title: str) -> None:
    print()
    print(title)
    print("-" * len(title))


def attempt(conn, sql: str) -> list[str]:
    """Выполнить и вернуть ответ сервера строками: OK или текст ошибки."""
    try:
        conn.execute(sql)
        return ["OK"]
    except psycopg.Error as e:
        lines = [f"ERROR:  {e.diag.message_primary}"]
        if e.diag.message_detail:
            lines.append(f"DETAIL:  {e.diag.message_detail}")
        if e.diag.message_hint:
            lines.append(f"HINT:  {e.diag.message_hint}")
        return lines


def main() -> None:
    print(f"{version()} | operations that touch more than one shard")
    print("  4 shards behind postgres_fdw, orders partitioned by hash of user_id")

    stand = Stand(4, 0.0, users=200, prefix="shx_")
    stand.build()

    show("1. A UNIQUE EMAIL ON A TABLE SHARDED BY THE USER ID")
    with stand.connect() as conn:
        print("  the sharded table: partitions are foreign tables on the shards")
        for sql in (
            "alter table orders add unique (email)",
            "alter table orders add unique (user_id, email)",
        ):
            print(f"    {sql}")
            for line in attempt(conn, sql):
                print(f"      {line}")
        conn.execute(
            "create table accounts (user_id bigint, email text) partition by hash (user_id)"
        )
        for i in range(4):
            conn.execute(
                f"create table accounts_{i} partition of accounts"
                f" for values with (modulus 4, remainder {i})"
            )
        print("  the same table with ordinary local partitions")
        for sql in (
            "alter table accounts add unique (email)",
            "alter table accounts add unique (user_id, email)",
        ):
            print(f"    {sql}")
            for line in attempt(conn, sql):
                print(f"      {line}")
    print()
    print("  Across shards there is no unique constraint at all. Even locally, a")
    print("  partitioned table can only enforce uniqueness that includes the")
    print("  partition key: each partition checks its own rows, and a duplicate")
    print("  email under another user_id lives in another partition.")

    # Пользователь, чьи строки лежат в шарде 0, и пользователь из шарда 1.
    owners = {}
    for i in (0, 1):
        with stand.shard_conn(i) as sc:
            owners[i] = sc.execute("select min(user_id) from orders").fetchone()[0]

    show("2. ONE TRANSACTION, TWO SHARDS, ONE SHARD REFUSES AT COMMIT")
    print(f"  insert one row for user {owners[0]} (shard 0) and one for user {owners[1]} (shard 1),")
    print("  then COMMIT; the refusing shard has a deferred trigger that raises at commit")
    print(f"  {'parallel_commit':<16} {'refusing shard':>15} {'client sees':>13} {'row on shard 0':>15} {'row on shard 1':>15}")
    outcomes = []
    for pc in ("false", "true"):
        stand.set_option("parallel_commit", pc)
        for refusing in (0, 1):
            for i in range(4):
                with stand.shard_conn(i) as sc:
                    sc.execute("drop trigger if exists refuse on orders")
                    sc.execute("drop function if exists refuse()")
                    sc.execute("delete from orders where order_id < 0")
            with stand.shard_conn(refusing) as sc:
                sc.execute(
                    "create function refuse() returns trigger language plpgsql as"
                    " $$ begin raise exception 'shard refuses at commit'; end $$"
                )
                sc.execute(
                    "create constraint trigger refuse after insert on orders"
                    " deferrable initially deferred for each row execute function refuse()"
                )
            seen = "committed"
            with stand.connect() as conn:
                try:
                    with conn.transaction():
                        conn.execute("insert into orders values (%s, -1, 'a@example.com', 1)", (owners[0],))
                        conn.execute("insert into orders values (%s, -2, 'b@example.com', 1)", (owners[1],))
                except psycopg.Error:
                    seen = "an error"
            present = []
            for i in (0, 1):
                with stand.shard_conn(i) as sc:
                    n = sc.execute("select count(*) from orders where order_id < 0").fetchone()[0]
                    present.append("kept" if n else "none")
            outcomes.append((pc, refusing, present))
            print(f"  {pc:<16} {refusing:>15} {seen:>13} {present[0]:>15} {present[1]:>15}")
    partial = sum(1 for _, _, p in outcomes if "kept" in p)
    print()
    print(f"  the client got an error every time; a row survived anyway in {partial} of {len(outcomes)} cases")
    print()
    print("  postgres_fdw commits each remote transaction in turn at the local commit.")
    print("  When the first shard to commit is the one that refuses, the rest are")
    print("  rolled back and nothing survives. When it is the other one, that one has")
    print("  already committed and stays committed. With parallel_commit every COMMIT")
    print("  goes out at once, so the shard that did not refuse has committed whichever")
    print("  shard refuses.")

    show("3. TWO-PHASE COMMIT ON A SHARD, DEFAULT SETTINGS")
    with stand.shard_conn(0) as sc:
        setting = sc.execute("show max_prepared_transactions").fetchone()[0]
        print(f"  max_prepared_transactions = {setting}")
        sc.execute("begin")
        sc.execute("insert into orders values (1, -9, 'c@example.com', 1)")
        print("    prepare transaction 'order-42'")
        for line in attempt(sc, "prepare transaction 'order-42'"):
            print(f"      {line}")
    print()
    print("  Two-phase commit is what would make block 2 atomic: every shard first")
    print("  promises to commit, and only then is anyone told to. PostgreSQL has it,")
    print("  turned off by default, and postgres_fdw does not use it.")

    stand.drop()


if __name__ == "__main__":
    main()