feat: synchronize restricted courses from discovery

ENT-9570

enterprise_catalog/apps/catalog/admin.py

@@ -158,13 +158,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/algolia_utils.py

@@ -10,7 +10,6 @@
 from django.utils.translation import gettext as _
 from pytz import UTC
 
-from enterprise_catalog.apps.api.v1.utils import is_course_run_active
 from enterprise_catalog.apps.api_client.algolia import AlgoliaSearchClient
 from enterprise_catalog.apps.api_client.constants import (
     COURSE_REVIEW_BASE_AVG_REVIEW_SCORE,
@@ -28,13 +27,17 @@
     PROGRAM_TYPES_MAP,
     VIDEO,
 )
+from enterprise_catalog.apps.catalog.content_metadata_utils import (
+    get_course_first_paid_enrollable_seat_price,
+    get_course_run_by_uuid,
+    is_course_run_active,
+)
 from enterprise_catalog.apps.catalog.models import ContentMetadata
 from enterprise_catalog.apps.catalog.serializers import (
     NormalizedContentMetadataSerializer,
 )
 from enterprise_catalog.apps.catalog.utils import (
     batch_by_pk,
-    get_course_run_by_uuid,
     localized_utcnow,
     to_timestamp,
 )
@@ -1286,33 +1289,6 @@ def _get_course_run_enroll_start_date_timestamp(normalized_content_metadata):
     return to_timestamp(enroll_start_date)
 
 
-def get_course_first_paid_enrollable_seat_price(course):
-    """
-    Gets the appropriate image to use for course cards.
-
-    Arguments:
-        course (dict): a dictionary representing a course
-
-    Returns:
-        str: the url for the course card image
-    """
-    # Use advertised course run.
-    # If that fails use one of the other active course runs. (The latter is what Discovery does)
-    advertised_course_run = get_course_run_by_uuid(course, course.get('advertised_course_run_uuid'))
-    if advertised_course_run and advertised_course_run.get('first_enrollable_paid_seat_price'):
-        return advertised_course_run.get('first_enrollable_paid_seat_price')
-
-    course_runs = course.get('course_runs') or []
-    active_course_runs = [run for run in course_runs if is_course_run_active(run)]
-    for course_run in sorted(
-        active_course_runs,
-        key=lambda active_course_run: active_course_run['key'].lower(),
-    ):
-        if 'first_enrollable_paid_seat_price' in course_run:
-            return course_run['first_enrollable_paid_seat_price']
-    return None
-
-
 def get_learning_type(content):
     """
     Gets the content's learning type, checking and returning if the content
@@ -1482,6 +1458,16 @@ def get_video_duration(video):
     return video.json_metadata.get('duration')
 
 
+def _first_enrollable_paid_seat_price(course_record):
+    """
+    Returns the course-level first_enrollable_paid_seat_price,
+    or computes it based on the course runs.
+    """
+    if course_value := course_record.get('first_enrollable_paid_seat_price'):
+        return course_value
+    return get_course_first_paid_enrollable_seat_price(course_record)
+
+
 def _algolia_object_from_product(product, algolia_fields):
     """
     Transforms a course or program into an Algolia object.
@@ -1508,7 +1494,7 @@ def _algolia_object_from_product(product, algolia_fields):
             'upcoming_course_runs': get_upcoming_course_runs(searchable_product),
             'skill_names': get_course_skill_names(searchable_product),
             'skills': get_course_skills(searchable_product),
-            'first_enrollable_paid_seat_price': get_course_first_paid_enrollable_seat_price(searchable_product),
+            'first_enrollable_paid_seat_price': _first_enrollable_paid_seat_price(searchable_product),
             'original_image_url': get_course_original_image_url(searchable_product),
             'marketing_url': get_course_marketing_url(searchable_product),
             'outcome': get_course_outcome(searchable_product),

enterprise_catalog/apps/catalog/constants.py

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

enterprise_catalog/apps/catalog/content_metadata_utils.py

@@ -41,3 +41,64 @@ def transform_course_metadata_to_visible(course_metadata):
         course_run_statuses.append(course_run.get('status'))
     course_metadata['course_run_statuses'] = course_run_statuses
     return course_metadata
+
+
+def get_course_run_by_uuid(course, course_run_uuid):
+    """
+    Find a course_run based on uuid
+    Arguments:
+        course (dict): course dict
+        course_run_uuid (str): uuid to lookup
+    Returns:
+        dict: a course_run or None
+    """
+    try:
+        course_run = [
+            run for run in course.get('course_runs', [])
+            if run.get('uuid') == course_run_uuid
+        ][0]
+    except IndexError:
+        return None
+    return course_run
+
+
+def is_course_run_active(course_run):
+    """
+    Checks whether a course run is active. That is, whether the course run is published,
+    enrollable, and marketable.
+    Arguments:
+        course_run (dict): The metadata about a course run.
+    Returns:
+        bool: True if course run is "active"
+    """
+    course_run_status = course_run.get('status') or ''
+    is_published = course_run_status.lower() == 'published'
+    is_enrollable = course_run.get('is_enrollable', False)
+    is_marketable = course_run.get('is_marketable', False)
+
+    return is_published and is_enrollable and is_marketable
+
+
+def get_course_first_paid_enrollable_seat_price(course):
+    """
+    Arguments:
+        course (dict): a dictionary representing a course
+    Returns:
+        The first enrollable paid seat price for the course.
+    """
+    # Use advertised course run.
+    # If that fails use one of the other active course runs.
+    # (The latter is what Discovery does)
+    advertised_course_run = get_course_run_by_uuid(course, course.get('advertised_course_run_uuid'))
+    if advertised_course_run and advertised_course_run.get('first_enrollable_paid_seat_price'):
+        return advertised_course_run.get('first_enrollable_paid_seat_price')
+
+    course_runs = course.get('course_runs') or []
+    active_course_runs = [run for run in course_runs if is_course_run_active(run)]
+    for course_run in sorted(
+        active_course_runs,
+        key=lambda active_course_run: active_course_run['key'].lower(),
+    ):
+        if 'first_enrollable_paid_seat_price' in course_run:
+            return course_run['first_enrollable_paid_seat_price']
+    return None

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,10 @@
     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,12 +36,17 @@
     CONTENT_TYPE_CHOICES,
     COURSE,
     COURSE_RUN,
+    COURSE_RUN_RESTRICTION_TYPE_KEY,
     EXEC_ED_2U_COURSE_TYPE,
     EXEC_ED_2U_ENTITLEMENT_MODE,
     PROGRAM,
+    QUERY_FOR_RESTRICTED_RUNS,
     RESTRICTED_RUNS_ALLOWED_KEY,
     json_serialized_course_modes,
 )
+from enterprise_catalog.apps.catalog.content_metadata_utils import (
+    get_course_first_paid_enrollable_seat_price,
+)
 from enterprise_catalog.apps.catalog.utils import (
     batch,
     enterprise_proxy_login_url,
@@ -685,12 +694,10 @@ class Meta:
 
     @property
     def is_exec_ed_2u_course(self):
-        # pylint: disable=no-member
         return self.content_type == COURSE and self.json_metadata.get('course_type') == EXEC_ED_2U_COURSE_TYPE
 
     @property
     def aggregation_key(self):
-        # pylint: disable=no-member
         return self.json_metadata.get('aggregation_key')
 
     @classmethod
@@ -818,7 +825,72 @@ 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, content_type=COURSE)
+        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_restricted_runs = catalog_query.restricted_runs_allowed.get(course_metadata_dict['key'], [])
+
+        allowed_runs = []
+        allowed_statuses = set()
+        allowed_keys = []
+
+        for run in filtered_metadata['course_runs']:
+            if run.get(COURSE_RUN_RESTRICTION_TYPE_KEY) is None or run['key'] in allowed_restricted_runs:
+                allowed_runs.append(run)
+                allowed_statuses.add(run['status'])
+                allowed_keys.append(run['key'])
+
+        filtered_metadata['course_runs'] = allowed_runs
+        filtered_metadata['course_run_keys'] = allowed_keys
+        filtered_metadata['course_run_statuses'] = sorted(list(allowed_statuses))
+        filtered_metadata['first_enrollable_paid_seat_price'] = get_course_first_paid_enrollable_seat_price(
+            filtered_metadata,
+        )
+
+        return filtered_metadata
 
 
 class RestrictedRunAllowedForRestrictedCourse(TimeStampedModel):
@@ -1180,7 +1252,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,29 +1379,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.keys())
+    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):