@@ -253,7 +253,7 @@ def on_user_creation_signal(sender, instance, created, **kwargs): from payments.services.payment_plan_service import PaymentPlanService if created: - free_plan = PaymentPlanSelector(instance).get_free_plan() + free_plan = PaymentPlanSelector(instance.pk).get_free_plan() PaymentPlanService(instance).subscribe_user_to_plan(free_plan) if not instance.profile_picture_name: instance.profile_picture_name = instance.LAST_NAME_AVATARS.get(instance.last_name, None) @@ -1,21 +1,18 @@ from authentication.models import ( - BusinessAccount, - BusinessUserHost, CustomUserModel, ) -from authentication.models.choices import AccountPrivileges class AccountStatusSelector: def __init__(self, user: CustomUserModel): self.user = user - def is_business_host(self) -> bool: - return BusinessUserHost.objects.filter(user=self.user).exists() + return getattr(self.user, 'host_account', None) is not None def is_business_account(self) -> bool: - return BusinessAccount.objects.filter(user=self.user).exists() + return getattr(self.user, 'business_account', None) is not None def is_admin(self) -> bool: - accounts = BusinessAccount.objects.filter(user=self.user) - return not accounts.exists() or (accounts.first().account_privileges == AccountPrivileges.ADMIN) + if getattr(self.user, 'business_account', None) is None: + return False + return self.user.business_account.account_privileges == 'admin' @@ -11,13 +11,10 @@ class BusinessAccountSelector: @classmethod def from_user(cls, user: CustomUserModel, company: BusinessUserHost | None = None): - account = BusinessAccount.objects.filter(user=user) - if company is not None: - account.filter(parent_company=company) - if not account.exists(): + account = getattr(user, 'business_account', None) + if not account or (company is not None and account.parent_company != company): return None - - return cls(account.first()) + return cls(account) @classmethod def filter_by_email(cls, email: str, company: BusinessUserHost | None = None): @@ -99,12 +99,14 @@ class UserSelector: return 'regular' account = BusinessAccountSelector.from_user(self.user) - if account.account_type() == 'admin': - return 'business_admin' - elif account.account_type() == 'regular': - return 'business_account' - elif account.account_type() == 'sec': - return 'business_security' + if account: + if account.account_type() == 'admin': + return 'business_admin' + elif account.account_type() == 'regular': + return 'business_account' + elif account.account_type() == 'sec': + return 'business_security' + return 'regular' def check_model_availability(self, model_title: str) -> bool: user_type = self.check_account_type() @@ -94,7 +94,7 @@ class BusinessHostService: raise AlreadyHost(user) if AccountStatusSelector(user).is_business_account(): raise AlreadyAccount(user) - if PaymentPlanSelector(user).is_plan_paid(): + if PaymentPlanSelector(user.pk).is_plan_paid(): raise AlreadyHasPlan(user) except CustomUserModel.DoesNotExist: return self.create_account( @@ -206,7 +206,7 @@ class BusinessHostService: serializer = NewBusinessHostSerializer(data=request.data) serializer.is_valid(raise_exception=True) - if PaymentPlanSelector(self.user).is_plan_paid(): + if PaymentPlanSelector(self.user.pk).is_plan_paid(): raise business_host_exceptions.AlreadyHasPlan(self.user) if AccountStatusSelector(self.user).is_business_host(): @@ -488,7 +488,7 @@ class Chatgpt(SimpleService): model: str = 'gpt-3.5-turbo', embedding_tokens: int = 0 ): - balance = PaymentPlanSelector(self.store.user).get_current_balance() + balance = PaymentPlanSelector(self.store.user.pk).get_current_balance() total_tokens = input_tokens if image_size: total_tokens += self.count_image_tokens(image_size) @@ -19,8 +19,17 @@ logger = logging.getLogger(__name__) class PaymentPlanSelector: - def __init__(self, user: CustomUserModel): - self.user = user + def __init__(self, user_pk: int): + self.user_pk = user_pk + + @property + def user(self) -> CustomUserModel: + return CustomUserModel.objects.select_related( + 'payment_plan', + 'business_account', + 'business_account__group', + 'business_account__parent_company__user__payment_plan' + ).get(pk=self.user_pk) def get_payment_plans( self, @@ -61,10 +70,9 @@ class PaymentPlanSelector: def get_current_balance(self) -> Decimal: user_type = UserSelector(self.user).check_account_type() if ( - user_type == 'business_account' - or user_type == 'business_admin' - or user_type == 'business_security' - ) and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED: + user_type in {'business_account', 'business_admin', 'business_security'} + and self.user.business_account.acceptance_status == InvitationStatus.ACCEPTED + ): balance = ( self.user.business_account.group.token_limit if self.user.business_account.group @@ -39,7 +39,7 @@ class PaymentPlanService: def create_payment_plan_invoice(self, payment_plan_uid: str): """""" - payment_plan = PaymentPlanSelector(self.user).get_payment_plan_by_id(payment_plan_uid) + payment_plan = PaymentPlanSelector(self.user.pk).get_payment_plan_by_id(payment_plan_uid) resulting_link = self.payment_service.create_payment_link(plan=payment_plan) result = PaymentLinkSerializer(data={'payment_url': resulting_link}) result.is_valid(raise_exception=True) @@ -72,7 +72,7 @@ class PaymentPlanService: task = plan_info.plan_schedule if task is not None: RecurrentPaymentService(task).delete() - plan_info.plan = PaymentPlanSelector(self.user).get_free_plan(corporate=self.user.is_corporate()) + plan_info.plan = PaymentPlanSelector(self.user.pk).get_free_plan(corporate=self.user.is_corporate()) plan_info.save() def update_per_token_plan_details(self, payment_amount: Decimal, model=None): @@ -68,7 +68,7 @@ class PaymentPlanAPIView(APIView): """List available payment plans""" try: plan_duration = request.query_params.get('duration', None) - result = PaymentPlanSelector(self.request.user).get_payment_plans(plan_duration=plan_duration) + result = PaymentPlanSelector(self.request.user.pk).get_payment_plans(plan_duration=plan_duration) return Response(result.data, status=status.HTTP_200_OK) except Exception as err: logger.exception(err) @@ -146,7 +146,7 @@ class UserPlanAPIView(APIView): def get(self, request, *args, **kwargs): """Get user balance""" try: - result = PaymentPlanSelector(self.request.user).get_current_balance() + result = PaymentPlanSelector(self.request.user.pk).get_current_balance() return Response({'current_token_balance': result}, status=status.HTTP_200_OK) except Exception as err: logger.exception(err)