feat: add tool to backpopulate parent_content_key

ENT-8389

enterprise_subsidy/apps/transaction/management/commands/backpopulate_transaction_parent_content_key.py

@@ -0,0 +1,101 @@
+"""
+Management command to backpopulate transaction parent_content_key
+"""
+import logging
+
+from django.core.management.base import BaseCommand
+from django.db.models import Q
+from openedx_ledger.models import Transaction
+
+from enterprise_subsidy.apps.subsidy.models import Subsidy
+from enterprise_subsidy.apps.transaction.utils import batch_by_pk
+
+logger = logging.getLogger(__name__)
+
+
+class Command(BaseCommand):
+    """
+    Management command for backpopulating transaction parent_content_key
+
+    ./manage.py backpopulate_transaction_parent_content_key
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.dry_run = False
+        self.include_internal_subsidies = False
+
+    def add_arguments(self, parser):
+        """
+        Entry point for subclassed commands to add custom arguments.
+        """
+        parser.add_argument(
+            '--dry-run',
+            action='store_true',
+            dest='dry_run',
+            help=(
+                'If set, no updates will occur; will instead log '
+                'the actions that would have been taken.'
+            ),
+        )
+        parser.add_argument(
+            '--include-internal-subsidies',
+            action='store_true',
+            dest='include_internal_subsidies',
+            help=(
+                'If set, internal subsidies will be included in the backpopulating'
+            ),
+        )
+
+    def process_transaction(self, subsidy, txn):
+        """
+        Given a transaction (and it's subsidy), backpopulate the parent_content_key.
+        """
+        logger.info(f"Processing subsidy={subsidy.uuid}, transaction={txn.uuid}")
+
+        try:
+            if txn.parent_content_key is None and txn.content_key is not None:
+                parent_content_key = subsidy.metadata_summary_for_content(txn.content_key).get('content_key')
+                txn.parent_content_key = parent_content_key
+                logger.info(
+                    f"Found parent_content_key={parent_content_key} "
+                    f"for subsidy={subsidy.uuid}, transaction={txn.uuid}"
+                )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+            logger.exception(
+                f"Error while processing parent_content_key for subsidy={subsidy.uuid}, transaction={txn.uuid}: {e}"
+            )
+
+        if not self.dry_run:
+            txn.save()
+            logger.info(f"Updated subsidy={subsidy.uuid}, transaction={txn.uuid}")
+
+    def handle(self, *args, **options):
+        """
+        Find all transactions that are missing parent_content_key and backpopulate them.
+        """
+        if options.get('dry_run'):
+            self.dry_run = True
+            logger.info("Running in dry-run mode. No updates will occur.")
+
+        if options.get('include_internal_subsidies'):
+            self.include_internal_subsidies = True
+            logger.info("Including internal_only subsidies while backpopulating.")
+
+        subsidy_filter = Q()
+        if not self.include_internal_subsidies:
+            subsidy_filter = Q(internal_only=False)
+
+        for subsidies in batch_by_pk(Subsidy, extra_filter=subsidy_filter):
+            for subsidy in subsidies:
+                logger.info(f"Processing subsidy {subsidy.uuid}")
+
+                subsidy_filter = Q(ledger=subsidy.ledger)
+                # We can only populate the parent_content_key when there's a child content key. Empty child content keys
+                # mayhappen with test data.
+                incomplete_parent_content_key = Q(parent_content_key__isnull=True) & Q(content_key__isnull=False)
+                txn_filter = subsidy_filter & incomplete_parent_content_key
+
+                for txns in batch_by_pk(Transaction, extra_filter=txn_filter):
+                    for txn in txns:
+                        self.process_transaction(subsidy, txn)

enterprise_subsidy/apps/transaction/tests/test_management.py

@@ -75,26 +75,41 @@ def setUp(self):
             external_fulfillment_provider=self.unknown_provider,
             transaction=self.unknown_transaction,
         )
+
         self.transaction_to_backpopulate = TransactionFactory(
             ledger=self.ledger,
             lms_user_email=None,
             content_title=None,
+            # We can't just set parent_content_key to None because it will break content_key (derived factory field).
+            # Do it after object creation.
+            # parent_content_key=None,
             quantity=100,
             fulfillment_identifier=self.fulfillment_identifier
         )
+        self.transaction_to_backpopulate.parent_content_key = None
+        self.transaction_to_backpopulate.save()
+
         self.internal_ledger = LedgerFactory()
         self.internal_subsidy = SubsidyFactory(ledger=self.internal_ledger, internal_only=True)
         self.internal_transaction_to_backpopulate = TransactionFactory(
             ledger=self.internal_ledger,
             lms_user_email=None,
             content_title=None,
         )
+        self.internal_transaction_to_backpopulate.parent_content_key = None
+        self.internal_transaction_to_backpopulate.save()
+
         self.transaction_not_to_backpopulate = TransactionFactory(
             ledger=self.ledger,
+
+            # Setting content_key or lms_user_id to None force-disables backpopulation.
+            content_key=None,
             lms_user_id=None,
+
+            # The target fields to backpopulate are empty, nevertheless.
             lms_user_email=None,
-            content_key=None,
             content_title=None,
+            parent_content_key=None,
         )
 
     @mock.patch('enterprise_subsidy.apps.api_client.base_oauth.OAuthAPIClient', return_value=mock.MagicMock())
@@ -980,3 +995,56 @@ def test_backpopulate_transaction_email_and_title_include_internal(
         assert self.internal_transaction_to_backpopulate.content_title == expected_content_title
         assert self.transaction_not_to_backpopulate.lms_user_email is None
         assert self.transaction_not_to_backpopulate.content_title is None
+
+    @mock.patch("enterprise_subsidy.apps.content_metadata.api.ContentMetadataApi.get_content_summary")
+    def test_backpopulate_transaction_parent_content_key(
+        self,
+        mock_get_content_summary,
+    ):
+        """
+        Test that the backpopulate_transaction_parent_content_key management command backpopulates the
+        parent_content_key.
+        """
+        expected_parent_content_key = 'edx+101'
+        mock_get_content_summary.return_value = {
+            'content_uuid': 'a content uuid',
+            'content_key': expected_parent_content_key,
+            'content_title': 'a content title',
+            'source': 'edX',
+            'mode': 'verified',
+            'content_price': 10000,
+            'geag_variant_id': None,
+        }
+        call_command('backpopulate_transaction_parent_content_key')
+        self.transaction_to_backpopulate.refresh_from_db()
+        self.internal_transaction_to_backpopulate.refresh_from_db()
+        self.transaction_not_to_backpopulate.refresh_from_db()
+        assert self.transaction_to_backpopulate.parent_content_key == expected_parent_content_key
+        assert self.internal_transaction_to_backpopulate.parent_content_key is None
+        assert self.transaction_not_to_backpopulate.parent_content_key is None
+
+    @mock.patch("enterprise_subsidy.apps.content_metadata.api.ContentMetadataApi.get_content_summary")
+    def test_backpopulate_transaction_parent_content_key_include_internal(
+        self,
+        mock_get_content_summary,
+    ):
+        """
+        Test backpopulate_transaction_parent_content_key while including internal subsidies.
+        """
+        expected_parent_content_key = 'edx+101'
+        mock_get_content_summary.return_value = {
+            'content_uuid': 'a content uuid',
+            'content_key': expected_parent_content_key,
+            'content_title': 'a content title',
+            'source': 'edX',
+            'mode': 'verified',
+            'content_price': 10000,
+            'geag_variant_id': None,
+        }
+        call_command('backpopulate_transaction_parent_content_key', include_internal_subsidies=True)
+        self.transaction_to_backpopulate.refresh_from_db()
+        self.internal_transaction_to_backpopulate.refresh_from_db()
+        self.transaction_not_to_backpopulate.refresh_from_db()
+        assert self.transaction_to_backpopulate.parent_content_key == expected_parent_content_key
+        assert self.internal_transaction_to_backpopulate.parent_content_key == expected_parent_content_key
+        assert self.transaction_not_to_backpopulate.parent_content_key is None