feat: backpopulate transaction email and content title

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

@@ -0,0 +1,93 @@
+"""
+Management command to backpopulate transaction email and title
+"""
+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 email and title
+
+    ./manage.py backpopulate_transaction_email_and_title
+    """
+
+    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 email and title
+        """
+        logger.info(f"Processing {subsidy.uuid} transaction {txn.uuid}")
+
+        if txn.lms_user_email is None:
+            lms_user_email = subsidy.email_for_learner(txn.lms_user_id)
+            txn.lms_user_email = lms_user_email
+            logger.info(f"Found {lms_user_email} for {subsidy.uuid} transaction {txn.uuid}")
+        if txn.content_title is None:
+            content_title = subsidy.title_for_content(txn.content_key)
+            txn.content_title = content_title
+            logger.info(f"Found {content_title} for {subsidy.uuid} transaction {txn.uuid}")
+
+        if not self.dry_run:
+            txn.save()
+            logger.info(f"Updated {subsidy.uuid} transaction {txn.uuid}")
+
+    def handle(self, *args, **options):
+        """
+        Find all transactions that are missing email or title 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)
+                incomplete_only_filter = Q(lms_user_email__isnull=True) | Q(content_title__isnull=True)
+                txn_filter = subsidy_filter & incomplete_only_filter
+                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,6 +75,22 @@ 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,
+            quantity=100,
+            fulfillment_identifier=self.fulfillment_identifier
+        )
+        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,
+            quantity=100,
+            fulfillment_identifier=self.fulfillment_identifier
+        )
 
     @mock.patch('enterprise_subsidy.apps.api_client.base_oauth.OAuthAPIClient', return_value=mock.MagicMock())
     def test_write_reversals_from_enterprise_unenrollment_with_existing_reversal(self, mock_oauth_client):
@@ -889,3 +905,67 @@ def test_write_reversals_from_geag_enterprise_unenrollments_unknown_provider(
         with self.assertRaises(FulfillmentException):
             call_command('write_reversals_from_enterprise_unenrollments')
         assert Reversal.objects.count() == 0
+
+    @mock.patch("enterprise_subsidy.apps.subsidy.models.Subsidy.lms_user_client")
+    @mock.patch("enterprise_subsidy.apps.content_metadata.api.ContentMetadataApi.get_content_summary")
+    def test_backpopulate_transaction_email_and_title(
+        self,
+        mock_get_content_summary,
+        mock_lms_user_client,
+    ):
+        """
+        Test that the backpopulate_transaction_email_and_title management command backpopulates the email and title
+        """
+        expected_email_address = '[email protected]'
+        mock_lms_user_client.return_value.best_effort_user_data.return_value = {
+            'email': expected_email_address,
+        }
+        expected_content_title = 'a content title'
+        mock_get_content_summary.return_value = {
+            'content_uuid': 'a content uuid',
+            'content_key': 'a content key',
+            'content_title': expected_content_title,
+            'source': 'edX',
+            'mode': 'verified',
+            'content_price': 10000,
+            'geag_variant_id': None,
+        }
+        call_command('backpopulate_transaction_email_and_title')
+        self.transaction_to_backpopulate.refresh_from_db()
+        self.internal_transaction_to_backpopulate.refresh_from_db()
+        assert self.transaction_to_backpopulate.lms_user_email == expected_email_address
+        assert self.transaction_to_backpopulate.content_title == expected_content_title
+        assert self.internal_transaction_to_backpopulate.lms_user_email is None
+        assert self.internal_transaction_to_backpopulate.content_title is None
+
+    @mock.patch("enterprise_subsidy.apps.subsidy.models.Subsidy.lms_user_client")
+    @mock.patch("enterprise_subsidy.apps.content_metadata.api.ContentMetadataApi.get_content_summary")
+    def test_backpopulate_transaction_email_and_title_include_internal(
+        self,
+        mock_get_content_summary,
+        mock_lms_user_client,
+    ):
+        """
+        Test that the backpopulate_transaction_email_and_title while including internal subsidies
+        """
+        expected_email_address = '[email protected]'
+        mock_lms_user_client.return_value.best_effort_user_data.return_value = {
+            'email': expected_email_address,
+        }
+        expected_content_title = 'a content title'
+        mock_get_content_summary.return_value = {
+            'content_uuid': 'a content uuid',
+            'content_key': 'a content key',
+            'content_title': expected_content_title,
+            'source': 'edX',
+            'mode': 'verified',
+            'content_price': 10000,
+            'geag_variant_id': None,
+        }
+        call_command('backpopulate_transaction_email_and_title', include_internal_subsidies=True)
+        self.transaction_to_backpopulate.refresh_from_db()
+        self.internal_transaction_to_backpopulate.refresh_from_db()
+        assert self.transaction_to_backpopulate.lms_user_email == expected_email_address
+        assert self.transaction_to_backpopulate.content_title == expected_content_title
+        assert self.internal_transaction_to_backpopulate.lms_user_email == expected_email_address
+        assert self.internal_transaction_to_backpopulate.content_title == expected_content_title

enterprise_subsidy/apps/transaction/utils.py

@@ -2,10 +2,35 @@
 Utility functions used in the implementation of subsidy Transactions.
 """
 
+from django.db.models import Q
+
 
 def generate_transaction_reversal_idempotency_key(fulfillment_uuid, enrollment_unenrolled_at):
     """
     Generates a unique idempotency key for a transaction reversal using the fulfillment uuid and time at which the
     unenrollment occurred.
     """
     return f'unenrollment-reversal-{fulfillment_uuid}-{enrollment_unenrolled_at}'
+
+
+def batch_by_pk(ModelClass, extra_filter=Q(), batch_size=10000):
+    """
+    yield per batch efficiently
+    using limit/offset does a lot of table scanning to reach higher offsets
+    this scanning can be slow on very large tables
+    if you order by pk, you can use the pk as a pivot rather than offset
+    this utilizes the index, which is faster than scanning to reach offset
+    Example usage:
+    filter_query = Q(column='value')
+    for items_batch in batch_by_pk(Model, extra_filter=filter_query):
+        for item in items_batch:
+            ...
+    """
+    qs = ModelClass.objects.filter(extra_filter).order_by('pk')[:batch_size]
+    while qs.exists():
+        yield qs
+        # qs.last() doesn't work here because we've already sliced
+        # loop through so we eventually grab the last one
+        for item in qs:
+            start_pk = item.pk
+        qs = ModelClass.objects.filter(pk__gt=start_pk).filter(extra_filter).order_by('pk')[:batch_size]