Platform: Add type annotations for primary key types by model [FC-0117]

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,10 +34,12 @@
 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
 from openedx.core.djangoapps.video_config.transcripts_utils import Transcript, build_components_import_path
+from openedx.core.types import AuthUser as UserType
 from xmodule.contentstore.content import StaticContent
 from xmodule.contentstore.django import contentstore
 from xmodule.exceptions import NotFoundError
@@ -280,7 +283,10 @@ def create_usage(self, def_id) -> UsageKey:
     def create_definition(self, block_type, slug=None) -> DefinitionLocator:
         """ Generate a new definition_id for an XBlock """
         # Note: Split modulestore will detect this temporary ID and create a new definition ID when the XBlock is saved.
-        return DefinitionLocator(block_type, LocalId(block_type))
+        # FIXME: The DefinitionLocator technically only accepts an ObjectId (or a str representing an ObjectId), but
+        # this code relies on passing a LocalId and having it save the LocalId object as its `definition_id`. We should
+        # either change this in the future or update DefinitionLocator to support LocalId-typed definition IDs.
+        return DefinitionLocator(block_type, LocalId(block_type))  # type: ignore[arg-type]
 
 
 @frozen
@@ -312,7 +318,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.
@@ -452,7 +458,7 @@ def _fetch_and_set_upstream_link(
     copied_from_block: str,
     copied_from_version_num: int,
     temp_xblock: XBlock,
-    user: User
+    user: UserType,
 ):
     """
     Fetch and set upstream link for the given xblock which is being pasted. This function handles following cases:
@@ -513,7 +519,7 @@ def _import_xml_node_to_parent(
     # The modulestore we're using
     store,
     # The user who is performing this operation
-    user: User,
+    user: UserType,
     # Hint to use as usage ID (block_id) for the new XBlock
     slug_hint: str | None = None,
     # Content tags applied to the source XBlock(s)
@@ -635,7 +641,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 +706,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 +766,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],
 ):
     """

cms/djangoapps/contentstore/models.py

@@ -485,7 +485,7 @@ def _annotate_query_with_ready_to_sync(cls, query_set: QuerySet["EntityLinkBase"
     @classmethod
     def update_or_create(
         cls,
-        upstream_container_id: int | None,
+        upstream_container_id: Container.ID | None,
         /,
         upstream_container_key: LibraryContainerLocator,
         upstream_context_key: str,

cms/djangoapps/contentstore/utils.py

@@ -2425,7 +2425,7 @@ def _create_or_update_container_link(created: datetime | None, xblock):
     """
     upstream_container_key = LibraryContainerLocator.from_string(xblock.upstream)
     try:
-        lib_component = get_container(upstream_container_key).container_pk
+        lib_component = get_container(upstream_container_key).container_id
     except ObjectDoesNotExist:
         log.error(f"Library component not found for {upstream_container_key}")
         lib_component = None

cms/djangoapps/modulestore_migrator/api/write_api.py

@@ -6,6 +6,7 @@
 from celery.result import AsyncResult
 from opaque_keys.edx.locator import LibraryLocatorV2
 from openedx_content.api import get_collection
+from openedx_content.models_api import LearningPackage
 
 from openedx.core.djangoapps.content_libraries.api import get_library
 from openedx.core.types.user import AuthUser
@@ -64,7 +65,8 @@ def start_bulk_migration_to_library(
     """
     target_library = get_library(target_library_key)
     # get_library ensures that the library is connected to a learning package.
-    target_package_id: int = target_library.learning_package_id  # type: ignore[assignment]
+    assert target_library.learning_package_id is not None
+    target_package_id: LearningPackage.ID = target_library.learning_package_id
 
     sources_pks: list[int] = []
     for source_key in source_key_list:

cms/djangoapps/modulestore_migrator/tasks.py

@@ -120,7 +120,7 @@ class _MigrationContext:
 
     # Fields that remain constant
     previous_block_migrations: dict[UsageKey, data.ModulestoreBlockMigrationResult]
-    target_package_id: int
+    target_package_id: LearningPackage.ID
     target_library_key: LibraryLocatorV2
     source_context_key: SourceContextKey
     content_by_filename: dict[str, int]
@@ -345,21 +345,21 @@ def _import_structure(
         used_component_keys=set(
             LibraryUsageLocatorV2(target_library.key, block_type, block_id)  # type: ignore[abstract]
             for block_type, block_id
-            in content_api.get_components(migration.target.pk).values_list(
+            in content_api.get_components(migration.target.id).values_list(
                 "component_type__name", "local_key"
             )
         ),
         used_container_slugs=set(
             content_api.get_containers(
-                migration.target.pk
+                migration.target.id
             ).values_list("publishable_entity__key", flat=True)
         ),
         previous_block_migrations=(
             get_migration_blocks(source_data.previous_migration.pk)
             if source_data.previous_migration
             else {}
         ),
-        target_package_id=migration.target.pk,
+        target_package_id=migration.target.id,
         target_library_key=target_library.key,
         source_context_key=source_data.source_root_usage_key.course_key,
         content_by_filename=content_by_filename,
@@ -882,12 +882,12 @@ def _migrate_container(
             )
     if container_exists and context.should_skip_strategy:
         return PublishableEntityVersion.objects.get(
-            entity_id=container.container_pk,
+            entity_id=container.container_id,
             version_num=container.draft_version_num,
         ), None
 
     container_publishable_entity_version = content_api.create_next_container_version(
-        container.container_pk,
+        container.container_id,
         title=title,
         entities=[child.entity for child in children],
         created=context.created_at,

cms/pytest.ini

@@ -17,6 +17,10 @@ filterwarnings =
     # ABC deprecation Warning comes from libsass
     ignore:Using or importing the ABCs from 'collections' instead of from 'collections.abc' is deprecated.*:DeprecationWarning:sass
     ignore:'etree' is deprecated. Use 'xml.etree.ElementTree' instead.:DeprecationWarning:wiki
+    # We sometimes use DeprecationWarning to enourage `.id` instead of `.pk` since there are some cases where we cannot
+    # override the type of `.pk`, but we can always customize `.id` (e.g. using TypedBigAutoField). But we don't want to
+    # see those warnings when Django is using `.pk` internally, which is fine.
+    ignore:Use \.id instead:DeprecationWarning:django\..*
 
 norecursedirs = envs
 python_classes =

mypy.ini

@@ -9,6 +9,7 @@ files =
     cms/lib/xblock,
     cms/djangoapps/contentstore/rest_api/v2/views,
     cms/djangoapps/contentstore/xblock_storage_handlers,
+    cms/djangoapps/contentstore/helpers.py,
     cms/djangoapps/modulestore_migrator,
     openedx/core/djangoapps/content/learning_sequences,
     # FIXME: need to solve type issues and add 'search' app here:

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,
@@ -243,7 +244,7 @@ def set_library_block_olx(usage_key: LibraryUsageLocatorV2, new_olx_str: str) ->
             created=now,
         )
         new_component_version = content_api.create_next_component_version(
-            component.pk,
+            component.id,
             title=new_title,
             media_to_replace={
                 'block.xml': new_content.pk,
@@ -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,
     *,
@@ -727,7 +728,7 @@ def send_block_deleted_signal():
     affected_collections = content_api.get_entity_collections(component.learning_package_id, component.key)
     affected_containers = get_containers_contains_item(usage_key)
 
-    content_api.soft_delete_draft(component.pk, deleted_by=user_id)
+    content_api.soft_delete_draft(component.id, deleted_by=user_id)
 
     send_block_deleted_signal()
 
@@ -773,7 +774,7 @@ def restore_library_block(usage_key: LibraryUsageLocatorV2, user_id: int | None
 
     # Set draft version back to the latest available component version id.
     content_api.set_draft_version(
-        component.pk,
+        component.id,
         component.versioning.latest.pk,
         set_by=user_id,
     )
@@ -909,7 +910,7 @@ def add_library_block_static_asset_file(
 
     with transaction.atomic():
         component_version = content_api.create_next_component_version(
-            component.pk,
+            component.id,
             media_to_replace={file_path: file_content},
             created=datetime.now(tz=timezone.utc),  # noqa: UP017
             created_by=user.id if user else None,
@@ -956,8 +957,8 @@ def delete_library_block_static_asset_file(usage_key, file_path, user=None):
     now = datetime.now(tz=timezone.utc)  # noqa: UP017
 
     with transaction.atomic():
-        component_version = content_api.create_next_component_version(  # noqa: F841
-            component.pk,
+        content_api.create_next_component_version(
+            component.id,
             media_to_replace={file_path: None},
             created=now,
             created_by=user.id if user else None,

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

@@ -71,7 +71,7 @@ class ContainerMetadata(PublishableItem):
 
     container_key: LibraryContainerLocator
     container_type_code: str
-    container_pk: int
+    container_id: Container.ID
 
     @classmethod
     def from_container(cls, library_key, container: Container, associated_collections=None):
@@ -99,7 +99,7 @@ def from_container(cls, library_key, container: Container, associated_collection
         return cls(
             container_key=container_key,
             container_type_code=container_key.container_type,
-            container_pk=container.pk,
+            container_id=container.id,
             display_name=draft.title,
             created=container.created,
             modified=draft.created,
@@ -110,7 +110,7 @@ def from_container(cls, library_key, container: Container, associated_collection
             published_by=published_by,
             last_draft_created=last_draft_created,
             last_draft_created_by=last_draft_created_by,
-            has_unpublished_changes=content_api.contains_unpublished_changes(container.pk),
+            has_unpublished_changes=content_api.contains_unpublished_changes(container.id),
             tags_count=tags.get(str(container_key), 0),
             collections=associated_collections or [],
         )
@@ -260,7 +260,7 @@ def _get_containers_with_entities(
     for member in members:
         qs = qs.union(
             content_api.get_containers_with_entity(
-                member.entity.pk,
+                member.entity.id,
                 ignore_pinned=ignore_pinned,
             )
         )
@@ -338,7 +338,7 @@ def create(
                 container=entity,
             ),
             display_name=entity.versioning.draft.title,
-            has_unpublished_changes=content_api.contains_unpublished_changes(entity.pk),
+            has_unpublished_changes=content_api.contains_unpublished_changes(entity.id),
             container=entity,
             component=None,
         )

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

@@ -234,7 +234,7 @@ def send_container_deleted_signal():
         container_key,
         published=False,
     )
-    content_api.soft_delete_draft(container.pk)
+    content_api.soft_delete_draft(container.id)
 
     send_container_deleted_signal()
 
@@ -294,7 +294,7 @@ def restore_container(container_key: LibraryContainerLocator) -> None:
         container.key,
     )
 
-    content_api.set_draft_version(container.pk, container.versioning.latest.pk)
+    content_api.set_draft_version(container.id, container.versioning.latest.pk)
     # Fetch related containers after restore
     affected_containers = get_containers_contains_item(container_key)
     # Get children containers or components to update their index data
@@ -441,7 +441,7 @@ def get_containers_contains_item(key: LibraryUsageLocatorV2 | LibraryContainerLo
     [ 🛑 UNSTABLE ] Get containers that contains the item, that can be a component or another container.
     """
     entity = get_entity_from_key(key)
-    containers = content_api.get_containers_with_entity(entity.pk).select_related("container_type")
+    containers = content_api.get_containers_with_entity(entity.id).select_related("container_type")
     return [ContainerMetadata.from_container(key.lib_key, container) for container in containers]
 
 
@@ -460,7 +460,7 @@ def publish_container_changes(
     learning_package = content_library.learning_package
     assert learning_package
     # The core publishing API is based on draft objects, so find the draft that corresponds to this container:
-    drafts_to_publish = content_api.get_all_drafts(learning_package.id).filter(entity__pk=container.pk)
+    drafts_to_publish = content_api.get_all_drafts(learning_package.id).filter(entity__pk=container.id)
     # Publish the container, which will also auto-publish any unpublished child components:
     publish_log = content_api.publish_from_drafts(
         learning_package.id,

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

@@ -124,7 +124,7 @@ class ContentLibraryMetadata:
     Class that represents the metadata about a content library.
     """
     key: LibraryLocatorV2
-    learning_package_id: int | None
+    learning_package_id: LearningPackage.ID | None
     title: str = ""
     description: str = ""
     num_blocks: int = 0
@@ -410,7 +410,7 @@ def get_library(library_key: LibraryLocatorV2) -> ContentLibraryMetadata:
         license=ref.license,
         created=learning_package.created,
         updated=learning_package.updated,
-        learning_package_id=learning_package.pk,
+        learning_package_id=learning_package.id,
     )
 
 
@@ -494,7 +494,7 @@ def create_library(
         allow_public_learning=ref.allow_public_learning,
         allow_public_read=ref.allow_public_read,
         license=library_license,
-        learning_package_id=ref.learning_package.pk,  # type: ignore[union-attr]
+        learning_package_id=ref.learning_package.id,  # type: ignore[union-attr]
     )