82 lines
2.3 KiB
Python
82 lines
2.3 KiB
Python
"""
|
|
Database connection module for Python service
|
|
Uses asyncpg for async PostgreSQL connections
|
|
"""
|
|
import os
|
|
import asyncpg
|
|
from typing import Optional
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
# Database connection pool
|
|
_pool: Optional[asyncpg.Pool] = None
|
|
|
|
|
|
async def get_pool() -> asyncpg.Pool:
|
|
"""Get or create database connection pool"""
|
|
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', '100.121.163.23'),
|
|
'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'
|
|
)
|
|
|
|
_pool = await asyncpg.create_pool(
|
|
**config,
|
|
min_size=5,
|
|
max_size=20,
|
|
command_timeout=60
|
|
)
|
|
|
|
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
|
|
|