refactor: Resume-able Apple migration mgmt cmds

common/djangoapps/third_party_auth/management/commands/generate_and_store_apple_transfer_ids.py

@@ -10,6 +10,7 @@
 
 from django.core.management.base import BaseCommand, CommandError
 from django.db import transaction
+from django.db.models import Q
 import jwt
 from social_django.models import UserSocialAuth
 from social_django.utils import load_strategy
@@ -21,6 +22,12 @@
 log = logging.getLogger(__name__)
 
 
+class AccessTokenExpiredException(Exception):
+    """
+    Raised when access token has been expired.
+    """
+
+
 class Command(BaseCommand):
     """
     Management command to generate transfer identifiers for apple users using their apple_id
@@ -78,37 +85,70 @@ def _generate_access_token(self, client_secret):
         access_token = response.json().get('access_token')
         return access_token
 
-    def add_arguments(self, parser):
-        parser.add_argument('target_team_id', help='Team ID to which the app is to be migrated to.')
+    def _get_token_and_secret(self):
+        """
+        Get access_token and client_secret
+        """
+        client_secret = self._generate_client_secret()
+        access_token = self._generate_access_token(client_secret)
+        return access_token, client_secret
 
-    @transaction.atomic
-    def handle(self, *args, **options):
-        target_team_id = options['target_team_id']
+    def _update_token_and_secret(self):
+        self.access_token, self.client_secret = self._get_token_and_secret()  # pylint: disable=W0201
 
+    def _fetch_transfer_id(self, apple_id, target_team_id):
+        """
+        Fetch Transfer ID for a given Apple ID from Apple API.
+        """
         migration_url = "https://appleid.apple.com/auth/usermigrationinfo"
         app_id = "org.edx.mobile"
-
-        client_secret = self._generate_client_secret()
-        access_token = self._generate_access_token(client_secret)
-        if not access_token:
-            raise CommandError('Failed to create access token.')
-
         headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Host": "appleid.apple.com",
-            "Authorization": "Bearer " + access_token
+            "Authorization": "Bearer " + self.access_token
         }
         payload = {
             "target": target_team_id,
             "client_id": app_id,
-            "client_secret": client_secret
+            "client_secret": self.client_secret,
+            "sub": apple_id
         }
+        response = requests.post(migration_url, data=payload, headers=headers)
+        if response.status_code == 400:
+            raise AccessTokenExpiredException
+
+        return response.json().get('transfer_sub')
+
+    def _get_transfer_id_for_apple_id(self, apple_id, target_team_id):
+        """
+        Given an Apple ID from the old transferring team,
+        create and return its respective transfer id.
+        """
+        try:
+            transfer_id = self._fetch_transfer_id(apple_id, target_team_id)
+        except AccessTokenExpiredException:
+            log.info('Access token expired. Re-creating access token.')
+            self._update_token_and_secret()
+            transfer_id = self._fetch_transfer_id(apple_id, target_team_id)
+        return transfer_id
+
+    def add_arguments(self, parser):
+        parser.add_argument('target_team_id', help='Team ID to which the app is to be migrated to.')
+
+    @transaction.atomic
+    def handle(self, *args, **options):
+        target_team_id = options['target_team_id']
+
+        self._update_token_and_secret()
+        if not self.access_token:
+            raise CommandError('Failed to create access token.')
 
-        apple_ids = UserSocialAuth.objects.filter(provider=AppleIdAuth.name).values_list('uid', flat=True)
+        already_processed_apple_ids = AppleMigrationUserIdInfo.objects.all().exclude(
+            Q(transfer_id__isnull=True) | Q(transfer_id="")).values_list('old_apple_id', flat=True)
+        apple_ids = UserSocialAuth.objects.filter(provider=AppleIdAuth.name).exclude(
+            uid__in=already_processed_apple_ids).values_list('uid', flat=True)
         for apple_id in apple_ids:
-            payload['sub'] = apple_id
-            response = requests.post(migration_url, data=payload, headers=headers)
-            transfer_id = response.json().get('transfer_sub')
+            transfer_id = self._get_transfer_id_for_apple_id(apple_id, target_team_id)
             if transfer_id:
                 apple_user_id_info, _ = AppleMigrationUserIdInfo.objects.get_or_create(old_apple_id=apple_id)
                 apple_user_id_info.transfer_id = transfer_id

common/djangoapps/third_party_auth/management/commands/generate_and_store_new_apple_ids.py

@@ -10,6 +10,7 @@
 
 from django.core.management.base import BaseCommand, CommandError
 from django.db import transaction
+from django.db.models import Q
 import jwt
 from social_django.utils import load_strategy
 
@@ -19,6 +20,12 @@
 log = logging.getLogger(__name__)
 
 
+class AccessTokenExpiredException(Exception):
+    """
+    Raised when access token has been expired.
+    """
+
+
 class Command(BaseCommand):
     """
     Management command to exchange transfer identifiers for new team-scoped identifier for
@@ -76,31 +83,63 @@ def _generate_access_token(self, client_secret):
         access_token = response.json().get('access_token')
         return access_token
 
-    @transaction.atomic
-    def handle(self, *args, **options):
-        migration_url = "https://appleid.apple.com/auth/usermigrationinfo"
-        app_id = "org.edx.mobile"
-
+    def _get_token_and_secret(self):
+        """
+        Get access_token and client_secret
+        """
         client_secret = self._generate_client_secret()
         access_token = self._generate_access_token(client_secret)
-        if not access_token:
-            raise CommandError('Failed to create access token.')
+        return access_token, client_secret
+
+    def _update_token_and_secret(self):
+        self.access_token, self.client_secret = self._get_token_and_secret()  # pylint: disable=W0201
 
+    def _fetch_new_apple_id(self, transfer_id):
+        """
+        Fetch Apple ID for a given transfer ID from Apple API.
+        """
+        migration_url = "https://appleid.apple.com/auth/usermigrationinfo"
+        app_id = "org.edx.mobile"
         headers = {
             "Content-Type": "application/x-www-form-urlencoded",
             "Host": "appleid.apple.com",
-            "Authorization": "Bearer " + access_token
+            "Authorization": "Bearer " + self.access_token
         }
         payload = {
             "client_id": app_id,
-            "client_secret": client_secret
+            "client_secret": self.client_secret,
+            "transfer_sub": transfer_id
         }
+        response = requests.post(migration_url, data=payload, headers=headers)
+        if response.status_code == 400:
+            raise AccessTokenExpiredException
+
+        return response.json().get('sub')
+
+    def _exchange_transfer_id_for_new_apple_id(self, transfer_id):
+        """
+        For a Transfer ID obtained from the transferring team,
+        return the correlating Apple ID belonging to the recipient team.
+        """
+        try:
+            new_apple_id = self._fetch_new_apple_id(transfer_id)
+        except AccessTokenExpiredException:
+            log.info('Access token expired. Re-creating access token.')
+            self._update_token_and_secret()
+            new_apple_id = self._fetch_new_apple_id(transfer_id)
+
+        return new_apple_id
+
+    @transaction.atomic
+    def handle(self, *args, **options):
+        self._update_token_and_secret()
+        if not self.access_token:
+            raise CommandError('Failed to create access token.')
 
-        apple_user_ids_info = AppleMigrationUserIdInfo.objects.all()
+        apple_user_ids_info = AppleMigrationUserIdInfo.objects.filter(Q(new_apple_id__isnull=True) | Q(new_apple_id=""),
+                                                                      ~Q(transfer_id=""), transfer_id__isnull=False)
         for apple_user_id_info in apple_user_ids_info:
-            payload['transfer_sub'] = apple_user_id_info.transfer_id
-            response = requests.post(migration_url, data=payload, headers=headers)
-            new_apple_id = response.json().get('sub')
+            new_apple_id = self._exchange_transfer_id_for_new_apple_id(apple_user_id_info.transfer_id)
             if new_apple_id:
                 apple_user_id_info.new_apple_id = new_apple_id
                 apple_user_id_info.save()