import asyncio
import gc
import time
import tracemalloc
from asyncio import current_task

from sqlalchemy.ext.asyncio import async_scoped_session, AsyncSession, create_async_engine
from sqlalchemy.orm import sessionmaker


DB_URL = 'postgresql+asyncpg://<your-url-here>'


def create_async_scoped_db_session(
    url: str, statement_timeout_ms: int = 60 * 1000
) -> async_scoped_session:
    session = async_scoped_session(sessionmaker(class_=AsyncSession), scopefunc=current_task)

    url += f'?command_timeout={statement_timeout_ms}'
    engine = create_async_engine(url)

    session.configure(bind=engine)

    return session


async def job(id, dbs, wait_for=1):
    # just do some dummy stuff
    print(f'Starting waiting job #{id} for {wait_for}s at {time.time()}. {dbs.bind.engine.pool.status()}')
    await dbs.execute(f'select 1;')
    await dbs.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public';")
    print(f'Waiting job #{id} ready at {time.time()}')
    await dbs.execute(f'select pg_sleep({wait_for});')
    print(f'Finished waiting job #{id} for {wait_for}s at {time.time()}')
    await dbs.close()
    await dbs.remove()
    print(f'Starting waiting job #{id} cleaned up {time.time()}. {dbs.bind.engine.pool.status()}')


async def run():
    tracemalloc.start()
    dbs = create_async_scoped_db_session(DB_URL)
    snapshot1 = tracemalloc.take_snapshot()

    num_rounds = 300
    num_jobs = 10

    for _ in range(num_rounds):
        await asyncio.gather(*(job(id=i+1, dbs=dbs) for i in range(num_jobs)))
        snapshot2 = tracemalloc.take_snapshot()
        top_stats = snapshot2.compare_to(snapshot1, 'lineno')

        print("[ Top 10 differences ]")
        for stat in top_stats[:10]:
            print(stat)

        snapshot1 = snapshot2

    snapshot1 = snapshot2

    await dbs.bind.dispose()

    print('Stats after disposal')

    snapshot2 = tracemalloc.take_snapshot()
    top_stats = snapshot2.compare_to(snapshot1, 'lineno')

    print("[ Top 10 differences ]")
    for stat in top_stats[:10]:
        print(stat)

    print('Running garbage collection')

    gc.collect()

    print('Stats after collection')

    snapshot2 = tracemalloc.take_snapshot()
    top_stats = snapshot2.compare_to(snapshot1, 'lineno')

    print("[ Top 10 differences ]")
    for stat in top_stats[:10]:
        print(stat)

    top_stats = snapshot2.statistics('lineno')

    print("[ Top 10 - FINAL stats ]")
    for stat in top_stats[:10]:
        print(stat)


if __name__ == "__main__":
    asyncio.run(
        run()
    )