From 089926616ab60c6f09b68b9c0ffe30b78ce93c09 Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Fri, 11 Aug 2023 11:34:56 -0700 Subject: [PATCH 1/2] Add billing query for getting all usages of an ID; add 'list' subcommand for billing_ids command --- .../management/commands/billing_ids.py | 92 ++++++++++++++++--- .../tests/test_commands/test_billing_ids.py | 2 + coldfront/core/billing/utils/queries.py | 85 +++++++++++++++++ 3 files changed, 164 insertions(+), 15 deletions(-) diff --git a/coldfront/core/billing/management/commands/billing_ids.py b/coldfront/core/billing/management/commands/billing_ids.py index 229e44670..be533ccea 100644 --- a/coldfront/core/billing/management/commands/billing_ids.py +++ b/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): diff --git a/coldfront/core/billing/tests/test_commands/test_billing_ids.py b/coldfront/core/billing/tests/test_commands/test_billing_ids.py index b310385ec..a36c48c33 100644 --- a/coldfront/core/billing/tests/test_commands/test_billing_ids.py +++ b/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 diff --git a/coldfront/core/billing/utils/queries.py b/coldfront/core/billing/utils/queries.py index 552cb4f03..cae62f125 100644 --- a/coldfront/core/billing/utils/queries.py +++ b/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.""" From 4d5cc4c78e97a0557bb212c777b6d55a49652cff Mon Sep 17 00:00:00 2001 From: Matthew Li Date: Tue, 15 Aug 2023 10:24:37 -0700 Subject: [PATCH 2/2] Update merging logic to handle UserProfile billing_activity and host_user on LRC --- .../user/utils_/merge_users/class_handlers.py | 63 ++++++++++++++++--- 1 file changed, 56 insertions(+), 7 deletions(-) diff --git a/coldfront/core/user/utils_/merge_users/class_handlers.py b/coldfront/core/user/utils_/merge_users/class_handlers.py index 6e3e36a0a..81b25ed11 100644 --- a/coldfront/core/user/utils_/merge_users/class_handlers.py +++ b/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):