Migrate order event backfill to a script

To make it more easy to control and handle when running

Author: Yopi

server/polar/order/tasks.py

@@ -1,24 +1,14 @@
-import asyncio
 import uuid
-from collections.abc import Sequence
-from typing import Any
 
 import stripe as stripe_lib
 import structlog
 from dramatiq import Retry
-from sqlalchemy import String, and_, func, or_, select
-from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload
 
-from polar.config import settings
-from polar.event.repository import EventRepository
-from polar.event.system import OrderPaidMetadata, OrderRefundedMetadata, SystemEvent
 from polar.exceptions import PolarTaskError
 from polar.logging import Logger
-from polar.models import Customer, Event, Order
-from polar.models.event import EventSource
-from polar.models.order import OrderBillingReasonInternal, OrderStatus
-from polar.models.refund import Refund, RefundStatus
+from polar.models import Customer, Order
+from polar.models.order import OrderBillingReasonInternal
 from polar.payment_method.repository import PaymentMethodRepository
 from polar.product.repository import ProductRepository
 from polar.subscription.repository import SubscriptionRepository
@@ -233,189 +223,3 @@ async def process_dunning_order(order_id: uuid.UUID) -> None:
             raise OrderDoesNotExist(order_id)
 
         await order_service.process_dunning_order(session, order)
-
-
-@actor(
-    actor_name="order.backfill_order_events",
-    priority=TaskPriority.LOW,
-    time_limit=3600_000,
-)
-async def backfill_order_events(
-    batch_size: int = settings.DATABASE_STREAM_YIELD_PER, rate_limit_delay: float = 1.0
-) -> None:
-    """
-    Backfill order.paid and order.refunded events for all existing orders.
-    Uses order.created_at for paid events and refund.created_at for refund events.
-
-    Args:
-        batch_size: The number of orders to process at a time between the delays (default 100)
-        rate_limit_delay: Seconds to wait between batch inserts (default: 1.0)
-    """
-
-    async with AsyncSessionMaker() as session:
-        last_created_at = None
-        last_id = None
-        total_orders = 0
-        total_events = 0
-
-        existing_order_ids_subquery = (
-            select(Event.user_metadata["order_id"].as_string().label("order_id"))
-            .where(Event.name.in_([SystemEvent.order_paid, SystemEvent.order_refunded]))
-            .distinct()
-            .subquery()
-        )
-
-        count_statement = (
-            select(func.count(Order.id))
-            .outerjoin(
-                existing_order_ids_subquery,
-                existing_order_ids_subquery.c.order_id == Order.id.cast(String),
-            )
-            .where(
-                Order.deleted_at.is_(None),
-                existing_order_ids_subquery.c.order_id.is_(None),
-                Order.status.in_(
-                    [
-                        OrderStatus.paid,
-                        OrderStatus.refunded,
-                        OrderStatus.partially_refunded,
-                    ]
-                ),
-            )
-        )
-        count_result = await session.execute(count_statement)
-        total_to_process = count_result.scalar() or 0
-
-        log.info(
-            "backfill_order_events.starting",
-            total_orders_to_process=total_to_process,
-        )
-
-        while True:
-            statement = (
-                select(Order)
-                .outerjoin(
-                    existing_order_ids_subquery,
-                    existing_order_ids_subquery.c.order_id == Order.id.cast(String),
-                )
-                .where(
-                    Order.deleted_at.is_(None),
-                    existing_order_ids_subquery.c.order_id.is_(None),
-                    Order.status.in_(
-                        [
-                            OrderStatus.paid,
-                            OrderStatus.refunded,
-                            OrderStatus.partially_refunded,
-                        ]
-                    ),
-                )
-                .options(selectinload(Order.customer))
-                .order_by(Order.created_at.asc(), Order.id.asc())
-                .limit(batch_size)
-            )
-
-            if last_created_at is not None:
-                statement = statement.where(
-                    or_(
-                        Order.created_at > last_created_at,
-                        and_(Order.created_at == last_created_at, Order.id > last_id),
-                    )
-                )
-
-            result = await session.execute(statement)
-            orders = result.scalars().all()
-            if not orders:
-                break
-
-            last_created_at = orders[-1].created_at
-            last_id = orders[-1].id
-
-            events = await _build_events_for_orders(session, orders)
-            total_orders += len(orders)
-
-            if events:
-                await EventRepository.from_session(session).insert_batch(events)
-                await session.commit()
-                total_events += len(events)
-                await asyncio.sleep(rate_limit_delay)
-
-            progress_pct = (
-                (total_orders / total_to_process * 100) if total_to_process > 0 else 100
-            )
-            log.info(
-                "backfill_order_events.progress",
-                progress_pct=f"{progress_pct:.1f}%",
-                orders=total_orders,
-                events=total_events,
-                total_to_process=total_to_process,
-            )
-
-        log.info(
-            "backfill_order_events.completed",
-            total_orders=total_orders,
-            total_events=total_events,
-        )
-
-
-async def _build_events_for_orders(
-    session: AsyncSession, orders: Sequence[Order]
-) -> list[dict[str, Any]]:
-    """Build events for a batch of orders using order and refund timestamps."""
-    order_ids = [order.id for order in orders]
-    orders_by_id = {order.id: order for order in orders}
-
-    refunds_result = await session.execute(
-        select(Refund)
-        .where(
-            Refund.order_id.in_(order_ids),
-            Refund.status == RefundStatus.succeeded,
-        )
-        .order_by(Refund.created_at.asc())
-    )
-    refunds = refunds_result.scalars().all()
-
-    refunds_by_order: dict[uuid.UUID, list[Refund]] = {}
-    for refund in refunds:
-        if refund.order_id is None:
-            continue
-        if refund.order_id not in refunds_by_order:
-            refunds_by_order[refund.order_id] = []
-        refunds_by_order[refund.order_id].append(refund)
-
-    events: list[dict[str, Any]] = []
-    for order_id, order in orders_by_id.items():
-        events.append(
-            {
-                "name": SystemEvent.order_paid,
-                "source": EventSource.system,
-                "timestamp": order.created_at,
-                "customer_id": order.customer_id,
-                "organization_id": order.customer.organization_id,
-                "user_metadata": OrderPaidMetadata(
-                    order_id=str(order.id),
-                    amount=order.total_amount,
-                    currency=order.currency,
-                    backfilled=True,
-                ),
-            }
-        )
-
-        order_refunds = refunds_by_order.get(order_id, [])
-        for refund in order_refunds:
-            events.append(
-                {
-                    "name": SystemEvent.order_refunded,
-                    "source": EventSource.system,
-                    "timestamp": refund.created_at,
-                    "customer_id": order.customer_id,
-                    "organization_id": order.customer.organization_id,
-                    "user_metadata": OrderRefundedMetadata(
-                        order_id=str(order.id),
-                        refunded_amount=refund.amount,
-                        currency=order.currency,
-                        backfilled=True,
-                    ),
-                }
-            )
-
-    return events