milestone runs in db

Author: ngc92

src/libkernelbot/leaderboard_db.py

@@ -293,6 +293,58 @@ def get_leaderboard_milestones(self, leaderboard_id: int) -> "list[MilestoneItem
             for row in self.cursor.fetchall()
         ]
 
+    def delete_milestone_runs(self, leaderboard_name: str):
+        self.cursor.execute(
+            """
+            DELETE FROM leaderboard.runs
+            WHERE milestone_id IN (
+                SELECT leaderboard.milestones.id
+                FROM leaderboard.milestones
+                WHERE leaderboard_id = %s
+            );
+            """,
+            (leaderboard_name,),
+        )
+        self.connection.commit()
+
+    def get_runs_generic(
+        self, *, milestone_id: Optional[int] = None, submission_id: Optional[int] = None
+    ) -> List["RunItem"]:
+        if milestone_id is not None:
+            key = "milestone_id"
+            value = milestone_id
+            if submission_id is not None:
+                logger.error("milestone_id and submission_id specified simultaneously")
+                raise KernelBotError("`milestone_id` and `submission_id` specified simultaneously")
+        else:
+            key = "submission_id"
+            value = submission_id
+        query = f"""
+                SELECT start_time, end_time, mode, secret, runner, score,
+                       passed, compilation, meta, result, system_info
+                FROM leaderboard.runs
+                WHERE {key} = %s
+                """
+        self.cursor.execute(query, (value,))
+        runs = self.cursor.fetchall()
+
+        return [
+            RunItem(
+                start_time=r[0],
+                end_time=r[1],
+                mode=r[2],
+                secret=r[3],
+                runner=r[4],
+                score=r[5],
+                passed=r[6],
+                compilation=r[7],
+                meta=r[8],
+                result=r[9],
+                system=r[10],
+            )
+            for r in runs
+        ]
+
     def create_submission(
         self,
         leaderboard: str,
@@ -397,7 +449,9 @@ def mark_submission_done(
 
     def create_submission_run(
         self,
-        submission: int,
+        *,
+        submission: Optional[int] = None,
+        milestone: Optional[int] = None,
         start: datetime.datetime,
         end: datetime.datetime,
         mode: str,
@@ -418,13 +472,14 @@ def create_submission_run(
             }
             self.cursor.execute(
                 """
-                INSERT INTO leaderboard.runs (submission_id, start_time, end_time, mode,
+                INSERT INTO leaderboard.runs (submission_id, milestone_id, start_time, end_time, mode,
                 secret, runner, score, passed, compilation, meta, result, system_info
                 )
-                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
+                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
                 """,
                 (
                     submission,
+                    milestone,
                     start,
                     end,
                     mode,
@@ -804,33 +859,6 @@ def get_submission_by_id(self, submission_id: int) -> Optional["SubmissionItem"]
         if submission is None:
             return None
 
-        # OK, now get the runs
-        query = """
-                SELECT start_time, end_time, mode, secret, runner, score,
-                       passed, compilation, meta, result, system_info
-                FROM leaderboard.runs
-                WHERE submission_id = %s
-                """
-        self.cursor.execute(query, (submission_id,))
-        runs = self.cursor.fetchall()
-
-        runs = [
-            RunItem(
-                start_time=r[0],
-                end_time=r[1],
-                mode=r[2],
-                secret=r[3],
-                runner=r[4],
-                score=r[5],
-                passed=r[6],
-                compilation=r[7],
-                meta=r[8],
-                result=r[9],
-                system=r[10],
-            )
-            for r in runs
-        ]
-
         return SubmissionItem(
             submission_id=submission_id,
             leaderboard_id=submission[0],
@@ -840,7 +868,7 @@ def get_submission_by_id(self, submission_id: int) -> Optional["SubmissionItem"]
             submission_time=submission[4],
             done=submission[5],
             code=submission[6],
-            runs=runs,
+            runs=self.get_runs_generic(submission_id=submission_id),
         )
 
     def get_leaderboard_submission_count(