Spaces:
Sleeping
Sleeping
File size: 1,431 Bytes
a23bdc6 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 |
import asyncio
import os
from contextlib import asynccontextmanager
import asyncpg
from dotenv import load_dotenv
# Global connection pool
_pool = None
load_dotenv()
async def get_pool(schema="talmudexplore", min_size=2, max_size=5):
"""Initialize and return the connection pool with the specified schema."""
global _pool
if _pool is not None:
current_loop = asyncio.get_running_loop()
if getattr(_pool, '_loop', None) != current_loop:
try:
await _pool.close()
except:
pass
_pool = None
if _pool is None:
_pool = await asyncpg.create_pool(
database=os.getenv("pg_dbname"),
user=os.getenv("pg_user"),
password=os.getenv("pg_password"),
host=os.getenv("pg_host"),
port=os.getenv("pg_port"),
min_size=min_size,
max_size=max_size,
setup=lambda conn: conn.execute(f'SET search_path TO {schema}')
)
return _pool
@asynccontextmanager
async def get_async_connection():
"""Get a connection from the pool as an async context manager."""
pool = await get_pool()
conn = await pool.acquire()
try:
yield conn
finally:
await pool.release(conn)
async def close_pool():
"""Close the connection pool."""
global _pool
if _pool:
await _pool.close()
_pool = None
|