"""Синхронный драйвер внутри асинхронного кода: `run_in_executor`.

ЧТО ПРОВЕРЯЕТСЯ. Самый частый способ «перейти на async, не переписывая доступ к
базе»: оставить синхронный драйвер и заворачивать каждый запрос в
`loop.run_in_executor`. Снаружи получается `await`, внутри — тот же поток.

ЧЕТЫРЕ УЧАСТНИКА, ОДНА И ТА ЖЕ РАБОТА (одинаковый пул соединений):

  1. потоки без цикла событий — синхронный psycopg в ThreadPoolExecutor;
  2. тот же psycopg, но через `run_in_executor` из корутины;
  3. асинхронный psycopg;
  4. asyncpg.

Разница между (1) и (2) — цена самой обёртки: работа одна и та же, потоки те же,
добавлен только цикл событий между ними.
"""

import asyncio
import sys
import time
from concurrent.futures import ThreadPoolExecutor

import asyncpg
import psycopg
from psycopg_pool import AsyncConnectionPool, ConnectionPool

sys.path.insert(0, __file__.rsplit("/", 1)[0])
import env  # noqa: E402

QUERIES = 600
SLEEP = 0.005
POOL = 8
REPEATS = 3


def threads_only() -> float:
    with ConnectionPool(env.DSN, min_size=POOL, max_size=POOL) as pool:
        pool.wait()

        def one(_: int) -> None:
            with pool.connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("select pg_sleep(%s)", (SLEEP,))
                    cur.fetchone()

        with ThreadPoolExecutor(max_workers=POOL) as ex:
            started = time.perf_counter()
            list(ex.map(one, range(QUERIES)))
            return time.perf_counter() - started


async def _executor() -> float:
    with ConnectionPool(env.DSN, min_size=POOL, max_size=POOL) as pool:
        pool.wait()

        def blocking(_: int) -> None:
            with pool.connection() as conn:
                with conn.cursor() as cur:
                    cur.execute("select pg_sleep(%s)", (SLEEP,))
                    cur.fetchone()

        loop = asyncio.get_running_loop()
        with ThreadPoolExecutor(max_workers=POOL) as ex:
            started = time.perf_counter()
            await asyncio.gather(
                *(loop.run_in_executor(ex, blocking, i) for i in range(QUERIES))
            )
            return time.perf_counter() - started


async def _async_psycopg() -> float:
    pool = AsyncConnectionPool(env.DSN, min_size=POOL, max_size=POOL, open=False)
    await pool.open(wait=True)
    try:

        async def one() -> None:
            async with pool.connection() as conn:
                async with conn.cursor() as cur:
                    await cur.execute("select pg_sleep(%s)", (SLEEP,))
                    await cur.fetchone()

        started = time.perf_counter()
        await asyncio.gather(*(one() for _ in range(QUERIES)))
        return time.perf_counter() - started
    finally:
        await pool.close()


async def _asyncpg() -> float:
    pool = await asyncpg.create_pool(env.DSN, min_size=POOL, max_size=POOL)
    try:

        async def one() -> None:
            async with pool.acquire() as conn:
                await conn.fetchval("select pg_sleep($1)", SLEEP)

        started = time.perf_counter()
        await asyncio.gather(*(one() for _ in range(QUERIES)))
        return time.perf_counter() - started
    finally:
        await pool.close()


def best(fn) -> float:
    return min(fn() for _ in range(REPEATS))


def main() -> None:
    env.describe()
    print(
        f"{QUERIES} запросов по {SLEEP * 1000:.0f} мс ожидания, пул {POOL} соединений, "
        f"лучшее из {REPEATS}"
    )
    print(f"идеал при полном совмещении ожиданий: {QUERIES * SLEEP / POOL:.2f}с\n")

    rows = [
        ("потоки, без цикла событий", best(threads_only)),
        ("потоки через run_in_executor", best(lambda: asyncio.run(_executor()))),
        ("psycopg async", best(lambda: asyncio.run(_async_psycopg()))),
        ("asyncpg", best(lambda: asyncio.run(_asyncpg()))),
    ]
    base = rows[0][1]
    print(f"{'модель':>30} | {'время':>8} | {'к потокам':>10}")
    print("-" * 54)
    for label, seconds in rows:
        print(f"{label:>30} | {seconds:>7.3f}с | {seconds / base:>9.2f}×")


if __name__ == "__main__":
    main()