Initial Commit

mvp-platform-services/vehicles/etl/connections.py

@@ -0,0 +1,152 @@
+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()