""" Database Setup Module Handles database initialization and table creation. """ from typing import List, Tuple import sys from pathlib import Path from loguru import logger class DatabaseSetup: """Manages database setup for ALwrity backend.""" def __init__(self, production_mode: bool = False): self.production_mode = production_mode def setup_essential_tables(self) -> bool: """Set up essential database tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" if verbose: print("📊 Setting up essential database tables...") try: from services.database import init_database, engine # Initialize database connection init_database() if verbose: print(" ✅ Database connection initialized") # Create essential tables self._create_monitoring_tables() self._create_subscription_tables() self._create_persona_tables() self._create_onboarding_tables() self._create_daily_workflow_tables() if verbose: print("✅ Essential database tables created") return True except Exception as e: if verbose: print(f"âš ī¸ Warning: Database setup failed: {e}") if self.production_mode: print(" Continuing in production mode...") else: print(" This may affect functionality") return True # Don't fail startup for database issues def _create_monitoring_tables(self) -> bool: """Create API monitoring tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" try: from models.api_monitoring import Base as MonitoringBase MonitoringBase.metadata.create_all(bind=engine) if verbose: print(" ✅ Monitoring tables created") return True except Exception as e: if verbose: print(f" âš ī¸ Monitoring tables failed: {e}") return True # Non-critical def _create_subscription_tables(self) -> bool: """Create subscription and billing tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" try: from models.subscription_models import Base as SubscriptionBase SubscriptionBase.metadata.create_all(bind=engine) if verbose: print(" ✅ Subscription tables created") return True except Exception as e: if verbose: print(f" âš ī¸ Subscription tables failed: {e}") return True # Non-critical def _create_persona_tables(self) -> bool: """Create persona analysis tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" try: from models.persona_models import Base as PersonaBase PersonaBase.metadata.create_all(bind=engine) if verbose: print(" ✅ Persona tables created") return True except Exception as e: if verbose: print(f" âš ī¸ Persona tables failed: {e}") return True # Non-critical def _create_onboarding_tables(self) -> bool: """Create onboarding tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" try: from models.onboarding import Base as OnboardingBase OnboardingBase.metadata.create_all(bind=engine) if verbose: print(" ✅ Onboarding tables created") return True except Exception as e: if verbose: print(f" âš ī¸ Onboarding tables failed: {e}") return True # Non-critical def _create_daily_workflow_tables(self) -> bool: """Create daily workflow tables.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" try: from models.enhanced_strategy_models import Base as StrategyBase StrategyBase.metadata.create_all(bind=engine) if verbose: print(" ✅ Daily workflow tables created") return True except Exception as e: if verbose: print(f" âš ī¸ Daily workflow tables failed: {e}") return True # Non-critical def verify_tables(self) -> bool: """Verify that essential tables exist.""" import os verbose = os.getenv("ALWRITY_VERBOSE", "false").lower() == "true" if self.production_mode: if verbose: print("âš ī¸ Skipping table verification in production mode") return True if verbose: print("🔍 Verifying database tables...") try: from services.database import engine from sqlalchemy import inspect if engine is None: if verbose: print(" âš ī¸ Global engine is None (Multi-tenant mode), skipping global table verification") return True inspector = inspect(engine) tables = inspector.get_table_names() essential_tables = [ 'api_monitoring_logs', 'subscription_plans', 'user_subscriptions', 'onboarding_sessions', 'persona_data' ] existing_tables = [table for table in essential_tables if table in tables] if verbose: print(f" ✅ Found tables: {existing_tables}") if len(existing_tables) < len(essential_tables): missing = [table for table in essential_tables if table not in existing_tables] if verbose: print(f" âš ī¸ Missing tables: {missing}") return True except Exception as e: print(f" âš ī¸ Table verification failed: {e}") return True # Non-critical def setup_advanced_tables(self) -> bool: """Set up advanced tables (non-critical).""" if self.production_mode: print("âš ī¸ Skipping advanced table setup in production mode") return True print("🔧 Setting up advanced database features...") try: # Set up monitoring tables self._setup_monitoring_tables() # Set up billing tables self._setup_billing_tables() logger.debug("✅ Advanced database features configured") return True except Exception as e: logger.warning(f"Advanced table setup failed: {e}") return True # Non-critical def _setup_monitoring_tables(self) -> bool: """Set up API monitoring tables.""" # Reuse the existing method that uses SQLAlchemy metadata # This avoids the script dependency that requires user_id return self._create_monitoring_tables() def _setup_billing_tables(self) -> bool: """Set up billing and subscription tables.""" try: sys.path.append(str(Path(__file__).parent.parent)) from scripts.create_billing_tables import create_billing_tables, check_existing_tables from services.database import engine # Check if engine is available (it might be None in multi-tenant mode) if engine is None: # In multi-tenant mode, we can't setup global billing tables # They will be created per-user when they are initialized return True # Check if tables already exist if check_existing_tables(engine): logger.debug("✅ Billing tables already exist") return True # For global setup, we can't call create_billing_tables() without user_id # But if engine is not None, it implies we have a global DB. # However, the script is designed for user_id. # We'll skip this call to avoid the TypeError and rely on per-user init. logger.debug("â„šī¸ Skipping global billing table creation (handled per-user)") return True except Exception as e: logger.warning(f"Billing setup failed: {e}") return True # Non-critical