Fetch assignment rosters

Same process as for course rosters, schedule fetch rosters for
assignments with recent launches and no recent roster in the DB.

lms/models/grouping.py

@@ -196,6 +196,7 @@ class Course(Grouping):
     lms_course: Mapped[LMSCourse] = sa.orm.relationship(
         LMSCourse,
         primaryjoin="Grouping.authority_provided_id == foreign(LMSCourse.h_authority_provided_id)",
+        backref=sa.orm.backref("course"),
     )
 
     def set_group_sets(self, group_sets: list[dict]):

lms/services/roster.py

@@ -4,6 +4,8 @@
 
 from lms.models import (
     ApplicationInstance,
+    Assignment,
+    AssignmentRoster,
     CourseRoster,
     LMSCourse,
     LMSCourseApplicationInstance,
@@ -33,19 +35,11 @@ def __init__(
         self._lti_role_service = lti_role_service
         self._h_authority = h_authority
 
-    def fetch_roster(self, lms_course: LMSCourse) -> None:
+    def fetch_course_roster(self, lms_course: LMSCourse) -> None:
         assert (
             lms_course.lti_context_memberships_url
         ), "Trying fetch roster for course without service URL."
-
-        lti_registration = self._db.scalars(
-            select(LTIRegistration)
-            .join(ApplicationInstance)
-            .where(LMSCourseApplicationInstance.lms_course_id == lms_course.id)
-            .join(LMSCourseApplicationInstance)
-            .order_by(LTIRegistration.updated.desc())
-        ).first()
-        assert lti_registration, "No LTI registration found for LMSCourse."
+        lti_registration = self._get_lti_registration(lms_course)
 
         roster = self._lti_names_roles_service.get_context_memberships(
             lti_registration, lms_course.lti_context_memberships_url
@@ -97,6 +91,75 @@ def fetch_roster(self, lms_course: LMSCourse) -> None:
             update_columns=["active", "updated"],
         )
 
+    def fetch_assignment_roster(self, assignment: Assignment) -> None:
+        assert (
+            assignment.lti_v13_resource_link_id
+        ), "Trying fetch roster for an assignment without LTI1.3 ID."
+
+        lms_course = self._db.scalars(
+            select(LMSCourse).where(
+                LMSCourse.h_authority_provided_id
+                == assignment.course.authority_provided_id
+            )
+        ).one()
+
+        assert (
+            lms_course.lti_context_memberships_url
+        ), "Trying fetch roster for course without service URL."
+        lti_registration = self._get_lti_registration(lms_course)
+
+        roster = self._lti_names_roles_service.get_context_memberships(
+            lti_registration,
+            lms_course.lti_context_memberships_url,
+            resource_link_id=assignment.lti_v13_resource_link_id,
+        )
+
+        # Insert any users we might be missing in the DB
+        lms_users_by_lti_user_id = {
+            u.lti_user_id: u
+            for u in self._get_roster_users(
+                roster, lms_course.tool_consumer_instance_guid
+            )
+        }
+        # Also insert any roles we might be missing
+        lti_roles_by_value: dict[str, LTIRole] = {
+            r.value: r for r in self._get_roster_roles(roster)
+        }
+
+        # Make sure any new rows have IDs
+        self._db.flush()
+
+        roster_upsert_elements = []
+
+        for member in roster:
+            lti_user_id = member.get("lti11_legacy_user_id") or member["user_id"]
+            # Now, for every user + role, insert a row  in the roster table
+            for role in member["roles"]:
+                roster_upsert_elements.append(
+                    {
+                        "assignment_id": assignment.id,
+                        "lms_user_id": lms_users_by_lti_user_id[lti_user_id].id,
+                        "lti_role_id": lti_roles_by_value[role].id,
+                        "active": member["status"] == "Active",
+                    }
+                )
+        # We'll first mark everyone as non-Active.
+        # We keep a record of who belonged to a course even if they are no longer present.
+        self._db.execute(
+            update(AssignmentRoster)
+            .where(AssignmentRoster.assignment_id == assignment.id)
+            .values(active=False)
+        )
+
+        # Insert and update roster rows.
+        bulk_upsert(
+            self._db,
+            AssignmentRoster,
+            values=roster_upsert_elements,
+            index_elements=["assignment_id", "lms_user_id", "lti_role_id"],
+            update_columns=["active", "updated"],
+        )
+
     def _get_roster_users(self, roster, tool_consumer_instance_guid):
         values = []
         for member in roster:
@@ -134,6 +197,17 @@ def _get_roster_roles(self, roster) -> list[LTIRole]:
         roles = {role for member in roster for role in member["roles"]}
         return self._lti_role_service.get_roles(list(roles))
 
+    def _get_lti_registration(self, lms_course) -> LTIRegistration:
+        lti_registration = self._db.scalars(
+            select(LTIRegistration)
+            .join(ApplicationInstance)
+            .where(LMSCourseApplicationInstance.lms_course_id == lms_course.id)
+            .join(LMSCourseApplicationInstance)
+            .order_by(LTIRegistration.updated.desc())
+        ).first()
+        assert lti_registration, "No LTI registration found for LMSCourse."
+        return lti_registration
+
 
 def factory(_context, request):
     return RosterService(

lms/tasks/roster.py

@@ -4,22 +4,37 @@
 
 from sqlalchemy import exists, select
 
-from lms.models import Course, CourseRoster, Event, LMSCourse
+from lms.models import (
+    Assignment,
+    AssignmentRoster,
+    Course,
+    CourseRoster,
+    Event,
+    LMSCourse,
+)
 from lms.services.roster import RosterService
 from lms.tasks.celery import app
 
-COURSE_LAUNCHED_WINDOW = timedelta(hours=24)
-"""How recent we need to have seen a launch from a course before we stop fetching rosters for it."""
+LAUNCHED_WINDOW = timedelta(hours=24)
+"""How recent we need to have seen a launch from a course/assignment before we stop fetching rosters for it."""
 
 ROSTER_REFRESH_WINDOW = timedelta(hours=24 * 7)
-"""How frequenly should we fetch roster for the same course"""
+"""How frequenly should we fetch roster for the same course/assignment"""
 
-ROSTER_COURSE_LIMIT = 5
+ROSTER_LIMIT = 5
 """How many roster should we fetch per executing of the schedule task."""
 
 
 @app.task()
 def schedule_fetching_rosters() -> None:
+    # Schedule fetching course rosters
+    schedule_fetching_course_rosters()
+
+    # Schedule fetching assignment rosters
+    schedule_fetching_assignment_rosters()
+
+
+def schedule_fetching_course_rosters() -> None:
     """Schedule fetching course rosters based on their last lunches and the most recent roster fetch."""
 
     # We use the python version (and not func.now()) for easier mocking during tests
@@ -38,7 +53,7 @@ def schedule_fetching_rosters() -> None:
         select(Event)
         .join(Course, Event.course_id == Course.id)
         .where(
-            Event.timestamp >= now - COURSE_LAUNCHED_WINDOW,
+            Event.timestamp >= now - LAUNCHED_WINDOW,
             Course.authority_provided_id == LMSCourse.h_authority_provided_id,
         )
     )
@@ -53,11 +68,58 @@ def schedule_fetching_rosters() -> None:
                     no_recent_roster_clause,
                     recent_launches_cluase,
                 )
-                # Schedule only a few roster per call to this method
-                .limit(ROSTER_COURSE_LIMIT)
+                # Schedule only a few rosters per call to this method
+                .limit(ROSTER_LIMIT)
             )
             for lms_course_id in request.db.scalars(query).all():
-                fetch_roster.delay(lms_course_id=lms_course_id)
+                fetch_course_roster.delay(lms_course_id=lms_course_id)
+
+
+def schedule_fetching_assignment_rosters() -> None:
+    """Schedule fetching assignment rosters based on their last lunches and the most recent roster fetch."""
+
+    # We use the python version (and not func.now()) for easier mocking during tests
+    now = datetime.now()
+
+    # Only fetch roster for assignments that don't have recent roster information
+    no_recent_roster_clause = ~exists(
+        select(AssignmentRoster).where(
+            AssignmentRoster.assignment_id == Assignment.id,
+            AssignmentRoster.updated >= now - ROSTER_REFRESH_WINDOW,
+        )
+    )
+
+    # Only fetch rosters for assignments that have been recently launched
+    recent_launches_cluase = exists(
+        select(Event)
+        .join(Assignment, Event.assignment_id == Assignment.id)
+        .where(
+            Event.timestamp >= now - LAUNCHED_WINDOW,
+        )
+    )
+
+    with app.request_context() as request:
+        with request.tm:
+            query = (
+                select(Assignment.id)
+                .join(Course)
+                .join(
+                    LMSCourse,
+                    LMSCourse.h_authority_provided_id == Course.authority_provided_id,
+                )
+                .where(
+                    # Assignments that belongs to courses for which we have a LTIA membership service URL
+                    LMSCourse.lti_context_memberships_url.is_not(None),
+                    # Assignments for which we have stored the LTI1.3 ID
+                    Assignment.lti_v13_resource_link_id.is_not(None),
+                    no_recent_roster_clause,
+                    recent_launches_cluase,
+                )
+                # Schedule only a few roster per call to this method
+                .limit(ROSTER_LIMIT)
+            )
+            for assignment_id in request.db.scalars(query).all():
+                fetch_assignment_roster.delay(assignment_id=assignment_id)
 
 
 @app.task(
@@ -67,10 +129,26 @@ def schedule_fetching_rosters() -> None:
     retry_backoff=3600,
     retry_backoff_max=7200,
 )
-def fetch_roster(*, lms_course_id) -> None:
+def fetch_course_roster(*, lms_course_id) -> None:
     """Fetch the roster for one course."""
     with app.request_context() as request:
         roster_service = request.find_service(RosterService)
         with request.tm:
             lms_course = request.db.get(LMSCourse, lms_course_id)
-            roster_service.fetch_roster(lms_course)
+            roster_service.fetch_course_roster(lms_course)
+
+
+@app.task(
+    acks_late=True,
+    autoretry_for=(Exception,),
+    max_retries=2,
+    retry_backoff=3600,
+    retry_backoff_max=7200,
+)
+def fetch_assignment_roster(*, assignment_id) -> None:
+    """Fetch the roster for one course."""
+    with app.request_context() as request:
+        roster_service = request.find_service(RosterService)
+        with request.tm:
+            assignment = request.db.get(Assignment, assignment_id)
+            roster_service.fetch_assignment_roster(assignment)

tests/factories/__init__.py

@@ -18,7 +18,6 @@
     TOOL_CONSUMER_INSTANCE_GUID,
     USER_ID,
 )
-from tests.factories.course_roster import CourseRoster
 from tests.factories.dashboard_admin import DashboardAdmin
 from tests.factories.event import Event
 from tests.factories.file import File
@@ -37,6 +36,7 @@
 from tests.factories.oauth2_token import OAuth2Token
 from tests.factories.organization import Organization
 from tests.factories.organization_usage import OrganizationUsageReport
+from tests.factories.roster import AssignmentRoster, CourseRoster
 from tests.factories.rsa_key import RSAKey
 from tests.factories.task_done import TaskDone
 from tests.factories.user import User

tests/factories/course_roster.py → tests/factories/roster.py

@@ -4,3 +4,6 @@
 from lms import models
 
 CourseRoster = make_factory(models.CourseRoster, FACTORY_CLASS=SQLAlchemyModelFactory)
+AssignmentRoster = make_factory(
+    models.AssignmentRoster, FACTORY_CLASS=SQLAlchemyModelFactory
+)