define milestones in db

Author: ngc92

src/libkernelbot/db_types.py

@@ -56,4 +56,12 @@ class SubmissionItem(TypedDict):
     runs: List[RunItem]
 
 
-__all__ = [LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem]
+class MilestoneItem(TypedDict):
+    id: int
+    name: str
+    code: str
+    description: str
+    created_at: datetime.datetime
+
+
+__all__ = [LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem, MilestoneItem]

src/libkernelbot/leaderboard_db.py

@@ -5,7 +5,13 @@
 
 import psycopg2
 
-from libkernelbot.db_types import LeaderboardItem, LeaderboardRankedEntry, RunItem, SubmissionItem
+from libkernelbot.db_types import (
+    LeaderboardItem,
+    LeaderboardRankedEntry,
+    MilestoneItem,
+    RunItem,
+    SubmissionItem,
+)
 from libkernelbot.run_eval import CompileResult, RunResult, SystemInfo
 from libkernelbot.task import LeaderboardDefinition, LeaderboardTask
 from libkernelbot.utils import (
@@ -238,6 +244,55 @@ def delete_leaderboard(self, leaderboard_name: str, force: bool = False):
             logger.exception("Could not delete leaderboard %s.", leaderboard_name, exc_info=e)
             raise KernelBotError(f"Could not delete leaderboard {leaderboard_name}.") from e
 
+    def create_milestone(
+        self,
+        leaderboard_id: int,
+        name: str,
+        code: str,
+        description: str = None,
+    ) -> int:
+        """Create a new milestone for a leaderboard"""
+        try:
+            self.cursor.execute(
+                """
+                INSERT INTO leaderboard.milestones (
+                    leaderboard_id, name, code, description
+                )
+                VALUES (%s, %s, %s, %s)
+                RETURNING id
+                """,
+                (leaderboard_id, name, code, description),
+            )
+            milestone_id = self.cursor.fetchone()[0]
+            self.connection.commit()
+            return milestone_id
+        except psycopg2.Error as e:
+            self.connection.rollback()
+            logger.exception("Error creating milestone", exc_info=e)
+            raise KernelBotError("Error creating milestone") from e
+
+    def get_leaderboard_milestones(self, leaderboard_id: int) -> "list[MilestoneItem]":
+        """Get all milestones for a leaderboard"""
+        self.cursor.execute(
+            """
+            SELECT id, name, code, description, created_at
+            FROM leaderboard.milestones
+            WHERE leaderboard_id = %s
+            ORDER BY created_at
+            """,
+            (leaderboard_id,),
+        )
+        return [
+            {
+                "id": row[0],
+                "name": row[1],
+                "code": row[2],
+                "description": row[3],
+                "created_at": row[4],
+            }
+            for row in self.cursor.fetchall()
+        ]
+
     def create_submission(
         self,
         leaderboard: str,

src/migrations/20250725_01_hwite-add-milestone-table.py

@@ -0,0 +1,61 @@
+"""
+Add milestone table for better milestone tracking
+"""
+
+from yoyo import step
+
+__depends__ = {"20250617_01_c5mrF-task-split"}  # Update to latest migration
+
+steps = [
+    step(
+        """
+         CREATE TABLE IF NOT EXISTS leaderboard.milestones (
+             id SERIAL PRIMARY KEY,
+             leaderboard_id INTEGER NOT NULL REFERENCES leaderboard.leaderboard(id) ON DELETE CASCADE,
+             name TEXT NOT NULL,
+             code TEXT NOT NULL,
+             description TEXT,
+             created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
+             UNIQUE(leaderboard_id, name)
+         )
+         """,
+        "DROP TABLE leaderboard.milestones",
+    ),
+    step("CREATE INDEX ON leaderboard.milestones (leaderboard_id)"),
+    # add alternative ID column that references milestones;
+    # we don't really care about being careful preserving milestone
+    # runs, so we can simply DELETE CASCADE
+    step(
+        """
+        ALTER TABLE leaderboard.runs
+        ADD COLUMN milestone_id INTEGER REFERENCES leaderboard.milestones(id) ON DELETE CASCADE;
+        """,
+        "ALTER TABLE leaderboard.runs DROP COLUMN milestone_id;",
+    ),
+    # as we now have two possible ids, exactly one of them can be NULL
+    step(
+        "ALTER TABLE leaderboard.runs ALTER COLUMN submission_id DROP NOT NULL;",
+        """
+         DELETE FROM leaderboard.runs WHERE submission_id IS NULL;
+         ALTER TABLE leaderboard.runs ALTER COLUMN submission_id SET NOT NULL;
+         """,
+    ),
+    step(
+        """
+        ALTER TABLE leaderboard.runs
+        ADD CONSTRAINT runs_single_parent CHECK (
+            (submission_id IS NOT NULL AND milestone_id IS NULL) OR
+            (submission_id IS NULL AND milestone_id IS NOT NULL)
+        );
+        """,
+        "ALTER TABLE leaderboard.runs DROP CONSTRAINT runs_single_parent;",
+    ),
+    # ensure we have fast indexing for regular submissions
+    step(
+        """
+        CREATE INDEX IF NOT EXISTS runs_submission_id_idx ON leaderboard.runs(submission_id)
+        WHERE submission_id IS NOT NULL;
+        """,
+        "DROP INDEX IF EXISTS leaderboard.runs_submission_id_idx",
+    ),
+]