feat: fully typed primary keys for StagedContent model

Author: bradenmacdonald

cms/djangoapps/contentstore/helpers.py

@@ -6,6 +6,7 @@
 Only Studio-specfic helper functions should be added here.
 Platform-wide Python APIs should be added to an appropriate api.py file instead.
 """
+
 from __future__ import annotations
 
 import json
@@ -33,6 +34,7 @@
 from cms.djangoapps.models.settings.course_grading import CourseGradingModel
 from cms.lib.xblock.upstream_sync import UpstreamLink, UpstreamLinkException
 from cms.lib.xblock.upstream_sync_block import fetch_customizable_fields_from_block
+from openedx.core.djangoapps.content_staging.api import StagedContentID
 from openedx.core.djangoapps.content_staging.data import LIBRARY_SYNC_PURPOSE
 from openedx.core.djangoapps.content_tagging.types import TagValuesByObjectIdDict
 from openedx.core.djangoapps.site_configuration import helpers as configuration_helpers
@@ -312,7 +314,7 @@ def _rewrite_static_asset_references(downstream_xblock: XBlock, substitutions: d
 
 
 def _insert_static_files_into_downstream_xblock(
-    downstream_xblock: XBlock, staged_content_id: int, request
+    downstream_xblock: XBlock, staged_content_id: StagedContentID, request
 ) -> StaticFileNotices:
     """
     Gets static files from staged content, and inserts them into the downstream XBlock.
@@ -635,7 +637,7 @@ def _import_xml_node_to_parent(
 
 def _import_files_into_course(
     course_key: CourseKey,
-    staged_content_id: int,
+    staged_content_id: StagedContentID,
     static_files: list[content_staging_api.StagedContentFileData],
     usage_key: UsageKey,
 ) -> tuple[StaticFileNotices, dict[str, str]]:
@@ -700,7 +702,7 @@ def _import_files_into_course(
 
 def _import_file_into_course(
     course_key: CourseKey,
-    staged_content_id: int,
+    staged_content_id: StagedContentID,
     file_data_obj: content_staging_api.StagedContentFileData,
     usage_key: UsageKey,
 ) -> tuple[bool | None, dict]:
@@ -760,7 +762,7 @@ def _import_file_into_course(
 
 def _import_transcripts(
     block: XBlock,
-    staged_content_id: int,
+    staged_content_id: StagedContentID,
     static_files: list[content_staging_api.StagedContentFileData],
 ):
     """

openedx/core/djangoapps/content_libraries/api/blocks.py

@@ -41,6 +41,7 @@
 )
 from xblock.core import XBlock
 
+from openedx.core.djangoapps.content_staging.data import StagedContentID
 from openedx.core.djangoapps.xblock.api import (
     get_component_from_usage_key,
     get_xblock_app_config,
@@ -382,7 +383,7 @@ def _import_staged_block(
     library_key: LibraryLocatorV2,
     source_context_key: LearningContextKey,
     user,
-    staged_content_id: int,
+    staged_content_id: StagedContentID,
     staged_content_files: list[StagedContentFileData],
     now: datetime,
 ) -> LibraryXBlockMetadata:
@@ -517,7 +518,7 @@ def _import_staged_block_as_container(
     library_key: LibraryLocatorV2,
     source_context_key: LearningContextKey,
     user,
-    staged_content_id: int,
+    staged_content_id: StagedContentID,
     staged_content_files: list[StagedContentFileData],
     now: datetime,
     *,

openedx/core/djangoapps/content_staging/api.py

@@ -19,7 +19,14 @@
 from xmodule.contentstore.content import StaticContent
 from xmodule.contentstore.django import contentstore
 
-from .data import CLIPBOARD_PURPOSE, StagedContentData, StagedContentFileData, StagedContentStatus, UserClipboardData
+from .data import (
+    CLIPBOARD_PURPOSE,
+    StagedContentData,
+    StagedContentFileData,
+    StagedContentID,
+    StagedContentStatus,
+    UserClipboardData,
+)
 from .models import StagedContent as _StagedContent
 from .models import StagedContentFile as _StagedContentFile
 from .models import UserClipboard as _UserClipboard
@@ -315,26 +322,26 @@ def _user_clipboard_model_to_data(clipboard: _UserClipboard) -> UserClipboardDat
     )
 
 
-def get_staged_content_olx(staged_content_id: int) -> str | None:
+def get_staged_content_olx(staged_content_id: StagedContentID) -> str | None:
     """
     Get the OLX (as a string) for the given StagedContent.
 
     Does not check permissions!
     """
     try:
-        sc = _StagedContent.objects.get(pk=staged_content_id)
+        sc = _StagedContent.objects.get(id=staged_content_id)
         return sc.olx
     except _StagedContent.DoesNotExist:
         return None
 
 
-def get_staged_content_static_files(staged_content_id: int) -> list[StagedContentFileData]:
+def get_staged_content_static_files(staged_content_id: StagedContentID) -> list[StagedContentFileData]:
     """
     Get the filenames and metadata for any static files used by the given staged content.
 
     Does not check permissions!
     """
-    sc = _StagedContent.objects.get(pk=staged_content_id)
+    sc = _StagedContent.objects.get(id=staged_content_id)
 
     def str_to_key(source_key_str: str):
         if not source_key_str:
@@ -360,13 +367,13 @@ def str_to_key(source_key_str: str):
     ]
 
 
-def get_staged_content_static_file_data(staged_content_id: int, filename: str) -> bytes | None:
+def get_staged_content_static_file_data(staged_content_id: StagedContentID, filename: str) -> bytes | None:
     """
     Get the data for the static asset associated with the given staged content.
 
     Does not check permissions!
     """
-    sc = _StagedContent.objects.get(pk=staged_content_id)
+    sc = _StagedContent.objects.get(id=staged_content_id)
     file_data_obj = sc.files.filter(filename=filename).first()
     if file_data_obj:
         return file_data_obj.data_file.open().read()

openedx/core/djangoapps/content_staging/data.py

@@ -4,6 +4,7 @@
 from __future__ import annotations
 
 from datetime import datetime
+from typing import NewType
 
 from attrs import field, frozen, validators
 from django.db.models import TextChoices
@@ -12,6 +13,9 @@
 
 from openedx.core.djangoapps.content_tagging.api import TagValuesByObjectIdDict
 
+# Typed primary key for the StagedContent model / StagedContentData
+StagedContentID = NewType("StagedContentID", int)
+
 
 class StagedContentStatus(TextChoices):
     """ The status of this staged content. """
@@ -42,7 +46,7 @@ class StagedContentData:
 
     (OLX content that isn't part of any course at the moment)
     """
-    id: int = field(validator=validators.instance_of(int))
+    id: StagedContentID = field(validator=validators.instance_of(int))  # type: ignore[assignment]
     user_id: int = field(validator=validators.instance_of(int))
     created: datetime = field(validator=validators.instance_of(datetime))
     purpose: str = field(validator=validators.instance_of(str))

openedx/core/djangoapps/content_staging/models.py

@@ -1,6 +1,7 @@
 """
 Models for content staging (and clipboard)
 """
+
 from __future__ import annotations
 
 import logging
@@ -11,11 +12,11 @@
 from django.utils.translation import gettext_lazy as _
 from opaque_keys import InvalidKeyError
 from opaque_keys.edx.keys import ContainerKey, LearningContextKey, UsageKey
-from openedx_django_lib.fields import MultiCollationTextField, case_insensitive_char_field
+from openedx_django_lib.fields import MultiCollationTextField, TypedAutoField, case_insensitive_char_field
 
 from openedx.core.djangoapps.content.course_overviews.api import get_course_overview_or_none
 
-from .data import CLIPBOARD_PURPOSE, StagedContentStatus
+from .data import CLIPBOARD_PURPOSE, StagedContentID, StagedContentStatus
 
 log = logging.getLogger(__name__)
 
@@ -42,7 +43,12 @@ class StagedContent(models.Model):
     class Meta:
         verbose_name_plural = _("Staged Content")
 
-    id = models.AutoField(primary_key=True)
+    type ID = StagedContentID
+
+    class IDField(TypedAutoField[ID]):
+        pass
+
+    id = IDField(primary_key=True)
     # The user that created and owns this staged content. Only this user can read it.
     user = models.ForeignKey(User, null=False, on_delete=models.CASCADE)
     created = models.DateTimeField(null=False, auto_now_add=True)

openedx/core/djangoapps/content_staging/tasks.py

@@ -17,7 +17,7 @@
 
 @shared_task(base=LoggedTask)
 @set_code_owner_attribute
-def delete_expired_clipboards(staged_content_ids: list[int]):
+def delete_expired_clipboards(staged_content_ids: list[StagedContent.ID]):
     """
     A Celery task to delete StagedContent clipboard entries that are no longer
     relevant.

openedx/core/djangoapps/content_staging/tests/test_clipboard.py

@@ -3,6 +3,7 @@
 Tests for the clipboard functionality
 """
 from textwrap import dedent
+from typing import cast
 from xml.etree import ElementTree
 
 from rest_framework.test import APIClient
@@ -294,7 +295,7 @@ def test_copy_static_assets(self):
         # Validate the response:
         assert response.status_code == 200
         response_data = response.json()
-        staged_content_id = response_data["content"]["id"]
+        staged_content_id = cast(python_api.StagedContentID, response_data["content"]["id"])
         olx_str = python_api.get_staged_content_olx(staged_content_id)
         assert '<img src="/static/foo_bar.jpg" />' in olx_str
         static_assets = python_api.get_staged_content_static_files(staged_content_id)