Spaces:
Running
Running
| import asyncio | |
| import os | |
| import logging | |
| import asyncpg | |
| from dotenv import load_dotenv | |
| # Configure logging | |
| logging.basicConfig( | |
| level=logging.INFO, | |
| format='%(asctime)s - %(name)s - %(levelname)s - %(message)s' | |
| ) | |
| logger = logging.getLogger("db_migration") | |
| # Load environment variables from .env file in the project root | |
| project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) | |
| env_path = os.path.join(project_root, '.env') | |
| if os.path.exists(env_path): | |
| load_dotenv(env_path) | |
| logger.info(f"Loaded environment variables from {env_path}") | |
| else: | |
| logger.warning(f".env file not found at {env_path}, using environment variables or defaults") | |
| # Database configuration | |
| DB_HOST = os.getenv("DB_HOST", "localhost") | |
| DB_PORT = os.getenv("DB_PORT", "5432") | |
| DB_NAME = os.getenv("DB_NAME", "postgres") | |
| DB_USER = os.getenv("DB_USER", "postgres") | |
| DB_PASSWORD = os.getenv("DB_PASSWORD", "") | |
| DB_SSLMODE = os.getenv("DB_SSLMODE", "disable") | |
| async def run_migration(): | |
| """ | |
| Connects to the database and executes migration scripts in order. | |
| """ | |
| # Construct DSN | |
| dsn = f"postgresql://{DB_USER}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}" | |
| logger.info(f"Connecting to database: {DB_HOST}:{DB_PORT}/{DB_NAME} as {DB_USER}") | |
| conn = None | |
| try: | |
| # SSL context | |
| ssl_ctx = False | |
| if DB_SSLMODE.lower() != "disable": | |
| import ssl | |
| ssl_ctx = ssl.create_default_context() | |
| ssl_ctx.check_hostname = False | |
| ssl_ctx.verify_mode = ssl.CERT_NONE | |
| logger.info(f"SSL enabled (mode: {DB_SSLMODE})") | |
| conn = await asyncpg.connect(dsn, ssl=ssl_ctx) | |
| logger.info("Successfully connected to database") | |
| # List of scripts to execute in order | |
| scripts = [ | |
| "01_schema.sql", | |
| "02_functions.sql", | |
| "03_views.sql" | |
| ] | |
| script_dir = os.path.dirname(os.path.abspath(__file__)) | |
| for script_name in scripts: | |
| script_path = os.path.join(script_dir, script_name) | |
| if not os.path.exists(script_path): | |
| logger.error(f"Script not found: {script_path}") | |
| continue | |
| logger.info(f"Executing script: {script_name}") | |
| with open(script_path, 'r') as f: | |
| sql_content = f.read() | |
| if not sql_content.strip(): | |
| logger.warning(f"Script {script_name} is empty, skipping") | |
| continue | |
| # Execute the script | |
| # asyncpg can execute multiple statements separated by semicolons | |
| try: | |
| await conn.execute(sql_content) | |
| logger.info(f"Successfully executed {script_name}") | |
| except Exception as e: | |
| logger.error(f"Error executing {script_name}: {e}") | |
| raise | |
| logger.info("Migration completed successfully") | |
| except Exception as e: | |
| logger.error(f"Migration failed: {e}") | |
| raise | |
| finally: | |
| if conn: | |
| await conn.close() | |
| logger.info("Database connection closed") | |
| if __name__ == "__main__": | |
| try: | |
| asyncio.run(run_migration()) | |
| except KeyboardInterrupt: | |
| logger.info("Migration interrupted by user") | |
| except Exception as e: | |
| logger.error("Migration script failed") | |
| exit(1) | |