forked from JJeshua/ddsm-backend-testing-tool
-
Notifications
You must be signed in to change notification settings - Fork 0
/
db_tools.py
46 lines (34 loc) · 1.1 KB
/
db_tools.py
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
import argparse
import subprocess
import cProfile
import pstats
def populate_database():
subprocess.run(["python", "./db/populate_db.py"])
def clear_database():
subprocess.run(["python", "./db/clear_db.py"])
def main():
parser = argparse.ArgumentParser(description="Test Script Menu")
parser.add_argument(
"action", choices=["populate", "clear"], help="Action to perform"
)
parser.add_argument(
"--profile", action="store_true", help="Enable profiling with cProfile."
)
args = parser.parse_args()
if args.profile:
with cProfile.Profile() as profile:
perform_action(args.action)
print("Successfully performed action:", args.action)
results = pstats.Stats(profile)
results.sort_stats(pstats.SortKey.TIME)
results.print_stats()
else:
perform_action(args.action)
print("Successfully performed action:", args.action)
def perform_action(action):
if action == "populate":
populate_database()
elif action == "clear":
clear_database()
if __name__ == "__main__":
main()