server: make Redis injected everywhere

Stop having a global instance that were causing problems in tests. We even have been able to replace it by fakeredis!

server/polar/app.py

@@ -40,6 +40,7 @@
 from polar.openapi import OPENAPI_PARAMETERS, APITag, set_openapi_generator
 from polar.postgres import create_async_engine, create_sync_engine
 from polar.posthog import configure_posthog
+from polar.redis import Redis, create_redis
 from polar.sentry import configure_sentry
 from polar.webhook.webhooks import document_webhooks
 from polar.worker import ArqRedis
@@ -90,6 +91,7 @@ class State(TypedDict):
     sync_engine: Engine
     sync_sessionmaker: SyncSessionMaker
     arq_pool: ArqRedis
+    redis: Redis
     ip_geolocation_client: ip_geolocation.IPGeolocationClient | None
 
 
@@ -98,40 +100,42 @@ async def lifespan(app: FastAPI) -> AsyncIterator[State]:
     log.info("Starting Polar API")
 
     async with worker_lifespan() as arq_pool:
-        async_engine = create_async_engine("app")
-        async_sessionmaker = create_async_sessionmaker(async_engine)
-        instrument_sqlalchemy(async_engine.sync_engine)
-
-        sync_engine = create_sync_engine("app")
-        sync_sessionmaker = create_sync_sessionmaker(sync_engine)
-        instrument_sqlalchemy(sync_engine)
-
-        try:
-            ip_geolocation_client = ip_geolocation.get_client()
-        except FileNotFoundError:
-            log.info(
-                "IP geolocation database not found. "
-                "Checkout won't automatically geolocate IPs."
-            )
-            ip_geolocation_client = None
-
-        log.info("Polar API started")
-
-        yield {
-            "async_engine": async_engine,
-            "async_sessionmaker": async_sessionmaker,
-            "sync_engine": sync_engine,
-            "sync_sessionmaker": sync_sessionmaker,
-            "arq_pool": arq_pool,
-            "ip_geolocation_client": ip_geolocation_client,
-        }
-
-        await async_engine.dispose()
-        sync_engine.dispose()
-        if ip_geolocation_client is not None:
-            ip_geolocation_client.close()
-
-        log.info("Polar API stopped")
+        async with create_redis() as redis:
+            async_engine = create_async_engine("app")
+            async_sessionmaker = create_async_sessionmaker(async_engine)
+            instrument_sqlalchemy(async_engine.sync_engine)
+
+            sync_engine = create_sync_engine("app")
+            sync_sessionmaker = create_sync_sessionmaker(sync_engine)
+            instrument_sqlalchemy(sync_engine)
+
+            try:
+                ip_geolocation_client = ip_geolocation.get_client()
+            except FileNotFoundError:
+                log.info(
+                    "IP geolocation database not found. "
+                    "Checkout won't automatically geolocate IPs."
+                )
+                ip_geolocation_client = None
+
+            log.info("Polar API started")
+
+            yield {
+                "async_engine": async_engine,
+                "async_sessionmaker": async_sessionmaker,
+                "sync_engine": sync_engine,
+                "sync_sessionmaker": sync_sessionmaker,
+                "arq_pool": arq_pool,
+                "redis": redis,
+                "ip_geolocation_client": ip_geolocation_client,
+            }
+
+            await async_engine.dispose()
+            sync_engine.dispose()
+            if ip_geolocation_client is not None:
+                ip_geolocation_client.close()
+
+            log.info("Polar API stopped")
 
 
 def create_app() -> FastAPI:

server/polar/backoffice/endpoints.py

@@ -20,6 +20,7 @@
 from polar.openapi import IN_DEVELOPMENT_ONLY
 from polar.pledge.service import pledge as pledge_service
 from polar.postgres import AsyncSession, get_db_session
+from polar.redis import Redis, get_redis
 from polar.repository.service import repository as repository_service
 from polar.reward.endpoints import to_resource as reward_to_resource
 from polar.reward.service import reward_service
@@ -201,6 +202,7 @@ async def manage_badge(
     badge: BackofficeBadge,
     auth_subject: AdminUser,
     session: AsyncSession = Depends(get_db_session),
+    redis: Redis = Depends(get_redis),
 ) -> BackofficeBadgeResponse:
     external_org = await external_organization_service.get_by_name(
         session, Platforms.github, badge.org_slug
@@ -233,6 +235,7 @@ async def manage_badge(
     if badge.action == "remove":
         issue = await github_issue.remove_polar_label(
             session,
+            redis,
             organization=external_org,
             repository=repo,
             issue=issue,
@@ -241,6 +244,7 @@ async def manage_badge(
     else:
         issue = await github_issue.add_polar_label(
             session,
+            redis,
             organization=external_org,
             repository=repo,
             issue=issue,

server/polar/benefit/benefits/__init__.py

@@ -7,6 +7,7 @@
 )
 from polar.models.benefit_grant import BenefitGrantPropertiesBase
 from polar.postgres import AsyncSession
+from polar.redis import Redis
 
 from .ads import BenefitAdsService
 from .articles import BenefitArticlesService
@@ -38,9 +39,9 @@
 
 
 def get_benefit_service(
-    type: BenefitType, session: AsyncSession
+    type: BenefitType, session: AsyncSession, redis: Redis
 ) -> BenefitServiceProtocol[Benefit, BenefitProperties, BenefitGrantPropertiesBase]:
-    return _SERVICE_CLASS_MAP[type](session)
+    return _SERVICE_CLASS_MAP[type](session, redis)
 
 
 __all__ = [

server/polar/benefit/benefits/articles.py

@@ -10,7 +10,6 @@
 from polar.models import ArticlesSubscription, Benefit, BenefitGrant, Organization, User
 from polar.models.benefit import BenefitArticles, BenefitArticlesProperties, BenefitType
 from polar.models.benefit_grant import BenefitGrantArticlesProperties
-from polar.redis import redis
 
 from .base import BenefitServiceProtocol
 
@@ -113,7 +112,7 @@ async def _get_articles_grants(
 
     @contextlib.asynccontextmanager
     async def _acquire_lock(self, user: User) -> AsyncGenerator[None, None]:
-        async with Locker(redis).lock(
+        async with Locker(self.redis).lock(
             f"articles-{user.id}", timeout=1, blocking_timeout=2
         ):
             yield

server/polar/benefit/benefits/base.py

@@ -8,6 +8,7 @@
     BenefitPreconditionErrorNotificationContextualPayload,
 )
 from polar.postgres import AsyncSession
+from polar.redis import Redis
 
 
 class BenefitServiceError(PolarError): ...
@@ -81,9 +82,11 @@ class BenefitServiceProtocol(Protocol[B, BP, BGP]):
     """
 
     session: AsyncSession
+    redis: Redis
 
-    def __init__(self, session: AsyncSession) -> None:
+    def __init__(self, session: AsyncSession, redis: Redis) -> None:
         self.session = session
+        self.redis = redis
 
     async def grant(
         self,

server/polar/benefit/benefits/github_repository.py

@@ -273,6 +273,7 @@ async def validate_properties(
         has_access = (
             await github_repository_benefit_user_service.user_has_access_to_repository(
                 oauth,
+                self.redis,
                 owner=repository_owner,
                 name=repository_name,
             )
@@ -292,6 +293,7 @@ async def validate_properties(
 
         installation = (
             await github_repository_benefit_user_service.get_repository_installation(
+                self.redis,
                 owner=repository_owner,
                 name=repository_name,
             )
@@ -312,7 +314,7 @@ async def validate_properties(
             "github-benefit-personal-org", user.posthog_distinct_id
         ):
             plan = await github_repository_benefit_user_service.get_billing_plan(
-                oauth, installation
+                self.redis, oauth, installation
             )
             if not plan or plan.is_personal:
                 raise BenefitPropertiesValidationError(
@@ -429,7 +431,7 @@ async def _get_github_app_client(
             installation_id = organization.installation_id
             assert installation_id is not None
             return github.get_app_installation_client(
-                installation_id, app=github.GitHubApp.polar
+                installation_id, redis=self.redis, app=github.GitHubApp.polar
             )
 
         # New integration, using the "Repository Benefit" GitHub App
@@ -439,10 +441,10 @@ async def _get_github_app_client(
         repository_name = benefit.properties["repository_name"]
         installation = (
             await github_repository_benefit_user_service.get_repository_installation(
-                owner=repository_owner, name=repository_name
+                self.redis, owner=repository_owner, name=repository_name
             )
         )
         assert installation is not None
         return github.get_app_installation_client(
-            installation.id, app=github.GitHubApp.repository_benefit
+            installation.id, redis=self.redis, app=github.GitHubApp.repository_benefit
         )

server/polar/benefit/endpoints.py

@@ -11,6 +11,7 @@
 from polar.organization.schemas import OrganizationID
 from polar.postgres import AsyncSession, get_db_session
 from polar.posthog import posthog
+from polar.redis import Redis, get_redis
 from polar.routing import APIRouter
 
 from . import auth
@@ -150,12 +151,13 @@ async def create(
     benefit_create: BenefitCreate = Body(..., discriminator="type"),
     authz: Authz = Depends(Authz.authz),
     session: AsyncSession = Depends(get_db_session),
+    redis: Redis = Depends(get_redis),
 ) -> Benefit:
     """
     Create a benefit.
     """
     benefit = await benefit_service.user_create(
-        session, authz, benefit_create, auth_subject
+        session, redis, authz, benefit_create, auth_subject
     )
 
     posthog.auth_subject_event(
@@ -188,6 +190,7 @@ async def update(
     auth_subject: auth.BenefitsWrite,
     authz: Authz = Depends(Authz.authz),
     session: AsyncSession = Depends(get_db_session),
+    redis: Redis = Depends(get_redis),
 ) -> Benefit:
     """
     Update a benefit.
@@ -209,7 +212,7 @@ async def update(
     )
 
     return await benefit_service.user_update(
-        session, authz, benefit, benefit_update, auth_subject
+        session, redis, authz, benefit, benefit_update, auth_subject
     )
 
 

server/polar/benefit/service/benefit.py

@@ -25,6 +25,7 @@
 from polar.models.webhook_endpoint import WebhookEventType
 from polar.organization.resolver import get_payload_organization
 from polar.postgres import sql
+from polar.redis import Redis
 from polar.webhook.service import webhook as webhook_service
 
 from ..benefits import get_benefit_service
@@ -132,6 +133,7 @@ async def get_by_id(
     async def user_create(
         self,
         session: AsyncSession,
+        redis: Redis,
         authz: Authz,
         create_schema: BenefitCreate,
         auth_subject: AuthSubject[User | Organization],
@@ -161,7 +163,7 @@ async def user_create(
         except AttributeError:
             is_tax_applicable = create_schema.type.is_tax_applicable()
 
-        benefit_service = get_benefit_service(create_schema.type, session)
+        benefit_service = get_benefit_service(create_schema.type, session, redis)
         properties = await benefit_service.validate_properties(
             auth_subject,
             create_schema.properties.model_dump(mode="json", by_alias=True),
@@ -195,6 +197,7 @@ async def user_create(
     async def user_update(
         self,
         session: AsyncSession,
+        redis: Redis,
         authz: Authz,
         benefit: Benefit,
         update_schema: BenefitUpdate,
@@ -211,7 +214,7 @@ async def user_update(
 
         properties_update: BaseModel | None = getattr(update_schema, "properties", None)
         if properties_update is not None:
-            benefit_service = get_benefit_service(benefit.type, session)
+            benefit_service = get_benefit_service(benefit.type, session, redis)
             update_dict["properties"] = await benefit_service.validate_properties(
                 auth_subject,
                 properties_update.model_dump(mode="json", by_alias=True),
@@ -224,7 +227,7 @@ async def user_update(
         session.add(benefit)
 
         await benefit_grant_service.enqueue_benefit_grant_updates(
-            session, benefit, previous_properties
+            session, redis, benefit, previous_properties
         )
 
         if benefit.organization: