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
109
110
|
import argparse
import datetime
import pathlib
import re
import uuid
MIGRATION_FILE_PATT = re.compile(
r"^(?P<year>\d{4})(?P<month>\d{2})(?P<day>\d{2})_(?P<from>\w+)_(?P<direction>(:?up)|(:?down))_(?P<to>\w+).(?P<format>(:?sql)|(:?py))$"
)
def parse_args():
parser = argparse.ArgumentParser(
description="Automatically generates a new migration file for asyncpg-trek."
)
parser.add_argument(
"-f", "--folder", default="migrations", help="the migration folder to check"
)
parser.add_argument(
"--dry-run",
action="store_true",
help="whether to perform a dry run of this script",
)
parser.add_argument(
"-t",
"--type",
choices=["sql", "py"],
default="sql",
help="the type of file to produce. defaults to .sql",
)
parser.add_argument(
"-d",
"--direction",
choices=["up", "down"],
default=None,
help="direction of migration to generate. defaults to up & down",
)
parser.add_argument(
"--preserve-env",
default=False,
action="store_true",
help="whether to not modify the .env with the new revision",
)
return parser.parse_args()
def main():
args = parse_args()
migration_folder = pathlib.Path(args.folder)
revisions = {}
possible_unlinked = []
for file in migration_folder.iterdir():
if not file.is_file():
continue
match = MIGRATION_FILE_PATT.match(file.name)
if not match:
continue
if match.group("direction") != "up":
continue
from_rev = match.group("from")
to_rev = match.group("to")
revisions[from_rev] = to_rev
if to_rev not in revisions:
possible_unlinked.append(to_rev)
if from_rev in possible_unlinked:
possible_unlinked.remove(from_rev)
from_: str
to = uuid.uuid4().hex[0:6]
if len(revisions) == 0:
from_ = "initial"
elif (
len(possible_unlinked) == 1
): # this most likely means the lone unlinked pair remains
from_ = possible_unlinked[0]
else:
keys = list(revisions.keys())
for revision in possible_unlinked:
if revision not in keys:
from_ = revision
break
else:
raise RuntimeError("last node in linked list not found")
today = datetime.datetime.now()
month = f"0{today.month}" if today.month < 10 else str(today.month)
day = f"0{today.day}" if today.day < 10 else str(today.day)
filename = f"{today.year}{month}{day}_{from_}_{{direction}}_{to}.{args.type}"
directions: tuple[str, ...] = (
(args.direction,) if args.direction is not None else ("up", "down")
)
for direction in directions:
file = migration_folder / filename.format(direction=direction)
if args.dry_run:
print("Made file", file)
else:
file.touch()
print("Done! New revision is now", to)
if __name__ == "__main__":
main()
|
