From c34a6db4e70f69445e37df4a1cc6b8abd3b40dfc Mon Sep 17 00:00:00 2001
From: Marcos Prieto <marcospri@gmail.com>
Date: Mon, 2 Sep 2024 18:18:33 +0200
Subject: [PATCH] Model changes for AssignmentRoster

---
 lms/models/__init__.py      |  2 +-
 lms/models/course_roster.py | 30 -----------------
 lms/models/roster.py        | 67 +++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 31 deletions(-)
 delete mode 100644 lms/models/course_roster.py
 create mode 100644 lms/models/roster.py

diff --git a/lms/models/__init__.py b/lms/models/__init__.py
index 242cbe5a4d..bd67960c53 100644
--- a/lms/models/__init__.py
+++ b/lms/models/__init__.py
@@ -4,7 +4,6 @@
 from lms.models.assignment_grouping import AssignmentGrouping
 from lms.models.assignment_membership import AssignmentMembership
 from lms.models.course_groups_exported_from_h import CourseGroupsExportedFromH
-from lms.models.course_roster import CourseRoster
 from lms.models.dashboard_admin import DashboardAdmin
 from lms.models.event import Event, EventData, EventType, EventUser
 from lms.models.exceptions import ReusedConsumerKey
@@ -37,6 +36,7 @@
 from lms.models.oauth2_token import OAuth2Token
 from lms.models.organization import Organization
 from lms.models.organization_usage import OrganizationUsageReport
+from lms.models.roster import AssignmentRoster, CourseRoster
 from lms.models.rsa_key import RSAKey
 from lms.models.task_done import TaskDone
 from lms.models.user import User
diff --git a/lms/models/course_roster.py b/lms/models/course_roster.py
deleted file mode 100644
index b39c2acb85..0000000000
--- a/lms/models/course_roster.py
+++ /dev/null
@@ -1,30 +0,0 @@
-from sqlalchemy import ForeignKey, UniqueConstraint
-from sqlalchemy.orm import Mapped, mapped_column, relationship
-
-from lms.db import Base
-from lms.models._mixins import CreatedUpdatedMixin
-
-
-class CourseRoster(Base, CreatedUpdatedMixin):
-    __tablename__ = "course_roster"
-
-    id: Mapped[int] = mapped_column(autoincrement=True, primary_key=True)
-
-    lms_course_id: Mapped[int] = mapped_column(
-        ForeignKey("lms_course.id", ondelete="cascade")
-    )
-    lms_course = relationship("LMSCourse")
-
-    lms_user_id: Mapped[int] = mapped_column(
-        ForeignKey("lms_user.id", ondelete="cascade")
-    )
-    lms_user = relationship("LMSUser")
-
-    lti_role_id: Mapped[int] = mapped_column(
-        ForeignKey("lti_role.id", ondelete="cascade")
-    )
-    lti_role = relationship("LTIRole")
-
-    active: Mapped[bool] = mapped_column()
-
-    __table_args__ = (UniqueConstraint("lms_course_id", "lms_user_id", "lti_role_id"),)
diff --git a/lms/models/roster.py b/lms/models/roster.py
new file mode 100644
index 0000000000..2034ac79ea
--- /dev/null
+++ b/lms/models/roster.py
@@ -0,0 +1,67 @@
+from sqlalchemy import ForeignKey, UniqueConstraint
+from sqlalchemy.orm import Mapped, mapped_column, relationship
+
+from lms.db import Base
+from lms.models._mixins import CreatedUpdatedMixin
+
+
+class CourseRoster(Base, CreatedUpdatedMixin):
+    """
+    Store roster information for courses.
+
+    Stores a row per unique: (course, user and role)
+    """
+
+    __tablename__ = "course_roster"
+
+    id: Mapped[int] = mapped_column(autoincrement=True, primary_key=True)
+
+    lms_course_id: Mapped[int] = mapped_column(
+        ForeignKey("lms_course.id", ondelete="cascade")
+    )
+    lms_course = relationship("LMSCourse")
+
+    lms_user_id: Mapped[int] = mapped_column(
+        ForeignKey("lms_user.id", ondelete="cascade")
+    )
+    lms_user = relationship("LMSUser")
+
+    lti_role_id: Mapped[int] = mapped_column(
+        ForeignKey("lti_role.id", ondelete="cascade")
+    )
+    lti_role = relationship("LTIRole")
+
+    active: Mapped[bool] = mapped_column()
+
+    __table_args__ = (UniqueConstraint("lms_course_id", "lms_user_id", "lti_role_id"),)
+
+
+class AssignmentRoster(Base, CreatedUpdatedMixin):
+    """
+    Store roster information for assignments.
+
+    Stores a row per unique: (assignment, user and role)
+    """
+
+    __tablename__ = "assignment_roster"
+
+    id: Mapped[int] = mapped_column(autoincrement=True, primary_key=True)
+
+    assignment_id: Mapped[int] = mapped_column(
+        ForeignKey("assignment.id", ondelete="cascade")
+    )
+    assignment = relationship("Assignment")
+
+    lms_user_id: Mapped[int] = mapped_column(
+        ForeignKey("lms_user.id", ondelete="cascade")
+    )
+    lms_user = relationship("LMSUser")
+
+    lti_role_id: Mapped[int] = mapped_column(
+        ForeignKey("lti_role.id", ondelete="cascade")
+    )
+    lti_role = relationship("LTIRole")
+
+    active: Mapped[bool] = mapped_column()
+
+    __table_args__ = (UniqueConstraint("assignment_id", "lms_user_id", "lti_role_id"),)