File size: 2,629 Bytes
4a2ab42
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Database migration rollback utilities"""

import sys

from alembic.config import Config

from alembic import command


def rollback_migration(steps: int = 1):
    """Rollback database migration by specified number of steps"""
    alembic_cfg = Config("alembic.ini")

    try:
        print(f"πŸ“‰ Rolling back {steps} migration(s)...")

        if steps == 1:
            command.downgrade(alembic_cfg, "-1")
        else:
            command.downgrade(alembic_cfg, f"-{steps}")

        print(f"βœ… Successfully rolled back {steps} migration(s)")
        return True
    except Exception as e:
        print(f"❌ Rollback failed: {e}")
        return False


def rollback_to_revision(revision: str):
    """Rollback to a specific revision"""
    alembic_cfg = Config("alembic.ini")

    try:
        print(f"πŸ“‰ Rolling back to revision: {revision}")
        command.downgrade(alembic_cfg, revision)
        print(f"βœ… Successfully rolled back to {revision}")
        return True
    except Exception as e:
        print(f"❌ Rollback failed: {e}")
        return False


def show_migration_history():
    """Display migration history"""
    alembic_cfg = Config("alembic.ini")

    try:
        print("πŸ“œ Migration History:")
        command.history(alembic_cfg)
        return True
    except Exception as e:
        print(f"❌ Failed to show history: {e}")
        return False


def get_current_revision():
    """Get current database revision"""
    alembic_cfg = Config("alembic.ini")

    try:
        command.current(alembic_cfg)
        return True
    except Exception as e:
        print(f"❌ Failed to get current revision: {e}")
        return False


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage:")
        print("  python rollback.py history           # Show migration history")
        print("  python rollback.py current           # Show current revision")
        print("  python rollback.py down [steps]      # Rollback N steps (default: 1)")
        print("  python rollback.py to <revision>     # Rollback to specific revision")
        sys.exit(1)

    action = sys.argv[1]

    if action == "history":
        show_migration_history()
    elif action == "current":
        get_current_revision()
    elif action == "down":
        steps = int(sys.argv[2]) if len(sys.argv) > 2 else 1
        rollback_migration(steps)
    elif action == "to":
        if len(sys.argv) < 3:
            print("❌ Please specify revision")
            sys.exit(1)
        rollback_to_revision(sys.argv[2])
    else:
        print(f"❌ Unknown action: {action}")
        sys.exit(1)