Spaces:
Running
Running
File size: 3,446 Bytes
69f2a47 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | 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)
|