Refactor/benefits

server/polar/advertisement/endpoints.py

@@ -4,7 +4,7 @@
 from pydantic import UUID4
 
 from polar.benefit.schemas import BenefitID
-from polar.benefit.service.benefit import benefit as benefit_service
+from polar.benefit.service import benefit as benefit_service
 from polar.exceptions import PolarRequestValidationError, ResourceNotFound
 from polar.kit.pagination import ListResource, PaginationParamsQuery
 from polar.kit.sorting import Sorting, SortingGetter

server/polar/benefit/endpoints.py

@@ -15,6 +15,7 @@
 from polar.routing import APIRouter
 
 from . import auth
+from .grant.service import benefit_grant as benefit_grant_service
 from .schemas import Benefit as BenefitSchema
 from .schemas import (
     BenefitCreate,
@@ -23,8 +24,7 @@
     BenefitUpdate,
     benefit_schema_map,
 )
-from .service.benefit import benefit as benefit_service
-from .service.benefit_grant import benefit_grant as benefit_grant_service
+from .service import benefit as benefit_service
 
 router = APIRouter(prefix="/benefits", tags=["benefits", APITag.documented])
 

server/polar/benefit/service/benefit_grant.py → server/polar/benefit/grant/service.py

@@ -6,9 +6,9 @@
 from sqlalchemy import select
 from sqlalchemy.orm import joinedload
 
-from polar.benefit.benefits import get_benefit_service
-from polar.benefit.benefits.base import BenefitActionRequiredError
+from polar.benefit.registry import get_benefit_strategy
 from polar.benefit.schemas import BenefitGrantWebhook
+from polar.benefit.strategies.base import BenefitActionRequiredError
 from polar.customer.service import customer as customer_service
 from polar.eventstream.service import publish as eventstream_publish
 from polar.exceptions import PolarError
@@ -25,7 +25,7 @@
 from polar.webhook.webhooks import WebhookPayloadTypeAdapter
 from polar.worker import enqueue_job
 
-from .benefit_grant_scope import scope_to_args
+from .scope import scope_to_args
 
 log: Logger = structlog.get_logger()
 
@@ -149,9 +149,9 @@ async def grant_benefit(
             return grant
 
         previous_properties = grant.properties
-        benefit_service = get_benefit_service(benefit.type, session, redis)
+        benefit_strategy = get_benefit_strategy(benefit.type, session, redis)
         try:
-            properties = await benefit_service.grant(
+            properties = await benefit_strategy.grant(
                 benefit,
                 customer,
                 grant.properties,
@@ -214,17 +214,17 @@ async def revoke_benefit(
 
         previous_properties = grant.properties
 
-        benefit_service = get_benefit_service(benefit.type, session, redis)
+        benefit_strategy = get_benefit_strategy(benefit.type, session, redis)
         # Call the revoke logic in two cases:
         # * If the service requires grants to be revoked individually
         # * If there is only one grant remaining for this benefit,
         #   so the benefit remains if other grants exist via other purchases
         other_grants = await self._get_granted_by_benefit_and_customer(
             session, benefit, customer
         )
-        if benefit_service.should_revoke_individually or len(other_grants) < 2:
+        if benefit_strategy.should_revoke_individually or len(other_grants) < 2:
             try:
-                properties = await benefit_service.revoke(
+                properties = await benefit_strategy.revoke(
                     benefit,
                     customer,
                     grant.properties,
@@ -296,8 +296,8 @@ async def enqueue_benefit_grant_updates(
         benefit: Benefit,
         previous_properties: BenefitProperties,
     ) -> None:
-        benefit_service = get_benefit_service(benefit.type, session, redis)
-        if not await benefit_service.requires_update(benefit, previous_properties):
+        benefit_strategy = get_benefit_strategy(benefit.type, session, redis)
+        if not await benefit_strategy.requires_update(benefit, previous_properties):
             return
 
         grants = await self._get_granted_by_benefit(session, benefit)
@@ -322,9 +322,9 @@ async def update_benefit_grant(
         assert customer is not None
 
         previous_properties = grant.properties
-        benefit_service = get_benefit_service(benefit.type, session, redis)
+        benefit_strategy = get_benefit_strategy(benefit.type, session, redis)
         try:
-            properties = await benefit_service.grant(
+            properties = await benefit_strategy.grant(
                 benefit,
                 customer,
                 grant.properties,
@@ -374,8 +374,8 @@ async def delete_benefit_grant(
         assert customer is not None
 
         previous_properties = grant.properties
-        benefit_service = get_benefit_service(benefit.type, session, redis)
-        properties = await benefit_service.revoke(
+        benefit_strategy = get_benefit_strategy(benefit.type, session, redis)
+        properties = await benefit_strategy.revoke(
             benefit,
             customer,
             grant.properties,

server/polar/benefit/registry.py

@@ -0,0 +1,87 @@
+from typing import Any
+
+from polar.models import Benefit
+from polar.models.benefit import (
+    BenefitProperties,
+    BenefitType,
+)
+from polar.models.benefit_grant import BenefitGrantPropertiesBase
+from polar.postgres import AsyncSession
+from polar.redis import Redis
+
+from .strategies.ads.service import BenefitAdsService
+from .strategies.base import (
+    BenefitActionRequiredError,
+    BenefitPropertiesValidationError,
+    BenefitRetriableError,
+    BenefitServiceError,
+    BenefitServiceProtocol,
+)
+from .strategies.custom.service import BenefitCustomService
+from .strategies.discord.service import BenefitDiscordService
+from .strategies.downloadables.service import BenefitDownloadablesService
+from .strategies.github_repository.service import BenefitGitHubRepositoryService
+from .strategies.license_keys.service import BenefitLicenseKeysService
+
+# class Config(NamedTuple):
+#     service: type[BenefitServiceProtocol[Any, Any, Any]]
+#     tax: bool
+#
+#
+# class BenefitType(StrEnum):
+#     custom = "custom"
+#     ads = "ads"
+#     discord = "discord"
+#     github_repository = "github_repository"
+#     downloadables = "downloadables"
+#     license_keys = "license_keys"
+#
+#     @classmethod
+#     @functools.cache
+#     def mapping(cls) -> dict["BenefitType", Config]:
+#         return {
+#             cls.custom: Config(BenefitCustomService, tax=True),
+#             cls.ads: Config(BenefitAdsService, tax=True),
+#             cls.discord: Config(BenefitDiscordService, tax=True),
+#             cls.github_repository: Config(BenefitGitHubRepositoryService, tax=True),
+#             cls.downloadables: Config(BenefitDownloadablesService, tax=True),
+#             cls.license_keys: Config(BenefitLicenseKeysService, tax=True),
+#         }
+#
+#     def is_tax_applicable(self) -> bool:
+#         return self.mapping()[self].tax
+#
+#     @classmethod
+#     def get_service(
+#         cls, type: "BenefitType", session: AsyncSession, redis: Redis
+#     ) -> BenefitServiceProtocol[Benefit, BenefitProperties, BenefitGrantPropertiesBase]:
+#         return cls.mapping()[type].service(session, redis)
+#
+
+_STRATEGY_CLASS_MAP: dict[
+    BenefitType,
+    type[BenefitServiceProtocol[Any, Any, Any]],
+] = {
+    BenefitType.custom: BenefitCustomService,
+    BenefitType.ads: BenefitAdsService,
+    BenefitType.discord: BenefitDiscordService,
+    BenefitType.github_repository: BenefitGitHubRepositoryService,
+    BenefitType.downloadables: BenefitDownloadablesService,
+    BenefitType.license_keys: BenefitLicenseKeysService,
+}
+
+
+def get_benefit_strategy(
+    type: BenefitType, session: AsyncSession, redis: Redis
+) -> BenefitServiceProtocol[Benefit, BenefitProperties, BenefitGrantPropertiesBase]:
+    return _STRATEGY_CLASS_MAP[type](session, redis)
+
+
+__all__ = [
+    "BenefitActionRequiredError",
+    "BenefitServiceProtocol",
+    "BenefitPropertiesValidationError",
+    "BenefitRetriableError",
+    "BenefitServiceError",
+    "get_benefit_strategy",
+]