feat: more restricted runs and contains_content_items logic

enterprise_catalog/apps/api/v1/tests/test_enterprise_customer_views.py

@@ -367,15 +367,23 @@ def test_contains_catalog_list_with_content_ids_param(self):
         # Create a two catalogs that have the content we're looking for
         content_key = 'fake-key+101x'
         second_catalog = EnterpriseCatalogFactory(enterprise_uuid=self.enterprise_uuid)
+        # Add some non-null, but irrelevant restricted runs to this catalog
+        second_catalog.catalog_query.content_filter[RESTRICTED_RUNS_ALLOWED_KEY] = {
+            'something+else1': ['course-v1:something+else+restrictedrun']
+        }
+        second_catalog.catalog_query.save()
+
         relevant_content = ContentMetadataFactory(content_key=content_key)
         self.add_metadata_to_catalog(second_catalog, [relevant_content])
         url = self._get_contains_content_base_url() + '?course_run_ids=' + content_key + \
             '&get_catalogs_containing_specified_content_ids=True'
         self.assert_correct_contains_response(url, True)
 
         response = self.client.get(url)
-        catalog_list = response.json()['catalog_list']
+        response_payload = response.json()
+        catalog_list = response_payload['catalog_list']
         assert set(catalog_list) == {str(second_catalog.uuid)}
+        self.assertIsNone(response_payload['restricted_runs_allowed'])
 
     def test_contains_catalog_key_restricted_runs_allowed(self):
         """
@@ -408,12 +416,16 @@ def test_contains_catalog_key_restricted_runs_allowed(self):
         # make sure to also request a course key that has no restricted runs,
         # and then assert that it is *not* included in the response payload.
         url = self._get_contains_content_base_url() + \
-            '?course_run_ids=org+key1&course_run_ids=org+key2&course_run_ids=org+key3'
+            '?course_run_ids=org+key1&course_run_ids=org+key2&course_run_ids=org+key3&get_catalog_list=true'
 
         response = self.client.get(url)
         response_payload = response.json()
 
         self.assertTrue(response_payload.get('contains_content_items'))
+        self.assertEqual(
+            set(response_payload['catalog_list']),
+            set([str(catalog.uuid), str(other_catalog.uuid)])
+        )
         self.assertEqual(
             response_payload['restricted_runs_allowed'],
             {
@@ -430,6 +442,67 @@ def test_contains_catalog_key_restricted_runs_allowed(self):
             }
         )
 
+    def test_restricted_course_disallowed_if_course_not_in_catalog(self):
+        """
+        Tests that a requested course with restricted runs is "disallowed"
+        if the course is not part of a customer's catalog.
+        """
+        catalog = EnterpriseCatalogFactory(enterprise_uuid=self.enterprise_uuid)
+        catalog.catalog_query.content_filter[RESTRICTED_RUNS_ALLOWED_KEY] = {
+            'org+key1': ['course-v1:org+key1+restrictedrun']
+        }
+        catalog.catalog_query.save()
+        ContentMetadataFactory(content_key='org+key1')
+        # don't add this content to the catalog
+
+        url = self._get_contains_content_base_url() + '?course_run_ids=org+key1'
+
+        response = self.client.get(url)
+        response_payload = response.json()
+
+        self.assertFalse(response_payload.get('contains_content_items'))
+        self.assertIsNone(response_payload['restricted_runs_allowed'])
+
+    def test_restricted_course_run_allowed_even_if_course_not_in_catalog(self):
+        """
+        Tests that a requested restricted course run is "allowed"
+        even if the course is not part of a customer's catalog. This is necessary
+        because typically restricted runs will not have corresponding
+        `ContentMetadata` records present in the DB, so a lookup via only
+        `EnterpriseCatalog.contains_content_keys` will fail. We rely
+        on the restricted run mapping to ascertain the *implicit* inclusion
+        of a restricted course run in a catalog.
+        """
+        catalog = EnterpriseCatalogFactory(enterprise_uuid=self.enterprise_uuid)
+        catalog.catalog_query.content_filter[RESTRICTED_RUNS_ALLOWED_KEY] = {
+            'org+key1': ['course-v1:org+key1+restrictedrun']
+        }
+        catalog.catalog_query.save()
+        ContentMetadataFactory(content_key='org+key1')
+        # don't add this content to the catalog
+
+        url = self._get_contains_content_base_url() + \
+            '?course_run_ids=course-v1:org+key1+restrictedrun&get_catalog_list=true'
+
+        response = self.client.get(url)
+        response_payload = response.json()
+
+        self.assertTrue(response_payload.get('contains_content_items'))
+        self.assertEqual(
+            response_payload['catalog_list'],
+            [str(catalog.uuid)],
+        )
+        self.assertEqual(
+            response_payload['restricted_runs_allowed'],
+            {
+                'org+key1': {
+                    'course-v1:org+key1+restrictedrun': {
+                        'catalog_uuids': [str(catalog.uuid)]
+                    },
+                },
+            }
+        )
+
     def test_contains_catalog_list_parent_key(self):
         """
         Verify the contains_content_items endpoint returns a list of catalogs the course is in

enterprise_catalog/apps/api/v1/views/enterprise_catalog_contains_content_items.py

@@ -17,6 +17,9 @@
 )
 from enterprise_catalog.apps.api.v1.utils import unquote_course_keys
 from enterprise_catalog.apps.api.v1.views.base import BaseViewSet
+from enterprise_catalog.apps.catalog.api import (
+    catalog_contains_any_restricted_course_run,
+)
 from enterprise_catalog.apps.catalog.models import EnterpriseCatalog
 
 
@@ -58,4 +61,9 @@ def contains_content_items(self, request, uuid, course_run_ids, program_uuids, *
 
         enterprise_catalog = self.get_object()
         contains_content_items = enterprise_catalog.contains_content_keys(course_run_ids + program_uuids)
-        return Response({'contains_content_items': contains_content_items})
+        contains_requested_restricted_items = catalog_contains_any_restricted_course_run(
+            enterprise_catalog, course_run_ids,
+        )
+        return Response({
+            'contains_content_items': contains_content_items or contains_requested_restricted_items
+        })

enterprise_catalog/apps/api/v1/views/enterprise_customer.py

@@ -1,6 +1,5 @@
 import logging
 import uuid
-from collections import defaultdict
 
 from django.utils.decorators import method_decorator
 from edx_rbac.utils import get_decoded_jwt
@@ -15,8 +14,11 @@
 from enterprise_catalog.apps.api.v1.serializers import ContentMetadataSerializer
 from enterprise_catalog.apps.api.v1.utils import unquote_course_keys
 from enterprise_catalog.apps.api.v1.views.base import BaseViewSet
+from enterprise_catalog.apps.catalog.api import (
+    catalog_contains_any_restricted_course_run,
+    get_restricted_runs_allowed_for_query,
+)
 from enterprise_catalog.apps.catalog.constants import (
-    COURSE_RUN_KEY_PREFIX,
     RESTRICTED_RUNS_ALLOWED_KEY,
 )
 from enterprise_catalog.apps.catalog.models import EnterpriseCatalog
@@ -93,7 +95,7 @@ def contains_content_items(self, request, enterprise_uuid, course_run_ids, progr
             'get_catalogs_containing_specified_content_ids', False
         )
         get_catalog_list = request.GET.get('get_catalog_list', False)
-        course_run_ids = unquote_course_keys(course_run_ids)
+        requested_course_or_run_keys = unquote_course_keys(course_run_ids)
 
         try:
             uuid.UUID(enterprise_uuid)
@@ -106,50 +108,33 @@ def contains_content_items(self, request, enterprise_uuid, course_run_ids, progr
                 f'Error: invalid enterprice customer uuid: "{enterprise_uuid}" provided.',
                 status=HTTP_400_BAD_REQUEST
             )
-        customer_catalogs = list(EnterpriseCatalog.objects.filter(enterprise_uuid=enterprise_uuid))
+        customer_catalogs = list(
+            EnterpriseCatalog.objects.select_related(
+                'catalog_query',
+            ).filter(
+                enterprise_uuid=enterprise_uuid,
+            ))
 
-        any_catalog_contains_content_items = False
         catalogs_that_contain_course = []
         for catalog in customer_catalogs:
-            contains_content_items = catalog.contains_content_keys(course_run_ids + program_uuids)
-            if contains_content_items:
-                any_catalog_contains_content_items = True
-                if not (get_catalogs_containing_specified_content_ids or get_catalog_list):
-                    # Break as soon as we find a catalog that contains the specified content
-                    break
-                catalogs_that_contain_course.append(catalog.uuid)
+            contains_content_items = catalog.contains_content_keys(requested_course_or_run_keys + program_uuids)
+            contains_requested_restricted_items = catalog_contains_any_restricted_course_run(
+                catalog, requested_course_or_run_keys,
+            )
+            if contains_content_items or contains_requested_restricted_items:
+                catalogs_that_contain_course.append(catalog)
 
         response_data = {
-            'contains_content_items': any_catalog_contains_content_items,
+            'contains_content_items': bool(catalogs_that_contain_course),
         }
         if (get_catalogs_containing_specified_content_ids or get_catalog_list):
-            response_data['catalog_list'] = catalogs_that_contain_course
+            response_data['catalog_list'] = [str(catalog.uuid) for catalog in catalogs_that_contain_course]
 
-        response_data[RESTRICTED_RUNS_ALLOWED_KEY] = self._get_restricted_runs_allowed_for_query(
-            course_run_ids, customer_catalogs,
+        response_data[RESTRICTED_RUNS_ALLOWED_KEY] = get_restricted_runs_allowed_for_query(
+            requested_course_or_run_keys, catalogs_that_contain_course,
         )
         return Response(response_data)
 
-    def _get_restricted_runs_allowed_for_query(self, course_run_ids, customer_catalogs):
-        # filter the set of restricted course keys down to only
-        # those requested by the client, and only if those requested keys
-        # are top-level course keys (NOT course run keys).
-        requested_course_keys = {
-            key for key in course_run_ids if not key.startswith(COURSE_RUN_KEY_PREFIX)
-        }
-        serialized_data = defaultdict(lambda: defaultdict(lambda: {'catalog_uuids': set()}))
-        for catalog in customer_catalogs:
-            if not catalog.restricted_runs_allowed:
-                continue
-            for restricted_course_key, restricted_runs in catalog.restricted_runs_allowed.items():
-                if restricted_course_key not in requested_course_keys:
-                    continue
-                course_dict = serialized_data[restricted_course_key]
-                for course_run_key in restricted_runs:
-                    run_dict = course_dict[course_run_key]
-                    run_dict['catalog_uuids'].add(str(catalog.uuid))
-        return serialized_data or None
-
     @action(detail=True, methods=['post'])
     def filter_content_items(self, request, enterprise_uuid, **kwargs):
         """

enterprise_catalog/apps/catalog/api.py

@@ -0,0 +1,73 @@
+"""
+Python interface for the ``catalog`` module.
+"""
+from collections import defaultdict
+
+from enterprise_catalog.apps.catalog.constants import COURSE_RUN_KEY_PREFIX
+
+
+def get_restricted_runs_allowed_for_query(course_run_ids, customer_catalogs):
+    """
+    Filter the set of restricted course keys down to only
+    those requested in ``course_run_ids``.
+
+    Params:
+        course_run_ids: A list of either course or course run keys.
+        customer_catalogs: A list of ``EnterpriseCatalog`` records.
+
+    Returns:
+        A dictionary mapping courses to allowed restricted runs for catalogs that
+        allow the restricted inclusion. That is:
+        ```
+        {
+            'org+key1': {
+                'course-v1:org+key1+restrictedrun': {
+                    'catalog_uuids': {'catalog-1.uuid'}
+                },
+            },
+            'org+key3': {
+                'course-v1:org+key3+restrictedrun': {
+                    'catalog_uuids': {'catalog-2.uuid'}
+                },
+            },
+        }
+        ```
+    """
+    requested_course_keys = {
+        key for key in course_run_ids if not key.startswith(COURSE_RUN_KEY_PREFIX)
+    }
+    requested_run_keys = {
+        key for key in course_run_ids if key.startswith(COURSE_RUN_KEY_PREFIX)
+    }
+    serialized_data = defaultdict(lambda: defaultdict(lambda: {'catalog_uuids': set()}))
+    for catalog in customer_catalogs:
+        if not catalog.restricted_runs_allowed:
+            continue
+
+        for restricted_course_key, restricted_runs in catalog.restricted_runs_allowed.items():
+            matching_runs = bool(set(restricted_runs).intersection(requested_run_keys))
+            matching_course = restricted_course_key in requested_course_keys
+            if not (matching_course or matching_runs):
+                continue
+
+            course_dict = serialized_data[restricted_course_key]
+            for course_run_key in restricted_runs:
+                run_dict = course_dict[course_run_key]
+                run_dict['catalog_uuids'].add(str(catalog.uuid))
+    return serialized_data or None
+
+
+def catalog_contains_any_restricted_course_run(enterprise_catalog, course_run_keys):
+    """
+    Params:
+        course_run_keys: A list of candidate course run keys.
+        customer_catalog: An ``EnterpriseCatalog`` record.
+
+    Returns:
+        A boolean indicating if at least one of the provided course_run_keys
+        is present in the set of restricted runs for the given enterprise_catalog.
+    """
+    for course_run_key in course_run_keys:
+        if course_run_key in enterprise_catalog.restricted_courses_by_run_key:
+            return True
+    return False

enterprise_catalog/apps/catalog/models.py

@@ -140,6 +140,29 @@ def restricted_runs_allowed(self):
             for course_key, course_run_list in mapping.items()
         }
 
+    @cached_property
+    def restricted_courses_by_run_key(self):
+        """
+        Returns a reverse mapping of self.restricted_runs_allowed, e.g.
+        ```
+        {
+            "course-v1:edX+FUN+3T2024": "edX+FUN",
+            "course-v1:edX+FUN+3T2025": "edX+FUN",
+            "course-v1:edX+GAMES+3T2024": "edX+GAMES",
+        }
+        ```
+
+        Returns an empty dict if no restricted runs are allowed for this CatalogQuery.
+        """
+        if not self.restricted_runs_allowed:
+            return {}
+
+        restricted_courses_by_run = {}
+        for course_key, restricted_run_list in self.restricted_runs_allowed.items():
+            for run_key in restricted_run_list:
+                restricted_courses_by_run[run_key] = course_key
+        return restricted_courses_by_run
+
     @classmethod
     def get_by_uuid(cls, uuid):
         try:
@@ -239,6 +262,10 @@ def content_metadata(self):
     def restricted_runs_allowed(self):
         return self.catalog_query.restricted_runs_allowed
 
+    @cached_property
+    def restricted_courses_by_run_key(self):
+        return self.catalog_query.restricted_courses_by_run_key
+
     @cached_property
     def enterprise_customer(self):
         """