Refactor all main.py

services/forecasting/app/core/database.py

@@ -7,67 +7,74 @@ Database configuration for forecasting service
 
 import structlog
 from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
-from sqlalchemy.pool import NullPool
 from sqlalchemy import text
+from typing import AsyncGenerator
 
 from app.core.config import settings
-from shared.database.base import Base
+from shared.database.base import Base, DatabaseManager
 
 logger = structlog.get_logger()
 
 # Create async engine
-engine = create_async_engine(
+async_engine = create_async_engine(
     settings.DATABASE_URL,
-    poolclass=NullPool,
     echo=settings.DEBUG,
-    future=True
+    pool_size=10,
+    max_overflow=20,
+    pool_pre_ping=True,
+    pool_recycle=3600
 )
 
-# Create session factory
+# Create async session factory
 AsyncSessionLocal = async_sessionmaker(
-    engine,
+    bind=async_engine,
     class_=AsyncSession,
-    expire_on_commit=False,
-    autoflush=False,
-    autocommit=False
+    expire_on_commit=False
 )
 
-class DatabaseManager:
-    """Database management operations"""
-    
-    async def create_tables(self):
-        """Create database tables"""
-        async with engine.begin() as conn:
+async def get_db() -> AsyncGenerator[AsyncSession, None]:
+    """Get database session"""
+    async with AsyncSessionLocal() as session:
+        try:
+            yield session
+        except Exception as e:
+            await session.rollback()
+            logger.error("Database session error", error=str(e))
+            raise
+        finally:
+            await session.close()
+
+async def init_database():
+    """Initialize database tables"""
+    try:
+        async with async_engine.begin() as conn:
+            # Import all models to ensure they are registered
+            from app.models.forecast import ForecastBatch, Forecast
+            from app.models.prediction import PredictionBatch, Prediction
+
+            # Create all tables
             await conn.run_sync(Base.metadata.create_all)
-        logger.info("Forecasting database tables created successfully")
-    
-    async def get_session(self) -> AsyncSession:
-        """Get database session"""
-        async with AsyncSessionLocal() as session:
-            try:
-                yield session
-                await session.commit()
-            except Exception as e:
-                await session.rollback()
-                logger.error(f"Database session error: {e}")
-                raise
-            finally:
-                await session.close()
 
-# Global database manager instance
-database_manager = DatabaseManager()
-
-async def get_db() -> AsyncSession:
-    """Database dependency"""
-    async for session in database_manager.get_session():
-        yield session
+        logger.info("Forecasting database initialized successfully")
+    except Exception as e:
+        logger.error("Failed to initialize forecasting database", error=str(e))
+        raise
 
 async def get_db_health() -> bool:
     """Check database health"""
     try:
-        async with AsyncSessionLocal() as session:
-            await session.execute(text("SELECT 1"))
-            return True
+        async with async_engine.begin() as conn:
+            await conn.execute(text("SELECT 1"))
+        return True
     except Exception as e:
-        logger.error(f"Database health check failed: {e}")
+        logger.error("Database health check failed", error=str(e))
         return False
+
+# Database manager instance for service_base compatibility
+database_manager = DatabaseManager(
+    database_url=settings.DATABASE_URL,
+    service_name="forecasting-service",
+    pool_size=10,
+    max_overflow=20,
+    echo=settings.DEBUG
+)