remove the BillingPlan enum and references to PlanData in favor of Plan properties or PlanName

Author: ajay-sentry

graphql_api/types/owner/owner.py

@@ -8,7 +8,7 @@
 from ariadne import ObjectType
 from django.conf import settings
 from graphql import GraphQLResolveInfo
-from shared.plan.constants import DEFAULT_FREE_PLAN, PlanData, convert_to_DTO
+from shared.plan.constants import DEFAULT_FREE_PLAN
 from shared.plan.service import PlanService
 
 import services.activation as activation
@@ -111,16 +111,16 @@ def resolve_plan(owner: Owner, info: GraphQLResolveInfo) -> PlanService:
 @owner_bindable.field("pretrialPlan")
 @require_part_of_org
 @sync_to_async
-def resolve_plan_representation(owner: Owner, info: GraphQLResolveInfo) -> PlanData:
+def resolve_plan_representation(owner: Owner, info: GraphQLResolveInfo) -> Plan:
     info.context["plan_service"] = PlanService(current_org=owner)
     free_plan = Plan.objects.select_related("tier").get(name=DEFAULT_FREE_PLAN)
-    return convert_to_DTO(free_plan)
+    return free_plan
 
 
 @owner_bindable.field("availablePlans")
 @require_part_of_org
 @sync_to_async
-def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[PlanData]:
+def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[Plan]:
     plan_service = PlanService(current_org=owner)
     info.context["plan_service"] = plan_service
     owner = info.context["request"].current_owner
@@ -130,7 +130,7 @@ def resolve_available_plans(owner: Owner, info: GraphQLResolveInfo) -> List[Plan
 @owner_bindable.field("hasPrivateRepos")
 @sync_to_async
 @require_part_of_org
-def resolve_has_private_repos(owner: Owner, info: GraphQLResolveInfo) -> List[PlanData]:
+def resolve_has_private_repos(owner: Owner, info: GraphQLResolveInfo) -> bool:
     return owner.has_private_repos
 
 

graphql_api/types/plan_representation/plan_representation.py

@@ -1,10 +1,10 @@
 from typing import List, Optional
 
 from ariadne import ObjectType
-from shared.plan.constants import PlanData
 from shared.plan.service import PlanService
 
 from codecov.db import sync_to_async
+from codecov_auth.models import Plan
 from graphql_api.helpers.ariadne import ariadne_load_local_graphql
 
 plan_representation = ariadne_load_local_graphql(
@@ -14,70 +14,70 @@
 
 
 @plan_representation_bindable.field("marketingName")
-def resolve_marketing_name(plan_data: PlanData, info) -> str:
-    return plan_data["marketing_name"]
+def resolve_marketing_name(plan_data: Plan, info) -> str:
+    return plan_data.marketing_name
 
 
 @plan_representation_bindable.field("value")
-def resolve_plan_value(plan_data: PlanData, info) -> str:
-    return plan_data["value"]
+def resolve_plan_value(plan_data: Plan, info) -> str:
+    return plan_data.name
 
 
 @plan_representation_bindable.field("billingRate")
-def resolve_billing_rate(plan_data: PlanData, info) -> Optional[str]:
-    return plan_data["billing_rate"]
+def resolve_billing_rate(plan_data: Plan, info) -> Optional[str]:
+    return plan_data.billing_rate
 
 
 @plan_representation_bindable.field("baseUnitPrice")
-def resolve_base_unit_price(plan_data: PlanData, info) -> int:
-    return plan_data["base_unit_price"]
+def resolve_base_unit_price(plan_data: Plan, info) -> int:
+    return plan_data.base_unit_price
 
 
 @plan_representation_bindable.field("benefits")
 @sync_to_async
-def resolve_benefits(plan_data: PlanData, info) -> List[str]:
+def resolve_benefits(plan_data: Plan, info) -> List[str]:
     plan_service: PlanService = info.context["plan_service"]
     if plan_service.is_org_trialing:
         benefits_with_pretrial_users = [
             benefit.replace(
                 "Up to 1 user", f"Up to {plan_service.pretrial_users_count} users"
             )
-            for benefit in plan_data["benefits"]
+            for benefit in plan_data.benefits
         ]
         return benefits_with_pretrial_users
-    return plan_data["benefits"]
+    return plan_data.benefits
 
 
 @plan_representation_bindable.field("monthlyUploadLimit")
-def resolve_monthly_uploads_limit(plan_data: PlanData, info) -> Optional[int]:
-    return plan_data["monthly_uploads_limit"]
+def resolve_monthly_uploads_limit(plan_data: Plan, info) -> Optional[int]:
+    return plan_data.monthly_uploads_limit
 
 
 @plan_representation_bindable.field("isEnterprisePlan")
-def resolve_is_enterprise(plan_data: PlanData, info) -> bool:
-    return plan_data["is_enterprise_plan"]
+def resolve_is_enterprise(plan_data: Plan, info) -> bool:
+    return plan_data.is_enterprise_plan
 
 
 @plan_representation_bindable.field("isFreePlan")
-def resolve_is_free(plan_data: PlanData, info) -> bool:
-    return plan_data["is_free_plan"]
+def resolve_is_free(plan_data: Plan, info) -> bool:
+    return plan_data.is_free_plan
 
 
 @plan_representation_bindable.field("isProPlan")
-def resolve_is_pro(plan_data: PlanData, info) -> bool:
-    return plan_data["is_pro_plan"]
+def resolve_is_pro(plan_data: Plan, info) -> bool:
+    return plan_data.is_pro_plan
 
 
 @plan_representation_bindable.field("isTeamPlan")
-def resolve_is_team(plan_data: PlanData, info) -> bool:
-    return plan_data["is_team_plan"]
+def resolve_is_team(plan_data: Plan, info) -> bool:
+    return plan_data.is_team_plan
 
 
 @plan_representation_bindable.field("isSentryPlan")
-def resolve_is_sentry(plan_data: PlanData, info) -> bool:
-    return plan_data["is_sentry_plan"]
+def resolve_is_sentry(plan_data: Plan, info) -> bool:
+    return plan_data.is_sentry_plan
 
 
 @plan_representation_bindable.field("isTrialPlan")
-def resolve_is_trial(plan_data: PlanData, info) -> bool:
-    return plan_data["is_trial_plan"]
+def resolve_is_trial(plan_data: Plan, info) -> bool:
+    return plan_data.is_trial_plan

services/task/task_router.py

@@ -1,6 +1,6 @@
 import shared.celery_config as shared_celery_config
-from shared.billing import BillingPlan
 from shared.celery_router import route_tasks_based_on_user_plan
+from shared.plan.constants import DEFAULT_FREE_PLAN
 
 from codecov_auth.models import Owner
 from compare.models import CommitComparison
@@ -14,14 +14,14 @@ def _get_user_plan_from_ownerid(ownerid, *args, **kwargs) -> str:
     owner = Owner.objects.filter(ownerid=ownerid).first()
     if owner:
         return owner.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_repoid(repoid, *args, **kwargs) -> str:
     repo = Repository.objects.filter(repoid=repoid).first()
     if repo and repo.author:
         return repo.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_profiling_commit(profiling_id, *args, **kwargs) -> str:
@@ -32,7 +32,7 @@ def _get_user_plan_from_profiling_commit(profiling_id, *args, **kwargs) -> str:
         and profiling_commit.repository.author
     ):
         return profiling_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_profiling_upload(profiling_upload_id, *args, **kwargs) -> str:
@@ -44,7 +44,7 @@ def _get_user_plan_from_profiling_upload(profiling_upload_id, *args, **kwargs) -
         and profiling_upload.profiling_commit.repository.author
     ):
         return profiling_upload.profiling_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_comparison_id(comparison_id, *args, **kwargs) -> str:
@@ -60,7 +60,7 @@ def _get_user_plan_from_comparison_id(comparison_id, *args, **kwargs) -> str:
         and compare_commit.compare_commit.repository.author
     ):
         return compare_commit.compare_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_label_request_id(request_id, *args, **kwargs) -> str:
@@ -76,7 +76,7 @@ def _get_user_plan_from_label_request_id(request_id, *args, **kwargs) -> str:
         and label_analysis_request.head_commit.repository.author
     ):
         return label_analysis_request.head_commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_suite_id(suite_id, *args, **kwargs) -> str:
@@ -92,7 +92,7 @@ def _get_user_plan_from_suite_id(suite_id, *args, **kwargs) -> str:
         and static_analysis_suite.commit.repository.author
     ):
         return static_analysis_suite.commit.repository.author.plan
-    return BillingPlan.users_basic.db_name
+    return DEFAULT_FREE_PLAN
 
 
 def _get_user_plan_from_task(task_name: str, task_kwargs: dict) -> str:
@@ -119,7 +119,7 @@ def _get_user_plan_from_task(task_name: str, task_kwargs: dict) -> str:
         shared_celery_config.static_analysis_task_name: _get_user_plan_from_suite_id,
     }
     func_to_use = owner_plan_lookup_funcs.get(
-        task_name, lambda *args, **kwargs: BillingPlan.users_basic.db_name
+        task_name, lambda *args, **kwargs: DEFAULT_FREE_PLAN
     )
     return func_to_use(**task_kwargs)