File size: 3,104 Bytes
8bab08d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
107
108
"""
Database Migration Management Script
Provides helper functions for managing database migrations with Alembic
"""
import os
import sys
import logging
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from alembic.config import Config
from alembic import command

logger = logging.getLogger(__name__)


def get_alembic_config() -> Config:
    """Get Alembic configuration"""
    # Path to alembic.ini
    alembic_ini = Path(__file__).parent.parent.parent / "alembic.ini"

    if not alembic_ini.exists():
        raise FileNotFoundError(f"alembic.ini not found at {alembic_ini}")

    config = Config(str(alembic_ini))
    return config


def create_migration(message: str):
    """Create a new migration"""
    config = get_alembic_config()
    command.revision(config, message=message, autogenerate=True)
    logger.info(f"Created migration: {message}")


def upgrade_database(revision: str = "head"):
    """Upgrade database to a revision"""
    config = get_alembic_config()
    command.upgrade(config, revision)
    logger.info(f"Upgraded database to {revision}")


def downgrade_database(revision: str):
    """Downgrade database to a revision"""
    config = get_alembic_config()
    command.downgrade(config, revision)
    logger.info(f"Downgraded database to {revision}")


def show_current_revision():
    """Show current database revision"""
    config = get_alembic_config()
    command.current(config)


def show_migration_history():
    """Show migration history"""
    config = get_alembic_config()
    command.history(config)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Database Migration Management")
    subparsers = parser.add_subparsers(dest="command", help="Command to run")

    # Create migration
    create_parser = subparsers.add_parser("create", help="Create a new migration")
    create_parser.add_argument("message", help="Migration message")

    # Upgrade database
    upgrade_parser = subparsers.add_parser("upgrade", help="Upgrade database")
    upgrade_parser.add_argument(
        "--revision",
        default="head",
        help="Revision to upgrade to (default: head)"
    )

    # Downgrade database
    downgrade_parser = subparsers.add_parser("downgrade", help="Downgrade database")
    downgrade_parser.add_argument("revision", help="Revision to downgrade to")

    # Show current revision
    subparsers.add_parser("current", help="Show current database revision")

    # Show history
    subparsers.add_parser("history", help="Show migration history")

    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO)

    if args.command == "create":
        create_migration(args.message)
    elif args.command == "upgrade":
        upgrade_database(args.revision)
    elif args.command == "downgrade":
        downgrade_database(args.revision)
    elif args.command == "current":
        show_current_revision()
    elif args.command == "history":
        show_migration_history()
    else:
        parser.print_help()