103 lines
3.4 KiB
Python
103 lines
3.4 KiB
Python
"""
|
|
Database connection module for Python service
|
|
Uses asyncpg for async PostgreSQL connections
|
|
"""
|
|
import os
|
|
import asyncio
|
|
import asyncpg
|
|
from typing import Optional
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
# Database connection pool
|
|
_pool: Optional[asyncpg.Pool] = None
|
|
|
|
# Connection timeout in seconds (reduced from default 60s to fail faster)
|
|
CONNECTION_TIMEOUT = 10
|
|
|
|
|
|
async def get_pool() -> asyncpg.Pool:
|
|
"""Get or create database connection pool with timeout"""
|
|
global _pool
|
|
|
|
if _pool is None:
|
|
# Check for local database first
|
|
use_local_db = os.getenv('USE_LOCAL_DB', 'false').lower() == 'true'
|
|
|
|
if use_local_db:
|
|
config = {
|
|
'host': os.getenv('LOCAL_DB_HOST', 'localhost'),
|
|
'port': int(os.getenv('LOCAL_DB_PORT', '5432')),
|
|
'user': os.getenv('LOCAL_DB_USER', 'postgres'),
|
|
'password': os.getenv('LOCAL_DB_PASSWORD', 'postgres'),
|
|
'database': os.getenv('LOCAL_DB_NAME', 'institutional_trader')
|
|
}
|
|
elif os.getenv('DATABASE_URL'):
|
|
# Parse connection string
|
|
import urllib.parse
|
|
url = urllib.parse.urlparse(os.getenv('DATABASE_URL'))
|
|
config = {
|
|
'host': url.hostname,
|
|
'port': url.port or 5432,
|
|
'user': url.username,
|
|
'password': url.password,
|
|
'database': url.path.lstrip('/')
|
|
}
|
|
else:
|
|
raise ValueError(
|
|
'No database configuration found. Set either:\n'
|
|
' - USE_LOCAL_DB=true with LOCAL_DB_* variables\n'
|
|
' - OR DATABASE_URL'
|
|
)
|
|
|
|
# Create pool with timeout to prevent hanging
|
|
# Use min_size=1 to speed up initial connection (pool will grow as needed)
|
|
try:
|
|
_pool = await asyncio.wait_for(
|
|
asyncpg.create_pool(
|
|
**config,
|
|
min_size=1, # Start with 1 connection, pool grows as needed
|
|
max_size=20,
|
|
command_timeout=60,
|
|
timeout=CONNECTION_TIMEOUT # Connection timeout per connection
|
|
),
|
|
timeout=CONNECTION_TIMEOUT # Overall timeout for pool creation
|
|
)
|
|
except asyncio.TimeoutError:
|
|
raise ConnectionError(
|
|
f'Database connection timeout after {CONNECTION_TIMEOUT}s. '
|
|
f'Check if database is reachable at {config.get("host", "unknown")}:{config.get("port", "unknown")}'
|
|
)
|
|
except Exception as e:
|
|
raise ConnectionError(
|
|
f'Failed to create database connection pool: {str(e)}. '
|
|
f'Check database configuration and network connectivity.'
|
|
)
|
|
|
|
return _pool
|
|
|
|
|
|
async def close_pool():
|
|
"""Close database connection pool"""
|
|
global _pool
|
|
if _pool:
|
|
await _pool.close()
|
|
_pool = None
|
|
|
|
|
|
async def query(sql: str, *args) -> list:
|
|
"""Execute a SQL query and return results"""
|
|
pool = await get_pool()
|
|
async with pool.acquire() as conn:
|
|
return await conn.fetch(sql, *args)
|
|
|
|
|
|
async def query_one(sql: str, *args) -> Optional[dict]:
|
|
"""Execute a SQL query and return first result"""
|
|
pool = await get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow(sql, *args)
|
|
return dict(row) if row else None
|
|
|