AsyncCConnection does not provide PgBouncer compatibility: unused with the default caches, and fails like asyncpg.Connection with caches > 0 #60

import asyncio
import re
from collections import Counter

import asyncpg
from pydantic import SecretStr
from sqlalchemy import text
from sqlalchemy_foundation_kit import create_async_session_manager
from sqlalchemy_foundation_kit.contrib.settings import BasePostgresConfig, ConnectionSettings, QuerySettings
from sqlalchemy_foundation_kit.session.connection import AsyncCConnection
from testcontainers.core.container import DockerContainer
from testcontainers.core.network import Network
from testcontainers.core.wait_strategies import LogMessageWaitStrategy
from testcontainers.postgres import PostgresContainer

names_generated = Counter()


class CountingAsyncCConnection(AsyncCConnection):
    def _get_unique_id(self, prefix):
        names_generated[prefix] += 1
        return super()._get_unique_id(prefix)


class CountingPlainConnection(asyncpg.Connection):
    def _get_unique_id(self, prefix):
        names_generated[prefix] += 1
        return super()._get_unique_id(prefix)


def start_pgbouncer(network: Network, **settings: object) -> DockerContainer:
    container = DockerContainer("edoburu/pgbouncer:v1.25.2-p0").with_network(network).with_exposed_ports(5432)
    env = {"DB_HOST": "db", "DB_USER": "test", "DB_PASSWORD": "test", "DB_NAME": "test",
           "AUTH_TYPE": "scram-sha-256", "POOL_MODE": "transaction",
           **{k.upper(): str(v) for k, v in settings.items()}}
    for key, value in env.items():
        container.with_env(key, value)
    container.waiting_for(LogMessageWaitStrategy("process up").with_startup_timeout(30))
    return container.start()


async def workload(bouncer: DockerContainer, caches: tuple[int, int], connection_class: type) -> str:
    config = BasePostgresConfig(
        connection=ConnectionSettings(host=bouncer.get_container_host_ip(), port=int(bouncer.get_exposed_port(5432)),
                                      user="test", password=SecretStr("test"), database="test"),
        application_name="probe",
        query=QuerySettings(statement_cache_size=caches[0], prepared_statement_cache_size=caches[1]),
        use_orjson_serialization=False,
    )
    manager = create_async_session_manager(config, connection_class=connection_class)
    names_generated.clear()
    outcomes = Counter()

    async def worker(n: int) -> None:
        for i in range(20):
            try:
                async with manager.get_transaction() as s:
                    await s.execute(text("SELECT CAST(:v AS integer) + 1"), {"v": n})
                    await s.execute(text("SELECT CAST(:v AS text) || 'x'"), {"v": str(i)})
                outcomes["ok"] += 1
            except Exception as exc:
                cause = getattr(getattr(exc, "orig", None), "__cause__", None) or exc
                message = re.sub(r'"__asyncpg_stmt_[^"]*"', '"__asyncpg_stmt_..."', str(cause).splitlines()[0])
                outcomes[f"{type(cause).__name__}: {message}"] += 1

    await asyncio.gather(*(worker(n) for n in range(4)))
    await manager.aclose()
    lines = [f"ok={outcomes.pop('ok', 0)}, statement names generated={sum(names_generated.values())}"]
    lines += [f"      {count:3} x {error}" for error, count in outcomes.most_common()]
    return "\n".join(lines)


async def main(untracked: DockerContainer, tracking: DockerContainer) -> None:
    print("PgBouncer 1.25.2, pool_mode=transaction, default_pool_size=2; 4 workers x 20 transactions x 2 statements")
    for bouncer, label in ((untracked, "max_prepared_statements=0"), (tracking, "max_prepared_statements=200 (default)")):
        print(f"\n{label}")
        for caches in ((0, 0), (100, 0), (100, 100), (0, 100)):
            for connection_class in (CountingAsyncCConnection, CountingPlainConnection):
                name = "AsyncCConnection" if connection_class is CountingAsyncCConnection else "asyncpg.Connection"
                result = await workload(bouncer, caches, connection_class)
                print(f"  caches {caches[0]:>3}/{caches[1]:<3} {name:18} -> {result}")


with Network() as network, PostgresContainer(
    "postgres:17-alpine", username="test", password="test", dbname="test"
).with_network(network).with_network_aliases("db"):
    untracked = start_pgbouncer(network, default_pool_size=2, max_client_conn=100, max_prepared_statements=0)
    tracking = start_pgbouncer(network, default_pool_size=2, max_client_conn=100)
    try:
        asyncio.run(main(untracked, tracking))
    finally:
        untracked.stop()
        tracking.stop()