diff --git a/common/djangoapps/student/management/commands/populate_users_emails_on_braze.py b/common/djangoapps/student/management/commands/populate_users_emails_on_braze.py new file mode 100644 index 0000000000..6c13392d22 --- /dev/null +++ b/common/djangoapps/student/management/commands/populate_users_emails_on_braze.py @@ -0,0 +1,139 @@ +""" Management command to add user emails data on Braze. """ +import logging +import time + +from braze.client import BrazeClient +from braze.exceptions import BrazeClientError +from django.conf import settings +from django.contrib.auth import get_user_model +from django.core.management.base import BaseCommand + +from common.djangoapps.util.query import use_read_replica_if_available + +User = get_user_model() + +MARKETING_EMAIL_ATTRIBUTE_NAME = 'is_marketable' +TRACK_USER_COMPONENT_CHUNK_SIZE = 75 + +logger = logging.getLogger(__name__) + + +class Command(BaseCommand): + """ + Command to add user email address on Braze against a user_id. + Example usage: + $ ./manage.py lms populate_user_emails_on_braze + """ + help = """ + Updates user accounts with email addresses on braze. + """ + + def add_arguments(self, parser): + """ + Function to get command arguments + """ + parser.add_argument( + '--batch-delay', + type=float, + dest='batch_delay', + default=0.5, + help='Time delay in seconds between each iteration' + ) + parser.add_argument( + '--batch-size', + type=int, + dest='batch_size', + default=10000, + help='Batch size' + ) + parser.add_argument( + '--starting-user-id', + type=int, + dest='starting_user_id', + default=0, + help='Starting user id to process a specific batch of users. ' + 'Both start and end id should be provided.', + ) + parser.add_argument( + '--ending-user-id', + type=int, + dest='ending_user_id', + default=0, + help='Ending user id (inclusive) to process a specific batch of users. ' + 'Both start and end id should be provided.', + ) + + def __init__(self): + super().__init__() + self.braze_client = BrazeClient( + api_key=settings.EDX_BRAZE_API_KEY, + api_url=settings.EDX_BRAZE_API_SERVER, + app_id='', + ) + + @staticmethod + def _chunks(users, chunk_size=TRACK_USER_COMPONENT_CHUNK_SIZE): + """ + Yields successive chunks of users. The size of each chunk is determined by + TRACK_USER_COMPONENT_CHUNK_SIZE which is set to 75. + Reference: https://www.braze.com/docs/api/endpoints/user_data/post_user_track/ + """ + for index in range(0, len(users), chunk_size): + yield users[index:index + chunk_size] + + def _get_user_batch(self, batch_start_id, batch_end_id): + """ + This returns the batch of users. + """ + query = User.objects.filter( + id__gte=batch_start_id, id__lt=batch_end_id, + ).select_related('profile').values_list( + 'id', 'email', named=True, + ).order_by('id') + + return use_read_replica_if_available(query) + + def _update_braze_attributes(self, users): + """ + Sends Braze API request to update user account. + Fields sent using the API include: + - external_id (user_id) + - email + """ + attributes = [] + for user in users: + attributes.append( + { + "external_id": user.id, + "email": user.email, + } + ) + + try: + self.braze_client.track_user(attributes=attributes) + except BrazeClientError as error: + logger.error(f'Failed to update attributes. Error: {error}') + + def handle(self, *args, **options): + """ + Handler to run the command. + """ + sleep_time = options['batch_delay'] + batch_size = options['batch_size'] + starting_user_id = options['starting_user_id'] + ending_user_id = options['ending_user_id'] + + all_users_query = use_read_replica_if_available(User.objects) + total_users_count = ending_user_id if ending_user_id else all_users_query.count() + + for index in range(starting_user_id, total_users_count + 1, batch_size): + users = self._get_user_batch(index, index + batch_size) + logger.info(f'Processing users with user ids in {index} - {(index + batch_size) - 1} range') + + # Force evaluating the query to avoid multiple hits to db + # when we evaluate the chunks. + evaluated_users = list(users) + for user_chunk in self._chunks(evaluated_users): + self._update_braze_attributes(user_chunk) + + time.sleep(sleep_time) diff --git a/common/djangoapps/student/management/tests/test_populate_users_emails_on_braze.py b/common/djangoapps/student/management/tests/test_populate_users_emails_on_braze.py new file mode 100644 index 0000000000..8d8ad31c4c --- /dev/null +++ b/common/djangoapps/student/management/tests/test_populate_users_emails_on_braze.py @@ -0,0 +1,91 @@ +""" +Unittests for populate_marketing_opt_in_user_attribute management command. +""" +from unittest.mock import patch, MagicMock + +from braze.exceptions import BrazeClientError +from django.core.management import call_command +from django.test import TestCase +from testfixtures import LogCapture + +from common.djangoapps.student.tests.factories import UserFactory +from openedx.core.djangolib.testing.utils import skip_unless_lms + +LOGGER_NAME = 'common.djangoapps.student.management.commands.populate_users_emails_on_braze' + + +@skip_unless_lms +class TestPopulateUsersEmailsOnBraze(TestCase): + """ + Tests for PopulateUsersEmailsOnBraze management command. + """ + + @classmethod + def setUpClass(cls): + super().setUpClass() + for index in range(15): + user = UserFactory() + + @patch('common.djangoapps.student.management.commands.populate_users_emails_on_braze.BrazeClient.track_user') + def test_command_updates_users_on_braze(self, track_user): + """ + Test that Braze API is called successfully for all users + """ + track_user.return_value = MagicMock() + call_command('populate_users_emails_on_braze', batch_delay=0) + assert track_user.called + + @patch('common.djangoapps.student.management.commands.populate_users_emails_on_braze.BrazeClient.track_user') + def test_logs_for_success(self, track_user): + """ + Test logs for a successful run of updating user accounts on Braze + """ + track_user.return_value = MagicMock() + with LogCapture(LOGGER_NAME) as log: + call_command( + 'populate_users_emails_on_braze', + batch_size=5, + batch_delay=0, + ) + log.check( + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 0 - 4 range'), + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 5 - 9 range'), + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 10 - 14 range'), + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 15 - 19 range'), + ) + + @patch('common.djangoapps.student.management.commands.populate_users_emails_on_braze.BrazeClient.track_user') + def test_logs_for_failure(self, track_user): + """ + Test logs for when the update to Braze fails + """ + track_user.side_effect = BrazeClientError('Update to attributes failed.') + with LogCapture(LOGGER_NAME) as log: + call_command( + 'populate_users_emails_on_braze', + batch_size=5, + batch_delay=0, + ) + log.check_present( + (LOGGER_NAME, 'ERROR', 'Failed to update attributes. Error: Update to attributes failed.'), + ) + + @patch('common.djangoapps.student.management.commands.populate_users_emails_on_braze.BrazeClient.track_user') + def test_running_a_specific_batch(self, track_user): + """ + Test running command for a specific batch of users + """ + track_user.return_value = MagicMock() + with LogCapture(LOGGER_NAME) as log: + call_command( + 'populate_users_emails_on_braze', + batch_size=5, + batch_delay=0, + starting_user_id=2, + ending_user_id=13, + ) + log.check( + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 2 - 6 range'), + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 7 - 11 range'), + (LOGGER_NAME, 'INFO', 'Processing users with user ids in 12 - 16 range'), + )