feat: synchronize restricted courses from discovery

enterprise_catalog/apps/catalog/admin.py

@@ -133,13 +133,19 @@ class RestrictedCourseMetadataAdmin(UnchangeableMixin):
         description='Catalog Query'
     )
     def get_catalog_query_for_list(self, obj):
+        if not obj.catalog_query:
+            return None
+
         link = reverse("admin:catalog_catalogquery_change", args=[obj.catalog_query.id])
         return format_html('<a href="{}">{}</a>', link, obj.catalog_query.short_str_for_listings())
 
     @admin.display(
         description='Catalog Query'
     )
     def get_catalog_query(self, obj):
+        if not obj.catalog_query:
+            return None
+
         link = reverse("admin:catalog_catalogquery_change", args=[obj.catalog_query.id])
         return format_html('<a href="{}">{}</a>', link, obj.catalog_query.pretty_print_content_filter())
 

enterprise_catalog/apps/catalog/constants.py

@@ -116,6 +116,8 @@
 LATE_ENROLLMENT_THRESHOLD_DAYS = 30
 
 RESTRICTED_RUNS_ALLOWED_KEY = 'restricted_runs_allowed'
+COURSE_RUN_RESTRICTION_TYPE_KEY = 'restriction_type'
+QUERY_FOR_RESTRICTED_RUNS = {'include_restricted': 'custom-b2b-enterprise'}
 
 AGGREGATION_KEY_PREFIX = 'course:'
 

enterprise_catalog/apps/catalog/models.py

@@ -1,4 +1,5 @@
 import collections
+import copy
 import json
 from logging import getLogger
 from uuid import uuid4
@@ -20,7 +21,7 @@
     get_most_recent_modified_time,
     update_query_parameters,
 )
-from enterprise_catalog.apps.api_client.discovery import CatalogQueryMetadata
+from enterprise_catalog.apps.api_client.discovery import CatalogQueryMetadata, DiscoveryApiClient
 from enterprise_catalog.apps.api_client.enterprise_cache import (
     EnterpriseCustomerDetails,
 )
@@ -32,8 +33,10 @@
     CONTENT_TYPE_CHOICES,
     COURSE,
     COURSE_RUN,
+    COURSE_RUN_RESTRICTION_TYPE_KEY,
     EXEC_ED_2U_COURSE_TYPE,
     EXEC_ED_2U_ENTITLEMENT_MODE,
+    QUERY_FOR_RESTRICTED_RUNS,
     PROGRAM,
     RESTRICTED_RUNS_ALLOWED_KEY,
     json_serialized_course_modes,
@@ -818,7 +821,58 @@ def __str__(self):
         """
         Return human-readable string representation.
         """
-        return f"<{self.__class__.__name__} for '{self.content_key}' and CatalogQuery ({self.catalog_query.id})>"
+        catalog_query_id = self.catalog_query.id if self.catalog_query else None
+        return f"<{self.__class__.__name__} for '{self.content_key}' and CatalogQuery ({catalog_query_id})>"
+
+    @classmethod
+    def _store_record(cls, course_metadata_dict, catalog_query=None):
+        """
+        Given a course metadata dictionary, stores a corresponding
+        ``RestrictedContentMetadata`` record. Raises if the content key
+        is not of type 'course', or if a corresponding unrestricted parent
+        record cannot be found.
+        """
+        content_type = course_metadata_dict.get('content_type')
+        if content_type != COURSE:
+            raise Exception('Can only store RestrictedContentMetadata with content type of course')
+
+        course_key = course_metadata_dict['key']
+        parent_record = ContentMetadata.objects.get(content_key=course_key)
+        record, _ = cls.objects.update_or_create(
+            content_key=course_key,
+            content_uuid=course_metadata_dict['uuid'],
+            content_type=COURSE,
+            unrestricted_parent=parent_record,
+            catalog_query=catalog_query,
+            defaults={
+                '_json_metadata': course_metadata_dict,
+            },
+        )
+        return record
+
+    @classmethod
+    def store_canonical_record(cls, course_metadata_dict):
+        return cls._store_record(course_metadata_dict)
+
+    @classmethod
+    def store_record_with_query(cls, course_metadata_dict, catalog_query):
+        filtered_metadata = cls.filter_restricted_runs(course_metadata_dict, catalog_query)
+        return cls._store_record(filtered_metadata, catalog_query)
+
+    @classmethod
+    def filter_restricted_runs(cls, course_metadata_dict, catalog_query):
+        """
+        Returns a copy of `course_metadata_dict` whose course_runs list
+        contains only unrestricted runs and restricted runs that are allowed
+        by the provided ``catalog_query``.
+        """
+        filtered_metadata = copy.deepcopy(course_metadata_dict)
+        allowed_runs = catalog_query.restricted_runs_allowed.get(course_metadata_dict['key'], [])
+        filtered_metadata['course_runs'] = [
+            run for run in filtered_metadata['course_runs']
+            if run[COURSE_RUN_RESTRICTION_TYPE_KEY] is None or run['key'] in allowed_runs
+        ]
+        return filtered_metadata
 
 
 class RestrictedRunAllowedForRestrictedCourse(TimeStampedModel):
@@ -1180,7 +1234,7 @@ def associate_content_metadata_with_query(metadata, catalog_query, dry_run=False
     metadata_list = create_content_metadata(metadata, catalog_query, dry_run)
     # Stop gap if the new metadata list is extremely different from the current one
     if _check_content_association_threshold(catalog_query, metadata_list):
-        return catalog_query.contentmetadata_set.values_list('content_key', flat=True)
+        return list(catalog_query.contentmetadata_set.values_list('content_key', flat=True))
     # Setting `clear=True` will remove all prior relationships between
     # the CatalogQuery's associated ContentMetadata objects
     # before setting all new relationships from `metadata_list`.
@@ -1307,30 +1361,65 @@ def update_contentmetadata_from_discovery(catalog_query, dry_run=False):
         LOGGER.exception(f'update_contentmetadata_from_discovery failed {catalog_query}')
         raise exc
 
+    if not metadata:
+        return []
+
     # associate content metadata with a catalog query only when we get valid results
     # back from the discovery service. if metadata is `None`, an error occurred while
     # calling discovery and we should not proceed with the below association logic.
-    if metadata:
-        metadata_content_keys = [get_content_key(entry) for entry in metadata]
-        LOGGER.info(
-            'Retrieved %d content items (%d unique) from course-discovery for catalog query %s',
-            len(metadata_content_keys),
-            len(set(metadata_content_keys)),
-            catalog_query,
-        )
+    metadata_content_keys = [get_content_key(entry) for entry in metadata]
+    LOGGER.info(
+        'Retrieved %d content items (%d unique) from course-discovery for catalog query %s',
+        len(metadata_content_keys),
+        len(set(metadata_content_keys)),
+        catalog_query,
+    )
 
-        associated_content_keys = associate_content_metadata_with_query(metadata, catalog_query, dry_run)
-        LOGGER.info(
-            'Associated %d content items (%d unique) with catalog query %s',
-            len(associated_content_keys),
-            len(set(associated_content_keys)),
-            catalog_query,
-        )
+    associated_content_keys = associate_content_metadata_with_query(metadata, catalog_query, dry_run)
+    LOGGER.info(
+        'Associated %d content items (%d unique) with catalog query %s',
+        len(associated_content_keys),
+        len(set(associated_content_keys)),
+        catalog_query,
+    )
+
+    restricted_content_keys = synchronize_restricted_content(catalog_query, dry_run=dry_run)
+    return associated_content_keys + restricted_content_keys
 
-        return associated_content_keys
 
-    return []
+def synchronize_restricted_content(catalog_query, dry_run=False):
+    """
+    Fetch and assoicate any permitted restricted couress for the given catalog_query.
+    """
+    if not getattr(settings, 'SHOULD_FETCH_RESTRICTED_COURSE_RUNS', False):
+        return []
+
+    if not catalog_query.restricted_runs_allowed:
+        return []
+
+    restricted_course_keys = list(catalog_query.restricted_runs_allowed)
+    content_filter = {
+        'content_type': 'course',
+        'key': restricted_course_keys,
+    }
+    discovery_client = DiscoveryApiClient()
+    course_payload = discovery_client.retrieve_metadata_for_content_filter(
+        content_filter, QUERY_FOR_RESTRICTED_RUNS,
+    )
+
+    restricted_course_keys = []
+    for course_dict in course_payload:
+        LOGGER.info('Storing restricted course %s for catalog_query %s', course_dict.get('key'), catalog_query.id)
+        if dry_run:
+            continue
+
+        RestrictedCourseMetadata.store_canonical_record(course_dict)
+        restricted_course_record = RestrictedCourseMetadata.store_record_with_query(
+            course_dict, catalog_query,
+        )
+        restricted_course_keys.append(restricted_course_record.content_key)
 
+    return restricted_course_keys
 
 class CatalogUpdateCommandConfig(ConfigurationModel):
     """

enterprise_catalog/settings/base.py

@@ -420,6 +420,10 @@
     DEFAULT_COURSE_FIELDS_TO_PLUCK_FROM_SEARCH_ALL,
 )
 
+# Whether to fetch restricted course runs from the course-discovery
+# /api/v1/courses endpoint
+SHOULD_FETCH_RESTRICTED_COURSE_RUNS = False
+
 # Set up system-to-feature roles mapping for edx-rbac
 SYSTEM_TO_FEATURE_ROLE_MAPPING = {
     # The enterprise catalog admin role is for users who need to perform state altering requests on catalogs