outpost: send reducer bucket totals over the WebSocket

server/polar/customer_meter/service.py

@@ -197,9 +197,11 @@ async def update_customer_meter(
 
         enqueue_outpost_event(
             customer.organization_id,
-            type="customer_meter",
-            customer_id=customer.id,
-            meter_id=meter.id,
+            {
+                "type": "customer_meter",
+                "customer_id": customer.id,
+                "meter_id": meter.id,
+            },
         )
 
         return await repository.update(customer_meter), True

server/polar/outpost/schemas.py

@@ -1,4 +1,5 @@
 import typing
+from datetime import datetime
 from enum import StrEnum
 
 from pydantic import UUID4, BaseModel, ConfigDict, Discriminator, TypeAdapter
@@ -64,6 +65,8 @@ class CustomerMeterIncomingMessage(IncomingMessageBase):
 class OutgoingMessageType(StrEnum):
     configuration = "configuration"
     customer_meter = "customer_meter"
+    reducer_bucket = "reducer_bucket"
+    reducer_seal = "reducer_seal"
 
 
 class OutgoingMessageBase(MessageBase):
@@ -100,8 +103,47 @@ class CustomerMeterOutgoingMessage(OutgoingMessageBase):
     payload: CustomerMeterOutgoingMessagePayload
 
 
+class ReducerBucketInterval(MessageBase):
+    bucket_start: datetime
+    bucket_end: datetime
+    sealed: bool
+    count: int
+    sum: float
+    min: float | None
+    max: float | None
+
+
+class ReducerBucketOutgoingMessagePayload(MessageBase):
+    reducer_id: UUID4
+    customer_id: UUID4 | None
+    external_customer_id: str | None
+    bucket_start: datetime
+    intervals: list[ReducerBucketInterval]
+
+
+class ReducerBucketOutgoingMessage(OutgoingMessageBase):
+    type: typing.Literal[OutgoingMessageType.reducer_bucket]
+    payload: ReducerBucketOutgoingMessagePayload
+
+
+class ReducerSealOutgoingMessagePayload(MessageBase):
+    reducer_id: UUID4
+    customer_id: UUID4 | None
+    external_customer_id: str | None
+    sealed_until: datetime
+
+
+class ReducerSealOutgoingMessage(OutgoingMessageBase):
+    type: typing.Literal[OutgoingMessageType.reducer_seal]
+    payload: ReducerSealOutgoingMessagePayload
+
+
 type OutgoingMessage = typing.Annotated[
-    ConfigurationOutgoingMessage | CustomerMeterOutgoingMessage, Discriminator("type")
+    ConfigurationOutgoingMessage
+    | CustomerMeterOutgoingMessage
+    | ReducerBucketOutgoingMessage
+    | ReducerSealOutgoingMessage,
+    Discriminator("type"),
 ]
 
 OutgoingMessageAdapter: TypeAdapter[OutgoingMessage] = TypeAdapter(OutgoingMessage)

server/polar/outpost/service.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 import structlog
 from anyio.streams.memory import MemoryObjectSendStream
 from pydantic import ValidationError
@@ -8,6 +10,8 @@
 from polar.logging import Logger
 from polar.models import Organization, Reducer
 from polar.reducer.repository import ReducerRepository
+from polar.reducer_bucket.repository import ReducerBucketRepository
+from polar.reducer_bucket.service import REDUCER_BUCKET_SIZE
 
 from .schemas import (
     CustomerMeterIncomingMessageCustomerPayload,
@@ -19,8 +23,13 @@
     OutgoingMessage,
     OutgoingMessageAdapter,
     OutgoingMessageType,
+    ReducerBucketInterval,
+    ReducerBucketOutgoingMessage,
+    ReducerBucketOutgoingMessagePayload,
+    ReducerSealOutgoingMessage,
+    ReducerSealOutgoingMessagePayload,
 )
-from .stream import OutpostEvent
+from .stream import OutpostEvent, ReducerBucketOutpostEvent
 
 log: Logger = structlog.get_logger(__name__)
 
@@ -69,6 +78,22 @@ async def handle_event(
                     ),
                     send_stream,
                 )
+            case "reducer_bucket":
+                return await self._send_reducer_bucket(
+                    session, organization, event, send_stream
+                )
+            case "reducer_seal":
+                return await send_stream.send(
+                    ReducerSealOutgoingMessage(
+                        type=OutgoingMessageType.reducer_seal,
+                        payload=ReducerSealOutgoingMessagePayload(
+                            reducer_id=event["reducer_id"],
+                            customer_id=event["customer_id"],
+                            external_customer_id=event["external_customer_id"],
+                            sealed_until=datetime.fromisoformat(event["sealed_until"]),
+                        ),
+                    )
+                )
 
     async def _send_configuration(
         self,
@@ -120,5 +145,37 @@ async def _send_customer_meter(
         )
         await send_stream.send(message)
 
+    async def _send_reducer_bucket(
+        self,
+        session: AsyncReadSession,
+        organization: Organization,
+        event: ReducerBucketOutpostEvent,
+        send_stream: MemoryObjectSendStream[OutgoingMessage],
+    ) -> None:
+        bucket_start = datetime.fromisoformat(event["bucket_start"])
+        repository = ReducerBucketRepository.from_session(session)
+        intervals = await repository.get_intervals(
+            organization.id,
+            event["reducer_id"],
+            bucket_start,
+            bucket_start + REDUCER_BUCKET_SIZE,
+            event["customer_id"],
+            event["external_customer_id"],
+        )
+        message = ReducerBucketOutgoingMessage(
+            type=OutgoingMessageType.reducer_bucket,
+            payload=ReducerBucketOutgoingMessagePayload(
+                reducer_id=event["reducer_id"],
+                customer_id=event["customer_id"],
+                external_customer_id=event["external_customer_id"],
+                bucket_start=bucket_start,
+                intervals=[
+                    ReducerBucketInterval.model_validate(interval._asdict())
+                    for interval in intervals
+                ],
+            ),
+        )
+        await send_stream.send(message)
+
 
 outpost = OutpostService()

server/polar/outpost/stream.py

@@ -16,20 +16,42 @@ class CustomerMeterOutpostEvent(typing.TypedDict):
     meter_id: uuid.UUID
 
 
-type OutpostEvent = CustomerMeterOutpostEvent
+class ReducerBucketOutpostEvent(typing.TypedDict):
+    """A customer's buckets are keyed by the pair (customer_id, external_customer_id):
+    (id, None), (None, external_id) and (id, external_id) are three distinct buckets
+    whose values add up."""
+
+    type: typing.Literal["reducer_bucket"]
+    reducer_id: uuid.UUID
+    bucket_start: str
+    customer_id: uuid.UUID | None
+    external_customer_id: str | None
+
+
+class ReducerSealOutpostEvent(typing.TypedDict):
+    type: typing.Literal["reducer_seal"]
+    reducer_id: uuid.UUID
+    customer_id: uuid.UUID | None
+    external_customer_id: str | None
+    sealed_until: str
+
+
+type OutpostEvent = (
+    CustomerMeterOutpostEvent | ReducerBucketOutpostEvent | ReducerSealOutpostEvent
+)
 
 
 async def publish(
     redis: Redis,
     organization_id: uuid.UUID,
-    **event: typing.Unpack[OutpostEvent],
+    event: OutpostEvent,
 ) -> None:
     await redis.publish(
         f"outpost:{organization_id}", json.dumps(event, default=json_obj_serializer)
     )
 
 
-def enqueue(organization_id: uuid.UUID, **event: typing.Unpack[OutpostEvent]) -> None:
+def enqueue(organization_id: uuid.UUID, event: OutpostEvent) -> None:
     enqueue_job(
         "outpost.publish",
         organization_id=organization_id,

server/polar/outpost/tasks.py

@@ -10,6 +10,8 @@
 @actor(actor_name="outpost.publish", priority=TaskPriority.HIGH)
 async def outpost_publish(
     organization_id: typing.Annotated[uuid.UUID, LoggableField],
-    **event: typing.Unpack[OutpostEvent],
+    **event: typing.Any,
 ) -> None:
-    await publish(RedisMiddleware.get(), organization_id, **event)
+    await publish(
+        RedisMiddleware.get(), organization_id, typing.cast(OutpostEvent, event)
+    )

server/polar/reducer_bucket/repository.py

@@ -146,6 +146,41 @@ async def get_latest_for_update(
         )
         return await self.get_one_or_none(statement)
 
+    async def get_intervals(
+        self,
+        organization_id: UUID,
+        reducer_id: UUID,
+        bucket_start: datetime,
+        bucket_end: datetime,
+        customer_id: UUID | None,
+        external_customer_id: str | None,
+    ) -> Sequence[Row[Any]]:
+        """count, sum, min and max of the customer identity's intervals starting in
+        the slot, with their sealed generations apart from the open one."""
+        sealed = ReducerBucket.sealed_at.is_not(None).label("sealed")
+        statement = (
+            select(
+                ReducerBucket.bucket_start,
+                ReducerBucket.bucket_end,
+                sealed,
+                func.sum(ReducerBucket.count).label("count"),
+                func.sum(ReducerBucket.sum).label("sum"),
+                func.min(ReducerBucket.min).label("min"),
+                func.max(ReducerBucket.max).label("max"),
+            )
+            .where(
+                ReducerBucket.organization_id == organization_id,
+                ReducerBucket.reducer_id == reducer_id,
+                ReducerBucket.bucket_start >= bucket_start,
+                ReducerBucket.bucket_start < bucket_end,
+                _is_customer(ReducerBucket, customer_id, external_customer_id),
+            )
+            .group_by(ReducerBucket.bucket_start, ReducerBucket.bucket_end, sealed)
+            .order_by(ReducerBucket.bucket_start, sealed.desc())
+        )
+        result = await self.session.execute(statement)
+        return result.all()
+
     async def add_contribution(
         self,
         reducer: Reducer,

server/polar/reducer_bucket/service.py

@@ -16,6 +16,7 @@
     UniqueAggregation,
 )
 from polar.models import Customer, Event, Reducer, ReducerBucket
+from polar.outpost.stream import enqueue as enqueue_outpost_event
 from polar.postgres import AsyncSession
 from polar.redis import Redis
 from polar.reducer.repository import ReducerRepository
@@ -143,6 +144,9 @@ async def add_to_closed_buckets(
                     count=len(values),
                     **statistics,
                 )
+            _enqueue_outpost_event(
+                reducer, bucket_start, customer_id, external_customer_id
+            )
 
     async def sync(
         self,
@@ -164,6 +168,7 @@ async def sync(
         await self._sync_bucket(
             repository, reducer, bucket_start, customer_id, external_customer_id
         )
+        _enqueue_outpost_event(reducer, bucket_start, customer_id, external_customer_id)
 
     async def seal(
         self,
@@ -273,6 +278,27 @@ async def seal(
                 span.set_attribute("sealed_buckets", len(sealed))
                 span.set_attribute("quantity", str(quantity))
             quantities[reducer.id] = quantity
+            for customer_id, external_customer_id in identities:
+                # Seal splits these slots. The ones between only change to sealed.
+                for bucket_start in sorted(
+                    {
+                        get_reducer_bucket_start(billable_from),
+                        get_reducer_bucket_start(cutoff),
+                    }
+                ):
+                    _enqueue_outpost_event(
+                        reducer, bucket_start, customer_id, external_customer_id
+                    )
+                enqueue_outpost_event(
+                    reducer.organization_id,
+                    {
+                        "type": "reducer_seal",
+                        "reducer_id": reducer.id,
+                        "customer_id": customer_id,
+                        "external_customer_id": external_customer_id,
+                        "sealed_until": cutoff.isoformat(),
+                    },
+                )
         return quantities
 
     async def _lock_buckets(
@@ -427,6 +453,24 @@ async def _split(
         return True
 
 
+def _enqueue_outpost_event(
+    reducer: Reducer,
+    bucket_start: datetime,
+    customer_id: uuid.UUID | None,
+    external_customer_id: str | None,
+) -> None:
+    enqueue_outpost_event(
+        reducer.organization_id,
+        {
+            "type": "reducer_bucket",
+            "reducer_id": reducer.id,
+            "bucket_start": bucket_start.isoformat(),
+            "customer_id": customer_id,
+            "external_customer_id": external_customer_id,
+        },
+    )
+
+
 def _get_bucket_key(bucket_id: redis_store.BucketId) -> str:
     reducer, bucket_start, customer_id, external_customer_id = bucket_id
     return redis_store.get_reducer_bucket_key(