refactor: infer the queue name from task priority

Author: psincraian

server/polar/checkout/tasks.py

@@ -1,7 +1,7 @@
 import uuid
 
 from polar.exceptions import PolarTaskError
-from polar.worker import AsyncSessionMaker, CronTrigger, TaskPriority, TaskQueue, actor
+from polar.worker import AsyncSessionMaker, CronTrigger, TaskPriority, actor
 
 from .repository import CheckoutRepository
 from .service import checkout as checkout_service
@@ -20,7 +20,6 @@ def __init__(self, checkout_id: uuid.UUID) -> None:
 @actor(
     actor_name="checkout.handle_free_success",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def handle_free_success(checkout_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:

server/polar/email/tasks.py

@@ -1,9 +1,9 @@
-from polar.worker import TaskPriority, TaskQueue, actor
+from polar.worker import TaskPriority, actor
 
 from .sender import email_sender
 
 
-@actor(actor_name="email.send", priority=TaskPriority.HIGH, queue_name=TaskQueue.HIGH_PRIORITY)
+@actor(actor_name="email.send", priority=TaskPriority.HIGH)
 async def email_send(
     to_email_addr: str,
     subject: str,

server/polar/eventstream/tasks.py

@@ -1,12 +1,11 @@
-from polar.worker import RedisMiddleware, TaskPriority, TaskQueue, actor
+from polar.worker import RedisMiddleware, TaskPriority, actor
 
 from .service import send_event
 
 
 @actor(
     actor_name="eventstream.publish",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def eventstream_publish(event: str, channels: list[str]) -> None:
     await send_event(RedisMiddleware.get(), event, channels)

server/polar/integrations/stripe/tasks.py

@@ -42,7 +42,6 @@
     AsyncSessionMaker,
     RedisMiddleware,
     TaskPriority,
-    TaskQueue,
     actor,
     can_retry,
     get_retries,
@@ -81,7 +80,6 @@ class StripeTaskError(PolarTaskError): ...
 @actor(
     actor_name="stripe.webhook.account.updated",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def account_updated(event_id: uuid.UUID) -> None:
@@ -96,7 +94,6 @@ async def account_updated(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.payment_intent.succeeded",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def payment_intent_succeeded(event_id: uuid.UUID) -> None:
@@ -141,7 +138,6 @@ async def payment_intent_succeeded(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.payment_intent.payment_failed",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def payment_intent_payment_failed(event_id: uuid.UUID) -> None:
@@ -167,7 +163,6 @@ async def payment_intent_payment_failed(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.setup_intent.succeeded",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def setup_intent_succeeded(event_id: uuid.UUID) -> None:
@@ -189,7 +184,6 @@ async def setup_intent_succeeded(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.setup_intent.setup_failed",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def setup_intent_setup_failed(event_id: uuid.UUID) -> None:
@@ -210,7 +204,6 @@ async def setup_intent_setup_failed(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.charge.pending",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def charge_pending(event_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:
@@ -234,7 +227,6 @@ async def charge_pending(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.charge.failed",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def charge_failed(event_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:
@@ -254,7 +246,6 @@ async def charge_failed(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.charge.succeeded",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def charge_succeeded(event_id: uuid.UUID) -> None:
@@ -276,7 +267,6 @@ async def charge_succeeded(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.refund.created",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def refund_created(event_id: uuid.UUID) -> None:
@@ -295,7 +285,6 @@ async def refund_created(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.refund.updated",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def refund_updated(event_id: uuid.UUID) -> None:
@@ -314,7 +303,6 @@ async def refund_updated(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.refund.failed",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def refund_failed(event_id: uuid.UUID) -> None:
@@ -333,7 +321,6 @@ async def refund_failed(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.charge.dispute.closed",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def charge_dispute_closed(event_id: uuid.UUID) -> None:
@@ -353,7 +340,6 @@ async def charge_dispute_closed(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.customer.subscription.updated",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def customer_subscription_updated(event_id: uuid.UUID) -> None:
@@ -380,7 +366,6 @@ async def customer_subscription_updated(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.customer.subscription.deleted",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def customer_subscription_deleted(event_id: uuid.UUID) -> None:
@@ -407,7 +392,6 @@ async def customer_subscription_deleted(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.invoice.created",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def invoice_created(event_id: uuid.UUID) -> None:
@@ -433,7 +417,6 @@ async def invoice_created(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.invoice.paid",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 @stripe_api_connection_error_retry
 async def invoice_paid(event_id: uuid.UUID) -> None:
@@ -474,7 +457,6 @@ async def payout_paid(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.identity.verification_session.verified",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def identity_verification_session_verified(event_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:
@@ -490,7 +472,6 @@ async def identity_verification_session_verified(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.identity.verification_session.processing",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def identity_verification_session_processing(event_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:
@@ -506,7 +487,6 @@ async def identity_verification_session_processing(event_id: uuid.UUID) -> None:
 @actor(
     actor_name="stripe.webhook.identity.verification_session.requires_input",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def identity_verification_session_requires_input(event_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:

server/polar/subscription/tasks.py

@@ -15,7 +15,7 @@
 )
 from polar.product.repository import ProductRepository
 from polar.subscription.repository import SubscriptionRepository
-from polar.worker import AsyncSessionMaker, TaskPriority, TaskQueue, actor, enqueue_job
+from polar.worker import AsyncSessionMaker, TaskPriority, actor, enqueue_job
 
 from .service import SubscriptionNotReadyForMigration
 from .service import subscription as subscription_service
@@ -89,7 +89,6 @@ async def subscription_update_meters(subscription_id: uuid.UUID) -> None:
 @actor(
     actor_name="subscription.cancel_customer",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def subscription_cancel_customer(customer_id: uuid.UUID) -> None:
     async with AsyncSessionMaker() as session:

server/polar/webhook/tasks.py

@@ -14,7 +14,7 @@
 from polar.kit.utils import utc_now
 from polar.logging import Logger
 from polar.models.webhook_delivery import WebhookDelivery
-from polar.worker import AsyncSessionMaker, TaskPriority, TaskQueue, actor, can_retry, enqueue_job
+from polar.worker import AsyncSessionMaker, TaskPriority, actor, can_retry, enqueue_job
 
 from .service import webhook as webhook_service
 
@@ -129,7 +129,6 @@ async def _webhook_event_send(
 @actor(
     actor_name="webhook_event.success",
     priority=TaskPriority.HIGH,
-    queue_name=TaskQueue.HIGH_PRIORITY,
 )
 async def webhook_event_success(webhook_event_id: UUID) -> None:
     async with AsyncSessionMaker() as session:

server/polar/worker/__init__.py

@@ -189,11 +189,15 @@ class TaskQueue:
 def actor[**P, R](
     actor_class: Callable[..., dramatiq.Actor[Any, Any]] = dramatiq.Actor,
     actor_name: str | None = None,
-    queue_name: str = "default",
+    queue_name: str | None = None,
     priority: TaskPriority = TaskPriority.LOW,
     broker: dramatiq.Broker | None = None,
     **options: Any,
 ) -> Callable[[Callable[P, Awaitable[R]]], Callable[P, Awaitable[R]]]:
+    # Auto-map priority to queue if not explicitly set
+    if queue_name is None:
+        queue_name = TaskQueue.HIGH_PRIORITY if priority == TaskPriority.HIGH else TaskQueue.DEFAULT
+
     def decorator(
         fn: Callable[P, Awaitable[R]],
     ) -> Callable[P, Awaitable[R]]: