Add seed cleanup script (#55)

* Add seed cleanup script

Signed-off-by: kerthcet <kerthcet@gmail.com>

* fix lint

Signed-off-by: kerthcet <kerthcet@gmail.com>

---------

Signed-off-by: kerthcet <kerthcet@gmail.com>

Author: kerthcet

Makefile

@@ -45,4 +45,8 @@ test-all: test test-integration
 
 .PHONY: seed
 seed:
-	python hack/seed.py
+	python hack/seed.py seed
+
+.PHONY: seed-cleanup
+seed-cleanup:
+	python hack/seed.py cleanup

hack/seed.py

@@ -2,6 +2,7 @@
 
 import os
 import random
+import sys
 import uuid
 from datetime import datetime
 from decimal import Decimal
@@ -115,30 +116,69 @@ def seed_all(
     session.add_all(projects)
     session.commit()
 
-    experiments = [generate_experiment(projects) for _ in range(num_exps_per_project)]
+    experiments = [
+        generate_experiment(projects)
+        for _ in range(num_exps_per_project)
+        for _ in range(len(projects))
+    ]
     session.add_all(experiments)
     session.commit()
 
-    trials = [generate_trial([exp]) for exp in experiments]
+    trials = [
+        generate_trial(experiments)
+        for _ in range(num_trials_per_exp)
+        for _ in range(len(experiments))
+    ]
     session.add_all(trials)
     session.commit()
 
-    runs = [generate_run(trials) for _ in range(num_runs_per_trial)]
+    runs = [
+        generate_run(trials)
+        for _ in range(num_runs_per_trial)
+        for _ in range(len(trials))
+    ]
     session.add_all(runs)
     session.commit()
 
-    metrics = [generate_metric(runs) for _ in range(num_metrics_per_run)]
+    metrics = [
+        generate_metric(runs)
+        for _ in range(num_metrics_per_run)
+        for _ in range(len(runs))
+    ]
     session.add_all(metrics)
     session.commit()
 
     print("🌳 seeding completed.")
 
 
+def cleanup():
+    print("🧹 cleaning up seeded data ...")
+    session.query(Metric).delete()
+    session.query(Run).delete()
+    session.query(Trial).delete()
+    session.query(Experiment).delete()
+    session.query(Project).delete()
+    session.commit()
+    print("🧼 cleanup completed.")
+
+
 if __name__ == "__main__":
-    seed_all(
-        num_projects=3,
-        num_exps_per_project=20,
-        num_trials_per_exp=50,
-        num_runs_per_trial=100,
-        num_metrics_per_run=100,
-    )
+    if len(sys.argv) < 2:  # noqa: PLR2004
+        print("Usage: python script.py [cleanup|seed]")
+        sys.exit(1)
+
+    action = sys.argv[1]
+
+    if action == "cleanup":
+        cleanup()
+    elif action == "seed":
+        seed_all(
+            num_projects=3,
+            num_exps_per_project=10,
+            num_trials_per_exp=10,
+            num_runs_per_trial=20,
+            num_metrics_per_run=30,
+        )
+    else:
+        print(f"Unknown action: {action}")
+        sys.exit(1)

site/docs/development.md

@@ -34,6 +34,12 @@ To seed the database with sample data, run:
 make seed
 ```
 
+To cleanup, run:
+
+```bash
+make seed-cleanup
+```
+
 ## How to Test
 
 To run tests, use: