Spaces:
Paused
Paused
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)
|