153 lines
5.3 KiB
Python
Executable File
153 lines
5.3 KiB
Python
Executable File
import pyodbc
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
import asyncpg
|
|
import redis
|
|
from contextlib import contextmanager
|
|
import logging
|
|
import time
|
|
from typing import Optional
|
|
from .config import config
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class DatabaseConnections:
|
|
"""Manage database connections with retry logic and timeouts"""
|
|
|
|
def __init__(self):
|
|
self.mssql_conn = None
|
|
self.postgres_conn = None
|
|
self.redis_client = None
|
|
self.pg_pool = None
|
|
self.max_retries = 3
|
|
self.retry_delay = 2 # seconds
|
|
|
|
def _retry_connection(self, connection_func, connection_type: str, max_retries: Optional[int] = None):
|
|
"""Retry connection with exponential backoff"""
|
|
max_retries = max_retries or self.max_retries
|
|
|
|
for attempt in range(max_retries):
|
|
try:
|
|
return connection_func()
|
|
except Exception as e:
|
|
if attempt == max_retries - 1:
|
|
logger.error(f"Failed to connect to {connection_type} after {max_retries} attempts: {e}")
|
|
raise
|
|
|
|
wait_time = self.retry_delay * (2 ** attempt)
|
|
logger.warning(f"{connection_type} connection failed (attempt {attempt + 1}/{max_retries}): {e}")
|
|
logger.info(f"Retrying in {wait_time} seconds...")
|
|
time.sleep(wait_time)
|
|
|
|
@contextmanager
|
|
def mssql_connection(self):
|
|
"""Context manager for MS SQL connection using pyodbc with retry logic"""
|
|
def _connect():
|
|
connection_string = (
|
|
f"DRIVER={{ODBC Driver 17 for SQL Server}};"
|
|
f"SERVER={config.MSSQL_HOST},{config.MSSQL_PORT};"
|
|
f"DATABASE={config.MSSQL_DATABASE};"
|
|
f"UID={config.MSSQL_USER};"
|
|
f"PWD={config.MSSQL_PASSWORD};"
|
|
f"TrustServerCertificate=yes;"
|
|
f"Connection Timeout=30;"
|
|
f"Command Timeout=300;"
|
|
)
|
|
return pyodbc.connect(connection_string)
|
|
|
|
conn = self._retry_connection(_connect, "MSSQL")
|
|
try:
|
|
yield conn
|
|
finally:
|
|
try:
|
|
conn.close()
|
|
except Exception as e:
|
|
logger.warning(f"Error closing MSSQL connection: {e}")
|
|
|
|
@contextmanager
|
|
def postgres_connection(self):
|
|
"""Context manager for PostgreSQL connection with retry logic"""
|
|
def _connect():
|
|
return psycopg2.connect(
|
|
host=config.POSTGRES_HOST,
|
|
port=config.POSTGRES_PORT,
|
|
database=config.POSTGRES_DATABASE,
|
|
user=config.POSTGRES_USER,
|
|
password=config.POSTGRES_PASSWORD,
|
|
cursor_factory=RealDictCursor,
|
|
connect_timeout=30,
|
|
options='-c statement_timeout=300000' # 5 minutes
|
|
)
|
|
|
|
conn = self._retry_connection(_connect, "PostgreSQL")
|
|
try:
|
|
yield conn
|
|
finally:
|
|
try:
|
|
conn.close()
|
|
except Exception as e:
|
|
logger.warning(f"Error closing PostgreSQL connection: {e}")
|
|
|
|
async def create_pg_pool(self):
|
|
"""Create async PostgreSQL connection pool"""
|
|
self.pg_pool = await asyncpg.create_pool(
|
|
host=config.POSTGRES_HOST,
|
|
port=config.POSTGRES_PORT,
|
|
database=config.POSTGRES_DATABASE,
|
|
user=config.POSTGRES_USER,
|
|
password=config.POSTGRES_PASSWORD,
|
|
min_size=10,
|
|
max_size=20
|
|
)
|
|
return self.pg_pool
|
|
|
|
def get_redis_client(self):
|
|
"""Get Redis client"""
|
|
if not self.redis_client:
|
|
self.redis_client = redis.Redis(
|
|
host=config.REDIS_HOST,
|
|
port=config.REDIS_PORT,
|
|
db=config.REDIS_DB,
|
|
decode_responses=True
|
|
)
|
|
return self.redis_client
|
|
|
|
def test_connections():
|
|
"""Test all database connections for health check"""
|
|
try:
|
|
# Test MSSQL connection (use master DB to avoid failures before restore)
|
|
db = DatabaseConnections()
|
|
mssql_master_conn_str = (
|
|
f"DRIVER={{ODBC Driver 17 for SQL Server}};"
|
|
f"SERVER={config.MSSQL_HOST},{config.MSSQL_PORT};"
|
|
f"DATABASE=master;"
|
|
f"UID={config.MSSQL_USER};"
|
|
f"PWD={config.MSSQL_PASSWORD};"
|
|
f"TrustServerCertificate=yes;"
|
|
)
|
|
import pyodbc as _pyodbc
|
|
with _pyodbc.connect(mssql_master_conn_str) as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT 1")
|
|
cursor.fetchone()
|
|
logger.info("MSSQL connection successful (master)")
|
|
|
|
# Test PostgreSQL connection
|
|
with db.postgres_connection() as conn:
|
|
cursor = conn.cursor()
|
|
cursor.execute("SELECT 1")
|
|
cursor.fetchone()
|
|
logger.info("PostgreSQL connection successful")
|
|
|
|
# Test Redis connection
|
|
redis_client = db.get_redis_client()
|
|
redis_client.ping()
|
|
logger.info("Redis connection successful")
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Connection test failed: {e}")
|
|
return False
|
|
|
|
db_connections = DatabaseConnections()
|