From 7a700dcfd19f574fd8e73ff53fdc4386f51a7adc Mon Sep 17 00:00:00 2001 From: Zainab Amir Date: Mon, 20 Jun 2022 10:19:09 +0500 Subject: [PATCH] feat: command to populate user data on braze (#30556) --- .../commands/populate_user_data_on_braze.py | 164 ++++++++++++++++++ .../tests/test_populate_user_data_on_braze.py | 93 ++++++++++ common/djangoapps/student/tests/factories.py | 10 ++ 3 files changed, 267 insertions(+) create mode 100644 common/djangoapps/student/management/commands/populate_user_data_on_braze.py create mode 100644 common/djangoapps/student/management/tests/test_populate_user_data_on_braze.py diff --git a/common/djangoapps/student/management/commands/populate_user_data_on_braze.py b/common/djangoapps/student/management/commands/populate_user_data_on_braze.py new file mode 100644 index 0000000000..f8fdbd8bf3 --- /dev/null +++ b/common/djangoapps/student/management/commands/populate_user_data_on_braze.py @@ -0,0 +1,164 @@ +""" Management command to add user 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.student.models import UserAttribute +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 is_marketable field on Braze against a user_id. If an account + with user_id doesn't exist on Braze, it automatically creates one. + + Example usage: + $ ./manage.py lms backfill_user_data_on_braze + """ + help = """ + Creates user accounts on Braze for missing users and adds the is_marketable + user attribute to all user accounts 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', 'profile__name', 'username', named=True, + ).order_by('id') + + return use_read_replica_if_available(query) + + def _get_marketable_users(self, batch_start_id, batch_end_id): + """ + For a given batch of users it gives the marketable users list + """ + get_marketable_users_query = UserAttribute.objects.filter( + user_id__gte=batch_start_id, + user_id__lt=batch_end_id, + name=MARKETING_EMAIL_ATTRIBUTE_NAME, + value__in=['True', 'true', '1'], + ).values_list('user_id', flat=True).order_by('user_id') + return set(use_read_replica_if_available(get_marketable_users_query)) + + def _update_braze_attributes(self, users, marketable_users): + """ + Sends Braze API request to update/create user account. + + Fields sent using the API include: + - external_id (user_id) + - email + - name + - username + - is_marketable + """ + attributes = [] + for user in users: + attributes.append( + { + "external_id": user.id, + "email": user.email, + "name": user.profile__name, + "username": user.username, + "is_marketable": user.id in marketable_users, + } + ) + + 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) + marketable_users = self._get_marketable_users(index, index + batch_size) + for user_chunk in self._chunks(evaluated_users): + self._update_braze_attributes(user_chunk, marketable_users) + + time.sleep(sleep_time) diff --git a/common/djangoapps/student/management/tests/test_populate_user_data_on_braze.py b/common/djangoapps/student/management/tests/test_populate_user_data_on_braze.py new file mode 100644 index 0000000000..9fe0b990c7 --- /dev/null +++ b/common/djangoapps/student/management/tests/test_populate_user_data_on_braze.py @@ -0,0 +1,93 @@ +""" +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.management.commands.populate_user_data_on_braze import MARKETING_EMAIL_ATTRIBUTE_NAME +from common.djangoapps.student.tests.factories import UserFactory, UserAttributeFactory +from openedx.core.djangolib.testing.utils import skip_unless_lms + +LOGGER_NAME = 'common.djangoapps.student.management.commands.populate_user_data_on_braze' + + +@skip_unless_lms +class TestPopulateUserDataOnBraze(TestCase): + """ + Tests for PopulateUserDataOnBraze management command. + """ + + @classmethod + def setUpClass(cls): + super().setUpClass() + for index in range(15): + user = UserFactory() + UserAttributeFactory(user=user, name=MARKETING_EMAIL_ATTRIBUTE_NAME, value=(index % 2 == 0)) + + @patch('common.djangoapps.student.management.commands.populate_user_data_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_user_data_on_braze', batch_delay=0) + assert track_user.called + + @patch('common.djangoapps.student.management.commands.populate_user_data_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_user_data_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_user_data_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_user_data_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_user_data_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_user_data_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'), + ) diff --git a/common/djangoapps/student/tests/factories.py b/common/djangoapps/student/tests/factories.py index b7eb31ce3f..21aac0a87f 100644 --- a/common/djangoapps/student/tests/factories.py +++ b/common/djangoapps/student/tests/factories.py @@ -21,6 +21,7 @@ from common.djangoapps.student.models import ( PendingEmailChange, Registration, User, + UserAttribute, UserProfile, UserStanding ) @@ -114,6 +115,15 @@ class UserFactory(DjangoModelFactory): # lint-amnesty, pylint: disable=missing- self.groups.add(GroupFactory.simple_generate(create, name=group_name)) # lint-amnesty, pylint: disable=no-member +class UserAttributeFactory(DjangoModelFactory): # lint-amnesty, pylint: disable=missing-class-docstring + class Meta: + model = UserAttribute + + user = factory.SubFactory(UserFactory) + name = factory.Sequence('{}'.format) + value = factory.Sequence('{}'.format) + + class RequestFactoryNoCsrf(RequestFactory): """ RequestFactory, which disables csrf checks.