Merge pull request #560 from ucb-rit/develop

Enable command-line listing of billing IDs and simple merging of LRC users

coldfront/core/billing/management/commands/billing_ids.py

@@ -9,6 +9,7 @@
 from coldfront.core.billing.utils import ProjectUserBillingActivityManager
 from coldfront.core.billing.utils import UserBillingActivityManager
 from coldfront.core.billing.utils.queries import get_billing_activity_from_full_id
+from coldfront.core.billing.utils.queries import get_billing_id_usages
 from coldfront.core.billing.utils.queries import get_or_create_billing_activity_from_full_id
 from coldfront.core.billing.utils.queries import is_billing_id_well_formed
 from coldfront.core.billing.utils.validation import is_billing_id_valid
@@ -34,15 +35,14 @@ def add_arguments(self, parser):
             title='subcommands')
         subparsers.required = True
         self._add_create_subparser(subparsers)
+        self._add_list_subparser(subparsers)
         self._add_set_subparser(subparsers)
 
     def handle(self, *args, **options):
         """Call the handler for the provided subcommand."""
         subcommand = options['subcommand']
-        if subcommand == 'create':
-            self._handle_create(*args, **options)
-        elif subcommand == 'set':
-            self._handle_set(*args, **options)
+        handler = getattr(self, f'_handle_{subcommand}')
+        handler(*args, **options)
 
     @staticmethod
     def _add_create_subparser(parsers):
@@ -52,6 +52,15 @@ def _add_create_subparser(parsers):
         add_ignore_invalid_argument(parser)
         add_argparse_dry_run_argument(parser)
 
+    @staticmethod
+    def _add_list_subparser(parsers):
+        """Add a subparser for the 'list' command."""
+        parser = parsers.add_parser(
+            'list', help='List billing IDs matching filters.')
+        add_billing_id_argument(parser, is_optional=True)
+        add_project_name_argument(parser, is_optional=True)
+        add_username_argument(parser, is_optional=True)
+
     @staticmethod
     def _add_set_subparser(parsers):
         """Add a subparser for the 'set' subcommand."""
@@ -168,6 +177,55 @@ def _handle_create(self, *args, **options):
             self.stdout.write(self.style.SUCCESS(message))
             self.logger.info(message)
 
+    def _handle_list(self, *args, **options):
+        """Handle the 'list' subcommand."""
+        kwargs = {'full_id': None, 'project_obj': None, 'user_obj': None}
+        billing_id = options['billing_id']
+        if billing_id is not None:
+            kwargs['full_id'] = billing_id
+        project_name = options['project_name']
+        if project_name is not None:
+            kwargs['project_obj'] = self._get_project_or_error(project_name)
+        username = options['username']
+        if username is not None:
+            kwargs['user_obj'] = self._get_user_or_error(username)
+        usages = get_billing_id_usages(**kwargs)
+
+        full_id_by_billing_activity_pk = {}
+
+        for allocation_attribute in usages.project_default:
+            pk = int(allocation_attribute.value)
+            if billing_id:
+                full_id = billing_id
+            elif pk in full_id_by_billing_activity_pk:
+                full_id = full_id_by_billing_activity_pk[pk]
+            else:
+                full_id = BillingActivity.objects.get(pk=pk).full_id()
+                full_id_by_billing_activity_pk[pk] = full_id
+            project_name = allocation_attribute.allocation.project.name
+            line = f'project_default,{project_name},{full_id}'
+            self.stdout.write(line)
+
+        for allocation_user_attribute in usages.recharge:
+            pk = int(allocation_user_attribute.value)
+            if billing_id:
+                full_id = billing_id
+            elif pk in full_id_by_billing_activity_pk:
+                full_id = full_id_by_billing_activity_pk[pk]
+            else:
+                full_id = BillingActivity.objects.get(pk=pk).full_id()
+                full_id_by_billing_activity_pk[pk] = full_id
+            project_name = allocation_user_attribute.allocation.project.name
+            username = allocation_user_attribute.allocation_user.user.username
+            line = f'recharge,{project_name},{username},{full_id}'
+            self.stdout.write(line)
+
+        for user_profile in usages.user_account:
+            full_id = user_profile.billing_activity.full_id()
+            username = user_profile.user.username
+            line = f'user_account,{username},{full_id}'
+            self.stdout.write(line)
+
     def _handle_set(self, *args, **options):
         """Handle the 'set' subcommand."""
         billing_activity = self._get_billing_activity_or_error(
@@ -272,11 +330,12 @@ def _validate_billing_id(self, billing_id, invalid_allowed=False):
                 raise CommandError(message)
 
 
-def add_billing_id_argument(parser):
+def add_billing_id_argument(parser, is_optional=False):
     """Add an argument 'billing_id' to the given argparse parser to
-    accept a billing ID."""
-    parser.add_argument(
-        'billing_id', help='A billing ID (e.g., 123456-789).', type=str)
+    accept a billing ID. Optionally make it an option rather than a
+    positional argument."""
+    name = int(is_optional) * '--' + 'billing_id'
+    parser.add_argument(name, help='A billing ID (e.g., 123456-789).', type=str)
 
 
 def add_ignore_invalid_argument(parser):
@@ -289,17 +348,20 @@ def add_ignore_invalid_argument(parser):
         help='Allow the billing ID to be invalid.')
 
 
-def add_project_name_argument(parser):
+def add_project_name_argument(parser, is_optional=False):
     """Add an argument 'project_name' to the given argparse parser to
-    accept the name of a Project."""
-    parser.add_argument(
-        'project_name', help='The name of a project.', type=str)
+    accept the name of a Project. Optionally make it an option rather
+    than a positional argument."""
+    name = int(is_optional) * '--' + 'project_name'
+    parser.add_argument(name, help='The name of a project.', type=str)
 
 
-def add_username_argument(parser):
+def add_username_argument(parser, is_optional=False):
     """Add an argument 'username' to the given argparse parser to accept
-    the username of a User."""
-    parser.add_argument('username', help='The username of a user.', type=str)
+    the username of a User. Optionally make it an option rather than a
+    positional argument."""
+    name = int(is_optional) * '--' + 'username'
+    parser.add_argument(name, help='The username of a user.', type=str)
 
 
 class Entity(object):

coldfront/core/billing/tests/test_commands/test_billing_ids.py

@@ -98,6 +98,8 @@ def test_create_success(self):
         billing_activity = get_billing_activity_from_full_id(billing_id)
         self.assertTrue(isinstance(billing_activity, BillingActivity))
 
+    # TODO: test_list
+
     def test_set_billing_id_invalid(self):
         """Test that, when the given billing ID is invalid, each of the
         subcommands of the 'set' subcommand raises an error, unless the

coldfront/core/billing/utils/queries.py

@@ -1,14 +1,21 @@
 import logging
 import re
 
+from collections import namedtuple
+
+from django.contrib.auth.models import User
+
 from flags.state import flag_enabled
 
+from coldfront.core.allocation.models import AllocationAttribute
 from coldfront.core.allocation.models import AllocationAttributeType
+from coldfront.core.allocation.models import AllocationUserAttribute
 from coldfront.core.allocation.utils import get_project_compute_allocation
 from coldfront.core.billing.models import BillingActivity
 from coldfront.core.billing.models import BillingProject
 from coldfront.core.project.models import Project
 from coldfront.core.resource.utils import get_computing_allowance_project_prefixes
+from coldfront.core.user.models import UserProfile
 
 
 logger = logging.getLogger(__name__)
@@ -31,6 +38,84 @@ def get_billing_activity_from_full_id(full_id):
     return billing_activity
 
 
+def get_billing_id_usages(full_id=None, project_obj=None, user_obj=None):
+    """Return all database objects storing billing IDs, with optional
+    filtering for a specific ID, for IDs associated with a specific
+    Project, or for IDs associated with a specific User.
+
+    Parameters:
+        - full_id (str): A fully-formed billing ID
+        - project_obj (Project): A Project instance
+        - user_obj (User): A User instance
+
+    Returns:
+        - BillingIdUsages (namedtuple): A namedtuple with the following
+          keys and values:
+              project_default: An AllocationAttribute queryset
+              recharge: An AllocationUserAttribute queryset
+              user_account: A UserProfile queryset
+    """
+    BillingIdUsages = namedtuple(
+        'BillingIdUsages', 'project_default recharge user_account')
+    output_kwargs = {
+        'project_default': AllocationAttribute.objects.none(),
+        'recharge': AllocationUserAttribute.objects.none(),
+        'user_account': UserProfile.objects.none(),
+    }
+
+    allocation_attribute_type = AllocationAttributeType.objects.get(
+        name='Billing Activity')
+
+    allocation_attribute_kwargs = {}
+    allocation_user_attribute_kwargs = {}
+    user_profile_kwargs = {}
+
+    if full_id is not None:
+        assert isinstance(full_id, str)
+        if not is_billing_id_well_formed(full_id):
+            return BillingIdUsages(**output_kwargs)
+        billing_activity = get_billing_activity_from_full_id(full_id)
+        if billing_activity is None:
+            return BillingIdUsages(**output_kwargs)
+        pk_str = str(billing_activity.pk)
+        allocation_attribute_kwargs['value'] = pk_str
+        allocation_user_attribute_kwargs['value'] = pk_str
+        user_profile_kwargs['billing_activity'] = billing_activity
+    if project_obj is not None:
+        assert isinstance(project_obj, Project)
+        allocation_attribute_kwargs['allocation__project'] = project_obj
+        allocation_user_attribute_kwargs['allocation__project'] = project_obj
+    if user_obj is not None:
+        assert isinstance(user_obj, User)
+        allocation_user_attribute_kwargs['allocation_user__user'] = user_obj
+        user_profile_kwargs['user'] = user_obj
+
+    if allocation_attribute_kwargs:
+        allocation_attributes = AllocationAttribute.objects.filter(
+            allocation_attribute_type=allocation_attribute_type,
+            **allocation_attribute_kwargs).order_by('id')
+    else:
+        allocation_attributes = AllocationAttribute.objects.none()
+    output_kwargs['project_default'] = allocation_attributes
+
+    if allocation_user_attribute_kwargs:
+        allocation_user_attributes = AllocationUserAttribute.objects.filter(
+            allocation_attribute_type=allocation_attribute_type,
+            **allocation_user_attribute_kwargs).order_by('id')
+    else:
+        allocation_user_attributes = AllocationUserAttribute.objects.none()
+    output_kwargs['recharge'] = allocation_user_attributes
+
+    if user_profile_kwargs:
+        user_profiles = UserProfile.objects.filter(
+            **user_profile_kwargs).order_by('id')
+    else:
+        user_profiles = UserProfile.objects.none()
+    output_kwargs['user_account'] = user_profiles
+
+    return BillingIdUsages(**output_kwargs)
+
+
 def get_or_create_billing_activity_from_full_id(full_id):
     """Given a fully-formed billing ID, get or create a matching
     BillingActivity, creating a BillingProject as needed."""

coldfront/core/user/utils_/merge_users/class_handlers.py

@@ -153,15 +153,64 @@ def _get_settable_if_falsy_attrs(self):
         ]
 
     def _run_special_handling(self):
-        self._set_host_user()
+        if flag_enabled('LRC_ONLY'):
+            # TODO: Refactor shared logic between these methods.
+            self._set_host_user()
+            self._set_billing_activity()
+
+    def _set_billing_activity(self):
+        attr_name = 'billing_activity'
+        src_activity = self._src_obj.billing_activity
+        dst_activity = self._dst_obj.billing_activity
+        if src_activity and dst_activity and src_activity != dst_activity:
+            src_full_id = src_activity.full_id()
+            dst_full_id = dst_activity.full_id()
+            prompt = (
+                f'{self._class_name}({self._dst_obj.pk}).{attr_name}: Conflict '
+                f'requiring manual resolution. Type 1 or 2 to keep the '
+                f'corresponding value.\n'
+                f'1 - {src_full_id} (Source)\n'
+                f'2 - {dst_full_id} (Destination)\n')
+            choice = input(prompt)
+            if choice == '1':
+                self._dst_obj.billing_activity = src_activity
+                self._record_update(
+                    self._dst_obj.pk, attr_name, dst_full_id, src_full_id)
+            elif choice == '2':
+                pass
+            else:
+                raise ValueError('Invalid choice.')
+        else:
+            self._set_attr_if_falsy(attr_name)
 
     def _set_host_user(self):
-        if flag_enabled('LRC_ONLY'):
-            raise NotImplementedError
-            # TODO
-            #  Deal with conflicts.
-            #  Handle LBL users.
-            # self._set_attr_if_falsy('host_user')
+        attr_name = 'host_user'
+        src_host = self._src_obj.host_user
+        dst_host = self._dst_obj.host_user
+        if src_host and dst_host and src_host != dst_host:
+            src_user_str = (
+                f'{src_host.username} ({src_host.pk}, {src_host.first_name} '
+                f'{src_host.last_name})')
+            dst_user_str = (
+                f'{dst_host.username} ({dst_host.pk}, {dst_host.first_name} '
+                f'{dst_host.last_name})')
+            prompt = (
+                f'{self._class_name}({self._dst_obj.pk}).{attr_name}: Conflict '
+                f'requiring manual resolution. Type 1 or 2 to keep the '
+                f'corresponding value.\n'
+                f'1 - {src_user_str} (Source)\n'
+                f'2 - {dst_user_str} (Destination)\n')
+            choice = input(prompt)
+            if choice == '1':
+                self._dst_obj.host_user = src_host
+                self._record_update(
+                    self._dst_obj.pk, attr_name, dst_host, src_host)
+            elif choice == '2':
+                pass
+            else:
+                raise ValueError('Invalid choice.')
+        else:
+            self._set_attr_if_falsy(attr_name)
 
 
 class SocialAccountHandler(ClassHandler):