From 53fe7b12a51c4047ba30d57bc004b964de652ed1 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 15:53:42 +0200 Subject: [PATCH 01/12] outpost: pass events to publish and enqueue as a dict --- server/polar/customer_meter/service.py | 8 +++++--- server/polar/outpost/stream.py | 4 ++-- server/polar/outpost/tasks.py | 6 ++++-- server/tests/outpost/test_endpoints.py | 8 +++++--- 4 files changed, 16 insertions(+), 10 deletions(-) diff --git a/server/polar/customer_meter/service.py b/server/polar/customer_meter/service.py index 3218517f58f..0836cad6a7d 100644 --- a/server/polar/customer_meter/service.py +++ b/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 diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index e02acc4a052..2b2ca620656 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -22,14 +22,14 @@ class CustomerMeterOutpostEvent(typing.TypedDict): 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, diff --git a/server/polar/outpost/tasks.py b/server/polar/outpost/tasks.py index c76dd8636cb..6c648edf8f0 100644 --- a/server/polar/outpost/tasks.py +++ b/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) + ) diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 16f5e3d6046..66e2e62fb36 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -104,9 +104,11 @@ async def test_event_customer_meter( await publish( redis, customer_meter.customer.organization_id, - type="customer_meter", - customer_id=customer_meter.customer_id, - meter_id=customer_meter.meter_id, + { + "type": "customer_meter", + "customer_id": customer_meter.customer_id, + "meter_id": customer_meter.meter_id, + }, ) response = await websocket.receive_json() From 3e033001aa8d829ef22bf036ce4eb6d57f545441 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 15:55:54 +0200 Subject: [PATCH 02/12] outpost: send reducer bucket totals over the WebSocket --- server/polar/outpost/schemas.py | 23 ++++++- server/polar/outpost/service.py | 44 +++++++++++++- server/polar/outpost/stream.py | 10 ++- server/polar/reducer_bucket/repository.py | 26 ++++++++ server/tests/outpost/test_endpoints.py | 74 ++++++++++++++++++++++- 5 files changed, 173 insertions(+), 4 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 9a26797ff05..ec647ec1c14 100644 --- a/server/polar/outpost/schemas.py +++ b/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,7 @@ class CustomerMeterIncomingMessage(IncomingMessageBase): class OutgoingMessageType(StrEnum): configuration = "configuration" customer_meter = "customer_meter" + reducer_bucket = "reducer_bucket" class OutgoingMessageBase(MessageBase): @@ -100,8 +102,27 @@ class CustomerMeterOutgoingMessage(OutgoingMessageBase): payload: CustomerMeterOutgoingMessagePayload +class ReducerBucketOutgoingMessagePayload(MessageBase): + reducer_id: UUID4 + customer_id: UUID4 | None + external_customer_id: str | None + bucket_start: datetime + count: int + sum: float + min: float | None + max: float | None + + +class ReducerBucketOutgoingMessage(OutgoingMessageBase): + type: typing.Literal[OutgoingMessageType.reducer_bucket] + payload: ReducerBucketOutgoingMessagePayload + + type OutgoingMessage = typing.Annotated[ - ConfigurationOutgoingMessage | CustomerMeterOutgoingMessage, Discriminator("type") + ConfigurationOutgoingMessage + | CustomerMeterOutgoingMessage + | ReducerBucketOutgoingMessage, + Discriminator("type"), ] OutgoingMessageAdapter: TypeAdapter[OutgoingMessage] = TypeAdapter(OutgoingMessage) diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index 8899ca59678..f357464623b 100644 --- a/server/polar/outpost/service.py +++ b/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,10 @@ OutgoingMessage, OutgoingMessageAdapter, OutgoingMessageType, + ReducerBucketOutgoingMessage, + ReducerBucketOutgoingMessagePayload, ) -from .stream import OutpostEvent +from .stream import OutpostEvent, ReducerBucketOutpostEvent log: Logger = structlog.get_logger(__name__) @@ -69,6 +75,10 @@ async def handle_event( ), send_stream, ) + case "reducer_bucket": + return await self._send_reducer_bucket( + session, organization, event, send_stream + ) async def _send_configuration( self, @@ -120,5 +130,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) + count, total, minimum, maximum = await repository.get_totals( + 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, + count=count, + sum=total, + min=minimum, + max=maximum, + ), + ) + await send_stream.send(message) + outpost = OutpostService() diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index 2b2ca620656..ca9052f37b5 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -16,7 +16,15 @@ class CustomerMeterOutpostEvent(typing.TypedDict): meter_id: uuid.UUID -type OutpostEvent = CustomerMeterOutpostEvent +class ReducerBucketOutpostEvent(typing.TypedDict): + type: typing.Literal["reducer_bucket"] + reducer_id: uuid.UUID + bucket_start: str + customer_id: uuid.UUID | None + external_customer_id: str | None + + +type OutpostEvent = CustomerMeterOutpostEvent | ReducerBucketOutpostEvent async def publish( diff --git a/server/polar/reducer_bucket/repository.py b/server/polar/reducer_bucket/repository.py index 1a05ebab50a..4b21abce9d1 100644 --- a/server/polar/reducer_bucket/repository.py +++ b/server/polar/reducer_bucket/repository.py @@ -146,6 +146,32 @@ async def get_latest_for_update( ) return await self.get_one_or_none(statement) + async def get_totals( + self, + organization_id: UUID, + reducer_id: UUID, + bucket_start: datetime, + bucket_end: datetime, + customer_id: UUID | None, + external_customer_id: str | None, + ) -> Row[Any]: + """count, sum, min and max of the customer identity's rows starting in the + slot, across all their generations.""" + statement = select( + func.coalesce(func.sum(ReducerBucket.count), 0).label("count"), + func.coalesce(func.sum(ReducerBucket.sum), 0).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), + ) + result = await self.session.execute(statement) + return result.one() + async def add_contribution( self, reducer: Reducer, diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 66e2e62fb36..96e2d7733ce 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -1,10 +1,17 @@ +from datetime import UTC, datetime, timedelta from decimal import Decimal import httpx import httpx_ws import pytest -from polar.models import Customer, CustomerMeter, Meter, Organization +from polar.models import ( + Customer, + CustomerMeter, + Meter, + Organization, + ReducerBucket, +) from polar.outpost.stream import publish from polar.redis import Redis from tests.fixtures import AuthSubjectFixture, SaveFixture @@ -123,3 +130,68 @@ async def test_event_customer_meter( assert response["payload"]["consumed_units"] == 90.0 assert response["payload"]["credited_units"] == 100 assert response["payload"]["balance"] == 10.0 + + @pytest.mark.auth(AuthSubjectFixture(subject="organization")) + async def test_event_reducer_bucket( + self, + save_fixture: SaveFixture, + client: httpx.AsyncClient, + redis: Redis, + organization: Organization, + meter: Meter, + ) -> None: + reducer = await create_reducer( + save_fixture, organization=organization, meters=[meter] + ) + bucket_start = datetime(2026, 10, 7, 12, 5, tzinfo=UTC) + split_at = bucket_start + timedelta(minutes=2) + bucket_end = bucket_start + timedelta(minutes=5) + # Two generations before the split, one row after it. + for start, end, generation, value in ( + (bucket_start, split_at, 1, 3), + (bucket_start, split_at, 2, 5), + (split_at, bucket_end, 1, 4), + ): + await save_fixture( + ReducerBucket( + organization=organization, + reducer=reducer, + external_customer_id="external", + bucket_start=start, + bucket_end=end, + count=1, + sum=Decimal(value), + min=Decimal(value), + max=Decimal(value), + generation=generation, + ) + ) + + async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] + await publish( + redis, + organization.id, + { + "type": "reducer_bucket", + "reducer_id": reducer.id, + "bucket_start": bucket_start.isoformat(), + "customer_id": None, + "external_customer_id": "external", + }, + ) + + response = await websocket.receive_json() + + assert response == { + "type": "reducer_bucket", + "payload": { + "reducer_id": str(reducer.id), + "customer_id": None, + "external_customer_id": "external", + "bucket_start": "2026-10-07T12:05:00Z", + "count": 3, + "sum": 12.0, + "min": 3.0, + "max": 5.0, + }, + } From abb3b3e5a7dd638c7ed70239ad58bd76c46856d2 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 15:57:08 +0200 Subject: [PATCH 03/12] reducer_bucket: push buckets to Outpost after sync and closed-bucket rollups --- server/polar/reducer_bucket/service.py | 23 ++++++++++++++++ server/tests/reducer_bucket/test_service.py | 30 +++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/server/polar/reducer_bucket/service.py b/server/polar/reducer_bucket/service.py index 07707092ef5..977582d4639 100644 --- a/server/polar/reducer_bucket/service.py +++ b/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, @@ -427,6 +432,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( diff --git a/server/tests/reducer_bucket/test_service.py b/server/tests/reducer_bucket/test_service.py index 3b24545b9f5..70970d4c7ac 100644 --- a/server/tests/reducer_bucket/test_service.py +++ b/server/tests/reducer_bucket/test_service.py @@ -127,12 +127,16 @@ async def test_routes_buckets_by_age( async def test_adds_to_closed_buckets( self, + mocker: MockerFixture, save_fixture: SaveFixture, session: AsyncSession, sum_reducer: Reducer, redis: Redis, organization: Organization, ) -> None: + enqueue_outpost_event = mocker.patch( + "polar.reducer_bucket.service.enqueue_outpost_event" + ) previous_start = BUCKET_START - REDUCER_BUCKET_SIZE for batch in ((3, 4), (2,)): await reducer_bucket_service.rollup( @@ -152,6 +156,17 @@ async def test_adds_to_closed_buckets( [bucket] = await get_buckets(session) assert (bucket.generation, bucket.count, bucket.sum) == (1, 3, Decimal(9)) assert (bucket.min, bucket.max) == (Decimal(2), Decimal(4)) + enqueue_outpost_event.assert_called_with( + organization.id, + { + "type": "reducer_bucket", + "reducer_id": sum_reducer.id, + "bucket_start": previous_start.isoformat(), + "customer_id": None, + "external_customer_id": "external", + }, + ) + assert enqueue_outpost_event.call_count == 2 async def test_adds_to_next_generation_once_billed( self, @@ -380,6 +395,7 @@ async def test_skips_inactive_reducers( async def test_writes_buckets( self, + mocker: MockerFixture, save_fixture: SaveFixture, session: AsyncSession, sum_reducer: Reducer, @@ -404,10 +420,24 @@ async def test_writes_buckets( metadata={"tokens": 6}, ) + enqueue_outpost_event = mocker.patch( + "polar.reducer_bucket.service.enqueue_outpost_event" + ) + await reducer_bucket_service.sync( session, sum_reducer.id, BUCKET_START, None, "external" ) + enqueue_outpost_event.assert_called_once_with( + organization.id, + { + "type": "reducer_bucket", + "reducer_id": sum_reducer.id, + "bucket_start": BUCKET_START.isoformat(), + "customer_id": None, + "external_customer_id": "external", + }, + ) [bucket] = await get_buckets(session) assert bucket.reducer_id == sum_reducer.id assert bucket.bucket_start == BUCKET_START From 64ecaac853841a6f48f73c8b28410869948a5ae0 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 16:08:11 +0200 Subject: [PATCH 04/12] outpost: send each interval of a reducer bucket slot --- server/polar/outpost/schemas.py | 15 ++++++-- server/polar/outpost/service.py | 11 +++--- server/polar/reducer_bucket/repository.py | 41 +++++++++++++-------- server/tests/outpost/test_endpoints.py | 45 ++++++++++++++++++----- 4 files changed, 78 insertions(+), 34 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index ec647ec1c14..36cc9c808fc 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -102,17 +102,24 @@ class CustomerMeterOutgoingMessage(OutgoingMessageBase): payload: CustomerMeterOutgoingMessagePayload -class ReducerBucketOutgoingMessagePayload(MessageBase): - reducer_id: UUID4 - customer_id: UUID4 | None - external_customer_id: str | None +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 diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index f357464623b..9e12163bb84 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -23,6 +23,7 @@ OutgoingMessage, OutgoingMessageAdapter, OutgoingMessageType, + ReducerBucketInterval, ReducerBucketOutgoingMessage, ReducerBucketOutgoingMessagePayload, ) @@ -139,7 +140,7 @@ async def _send_reducer_bucket( ) -> None: bucket_start = datetime.fromisoformat(event["bucket_start"]) repository = ReducerBucketRepository.from_session(session) - count, total, minimum, maximum = await repository.get_totals( + intervals = await repository.get_intervals( organization.id, event["reducer_id"], bucket_start, @@ -154,10 +155,10 @@ async def _send_reducer_bucket( customer_id=event["customer_id"], external_customer_id=event["external_customer_id"], bucket_start=bucket_start, - count=count, - sum=total, - min=minimum, - max=maximum, + intervals=[ + ReducerBucketInterval.model_validate(interval._asdict()) + for interval in intervals + ], ), ) await send_stream.send(message) diff --git a/server/polar/reducer_bucket/repository.py b/server/polar/reducer_bucket/repository.py index 4b21abce9d1..f1f8fb64a45 100644 --- a/server/polar/reducer_bucket/repository.py +++ b/server/polar/reducer_bucket/repository.py @@ -146,7 +146,7 @@ async def get_latest_for_update( ) return await self.get_one_or_none(statement) - async def get_totals( + async def get_intervals( self, organization_id: UUID, reducer_id: UUID, @@ -154,23 +154,32 @@ async def get_totals( bucket_end: datetime, customer_id: UUID | None, external_customer_id: str | None, - ) -> Row[Any]: - """count, sum, min and max of the customer identity's rows starting in the - slot, across all their generations.""" - statement = select( - func.coalesce(func.sum(ReducerBucket.count), 0).label("count"), - func.coalesce(func.sum(ReducerBucket.sum), 0).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), + ) -> 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.one() + return result.all() async def add_contribution( self, diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 96e2d7733ce..0a4e587f115 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -146,11 +146,12 @@ async def test_event_reducer_bucket( bucket_start = datetime(2026, 10, 7, 12, 5, tzinfo=UTC) split_at = bucket_start + timedelta(minutes=2) bucket_end = bucket_start + timedelta(minutes=5) - # Two generations before the split, one row after it. - for start, end, generation, value in ( - (bucket_start, split_at, 1, 3), - (bucket_start, split_at, 2, 5), - (split_at, bucket_end, 1, 4), + billed_at = datetime(2026, 10, 7, 12, 20, tzinfo=UTC) + for start, end, generation, sealed_at, value in ( + (bucket_start, split_at, 1, billed_at, 3), + (bucket_start, split_at, 2, billed_at, 4), + (bucket_start, split_at, 3, None, 5), + (split_at, bucket_end, 1, None, 6), ): await save_fixture( ReducerBucket( @@ -164,6 +165,7 @@ async def test_event_reducer_bucket( min=Decimal(value), max=Decimal(value), generation=generation, + sealed_at=sealed_at, ) ) @@ -189,9 +191,34 @@ async def test_event_reducer_bucket( "customer_id": None, "external_customer_id": "external", "bucket_start": "2026-10-07T12:05:00Z", - "count": 3, - "sum": 12.0, - "min": 3.0, - "max": 5.0, + "intervals": [ + { + "bucket_start": "2026-10-07T12:05:00Z", + "bucket_end": "2026-10-07T12:07:00Z", + "sealed": True, + "count": 2, + "sum": 7.0, + "min": 3.0, + "max": 4.0, + }, + { + "bucket_start": "2026-10-07T12:05:00Z", + "bucket_end": "2026-10-07T12:07:00Z", + "sealed": False, + "count": 1, + "sum": 5.0, + "min": 5.0, + "max": 5.0, + }, + { + "bucket_start": "2026-10-07T12:07:00Z", + "bucket_end": "2026-10-07T12:10:00Z", + "sealed": False, + "count": 1, + "sum": 6.0, + "min": 6.0, + "max": 6.0, + }, + ], }, } From 9631ebcd15574bc1ac3b0641fae695478d6d6878 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 16:09:26 +0200 Subject: [PATCH 05/12] outpost: document the customer identity of reducer bucket events --- server/polar/outpost/stream.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index ca9052f37b5..2bdf8695d16 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -17,6 +17,10 @@ class CustomerMeterOutpostEvent(typing.TypedDict): 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 From 85c319434c752447d1889133fd819b7e2727b6ae Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Tue, 13 Oct 2026 16:41:20 +0200 Subject: [PATCH 06/12] outpost: tell Outpost what a seal changed --- server/polar/outpost/schemas.py | 16 ++++++- server/polar/outpost/service.py | 14 ++++++ server/polar/outpost/stream.py | 12 ++++- server/polar/reducer_bucket/service.py | 21 ++++++++ server/tests/outpost/test_endpoints.py | 31 ++++++++++++ server/tests/reducer_bucket/test_service.py | 53 +++++++++++++++++++++ 6 files changed, 145 insertions(+), 2 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 36cc9c808fc..471ea2ab47b 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -66,6 +66,7 @@ class OutgoingMessageType(StrEnum): configuration = "configuration" customer_meter = "customer_meter" reducer_bucket = "reducer_bucket" + reducer_seal = "reducer_seal" class OutgoingMessageBase(MessageBase): @@ -125,10 +126,23 @@ class ReducerBucketOutgoingMessage(OutgoingMessageBase): 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 - | ReducerBucketOutgoingMessage, + | ReducerBucketOutgoingMessage + | ReducerSealOutgoingMessage, Discriminator("type"), ] diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index 9e12163bb84..98c2d11f922 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -26,6 +26,8 @@ ReducerBucketInterval, ReducerBucketOutgoingMessage, ReducerBucketOutgoingMessagePayload, + ReducerSealOutgoingMessage, + ReducerSealOutgoingMessagePayload, ) from .stream import OutpostEvent, ReducerBucketOutpostEvent @@ -80,6 +82,18 @@ async def handle_event( 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, diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index 2bdf8695d16..0a646a30a59 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -28,7 +28,17 @@ class ReducerBucketOutpostEvent(typing.TypedDict): external_customer_id: str | None -type OutpostEvent = CustomerMeterOutpostEvent | ReducerBucketOutpostEvent +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( diff --git a/server/polar/reducer_bucket/service.py b/server/polar/reducer_bucket/service.py index 977582d4639..36bacfee1f5 100644 --- a/server/polar/reducer_bucket/service.py +++ b/server/polar/reducer_bucket/service.py @@ -278,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( diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 0a4e587f115..94d69a47708 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -1,3 +1,4 @@ +import uuid from datetime import UTC, datetime, timedelta from decimal import Decimal @@ -222,3 +223,33 @@ async def test_event_reducer_bucket( ], }, } + + @pytest.mark.auth(AuthSubjectFixture(subject="organization")) + async def test_event_reducer_seal( + self, client: httpx.AsyncClient, redis: Redis, organization: Organization + ) -> None: + reducer_id = uuid.uuid4() + async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] + await publish( + redis, + organization.id, + { + "type": "reducer_seal", + "reducer_id": reducer_id, + "customer_id": None, + "external_customer_id": "external", + "sealed_until": "2026-10-07T12:07:00+00:00", + }, + ) + + response = await websocket.receive_json() + + assert response == { + "type": "reducer_seal", + "payload": { + "reducer_id": str(reducer_id), + "customer_id": None, + "external_customer_id": "external", + "sealed_until": "2026-10-07T12:07:00Z", + }, + } diff --git a/server/tests/reducer_bucket/test_service.py b/server/tests/reducer_bucket/test_service.py index 70970d4c7ac..aba4db24419 100644 --- a/server/tests/reducer_bucket/test_service.py +++ b/server/tests/reducer_bucket/test_service.py @@ -959,3 +959,56 @@ async def test_returns_sealed_quantity( assert quantities == {reducer.id: expected} assert len(await get_buckets(session)) == 2 + + async def test_notifies_outpost( + self, + mocker: MockerFixture, + session: AsyncSession, + max_reducer: Reducer, + external_customer: Customer, + ) -> None: + enqueue_outpost_event = mocker.patch( + "polar.reducer_bucket.service.enqueue_outpost_event" + ) + + await reducer_bucket_service.seal( + session, + [max_reducer], + external_customer, + billable_from=START, + cutoff=CUTOFF, + ) + + assert enqueue_outpost_event.call_args_list == [ + call + for customer_id, external_customer_id in ( + (external_customer.id, None), + (None, "external"), + (external_customer.id, "external"), + ) + for call in ( + *( + mocker.call( + max_reducer.organization_id, + { + "type": "reducer_bucket", + "reducer_id": max_reducer.id, + "bucket_start": bucket_start.isoformat(), + "customer_id": customer_id, + "external_customer_id": external_customer_id, + }, + ) + for bucket_start in (START, BUCKET_START) + ), + mocker.call( + max_reducer.organization_id, + { + "type": "reducer_seal", + "reducer_id": max_reducer.id, + "customer_id": customer_id, + "external_customer_id": external_customer_id, + "sealed_until": CUTOFF.isoformat(), + }, + ), + ) + ] From daecb9e20948e631db449ca8f6af5630b0a09dc3 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 09:42:02 +0200 Subject: [PATCH 07/12] Revert "outpost: tell Outpost what a seal changed" --- server/polar/outpost/schemas.py | 16 +------ server/polar/outpost/service.py | 14 ------ server/polar/outpost/stream.py | 12 +---- server/polar/reducer_bucket/service.py | 21 -------- server/tests/outpost/test_endpoints.py | 31 ------------ server/tests/reducer_bucket/test_service.py | 53 --------------------- 6 files changed, 2 insertions(+), 145 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 471ea2ab47b..36cc9c808fc 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -66,7 +66,6 @@ class OutgoingMessageType(StrEnum): configuration = "configuration" customer_meter = "customer_meter" reducer_bucket = "reducer_bucket" - reducer_seal = "reducer_seal" class OutgoingMessageBase(MessageBase): @@ -126,23 +125,10 @@ class ReducerBucketOutgoingMessage(OutgoingMessageBase): 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 - | ReducerBucketOutgoingMessage - | ReducerSealOutgoingMessage, + | ReducerBucketOutgoingMessage, Discriminator("type"), ] diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index 98c2d11f922..9e12163bb84 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -26,8 +26,6 @@ ReducerBucketInterval, ReducerBucketOutgoingMessage, ReducerBucketOutgoingMessagePayload, - ReducerSealOutgoingMessage, - ReducerSealOutgoingMessagePayload, ) from .stream import OutpostEvent, ReducerBucketOutpostEvent @@ -82,18 +80,6 @@ async def handle_event( 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, diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index 0a646a30a59..2bdf8695d16 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -28,17 +28,7 @@ class ReducerBucketOutpostEvent(typing.TypedDict): 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 -) +type OutpostEvent = CustomerMeterOutpostEvent | ReducerBucketOutpostEvent async def publish( diff --git a/server/polar/reducer_bucket/service.py b/server/polar/reducer_bucket/service.py index 36bacfee1f5..977582d4639 100644 --- a/server/polar/reducer_bucket/service.py +++ b/server/polar/reducer_bucket/service.py @@ -278,27 +278,6 @@ 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( diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 94d69a47708..0a4e587f115 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -1,4 +1,3 @@ -import uuid from datetime import UTC, datetime, timedelta from decimal import Decimal @@ -223,33 +222,3 @@ async def test_event_reducer_bucket( ], }, } - - @pytest.mark.auth(AuthSubjectFixture(subject="organization")) - async def test_event_reducer_seal( - self, client: httpx.AsyncClient, redis: Redis, organization: Organization - ) -> None: - reducer_id = uuid.uuid4() - async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] - await publish( - redis, - organization.id, - { - "type": "reducer_seal", - "reducer_id": reducer_id, - "customer_id": None, - "external_customer_id": "external", - "sealed_until": "2026-10-07T12:07:00+00:00", - }, - ) - - response = await websocket.receive_json() - - assert response == { - "type": "reducer_seal", - "payload": { - "reducer_id": str(reducer_id), - "customer_id": None, - "external_customer_id": "external", - "sealed_until": "2026-10-07T12:07:00Z", - }, - } diff --git a/server/tests/reducer_bucket/test_service.py b/server/tests/reducer_bucket/test_service.py index aba4db24419..70970d4c7ac 100644 --- a/server/tests/reducer_bucket/test_service.py +++ b/server/tests/reducer_bucket/test_service.py @@ -959,56 +959,3 @@ async def test_returns_sealed_quantity( assert quantities == {reducer.id: expected} assert len(await get_buckets(session)) == 2 - - async def test_notifies_outpost( - self, - mocker: MockerFixture, - session: AsyncSession, - max_reducer: Reducer, - external_customer: Customer, - ) -> None: - enqueue_outpost_event = mocker.patch( - "polar.reducer_bucket.service.enqueue_outpost_event" - ) - - await reducer_bucket_service.seal( - session, - [max_reducer], - external_customer, - billable_from=START, - cutoff=CUTOFF, - ) - - assert enqueue_outpost_event.call_args_list == [ - call - for customer_id, external_customer_id in ( - (external_customer.id, None), - (None, "external"), - (external_customer.id, "external"), - ) - for call in ( - *( - mocker.call( - max_reducer.organization_id, - { - "type": "reducer_bucket", - "reducer_id": max_reducer.id, - "bucket_start": bucket_start.isoformat(), - "customer_id": customer_id, - "external_customer_id": external_customer_id, - }, - ) - for bucket_start in (START, BUCKET_START) - ), - mocker.call( - max_reducer.organization_id, - { - "type": "reducer_seal", - "reducer_id": max_reducer.id, - "customer_id": customer_id, - "external_customer_id": external_customer_id, - "sealed_until": CUTOFF.isoformat(), - }, - ), - ) - ] From 152e0d984e8dfac2a31b2dcccfdd8111d255d565 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 09:45:26 +0200 Subject: [PATCH 08/12] outpost: send the customer's unbilled reducer value --- server/polar/outpost/schemas.py | 17 ++---- server/polar/outpost/service.py | 48 +++++++++------ server/polar/outpost/stream.py | 5 -- server/polar/reducer_bucket/repository.py | 49 +++++++--------- server/polar/reducer_bucket/service.py | 28 +++++---- server/tests/outpost/test_endpoints.py | 65 +++++++-------------- server/tests/reducer_bucket/test_service.py | 2 - 7 files changed, 89 insertions(+), 125 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 36cc9c808fc..115ba7ab640 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -1,5 +1,4 @@ import typing -from datetime import datetime from enum import StrEnum from pydantic import UUID4, BaseModel, ConfigDict, Discriminator, TypeAdapter @@ -102,22 +101,14 @@ 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] + count: int + sum: float + min: float | None + max: float | None class ReducerBucketOutgoingMessage(OutgoingMessageBase): diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index 9e12163bb84..b0044bb0c61 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -1,17 +1,16 @@ -from datetime import datetime - import structlog from anyio.streams.memory import MemoryObjectSendStream from pydantic import ValidationError from sqlalchemy.orm import selectinload +from polar.customer.repository import CustomerRepository from polar.customer_meter.repository import CustomerMeterRepository from polar.kit.db.postgres import AsyncReadSession 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 polar.reducer_bucket.service import get_customer_identities from .schemas import ( CustomerMeterIncomingMessageCustomerPayload, @@ -23,7 +22,6 @@ OutgoingMessage, OutgoingMessageAdapter, OutgoingMessageType, - ReducerBucketInterval, ReducerBucketOutgoingMessage, ReducerBucketOutgoingMessagePayload, ) @@ -138,27 +136,39 @@ async def _send_reducer_bucket( event: ReducerBucketOutpostEvent, send_stream: MemoryObjectSendStream[OutgoingMessage], ) -> None: - bucket_start = datetime.fromisoformat(event["bucket_start"]) + customer_id = event["customer_id"] + external_customer_id = event["external_customer_id"] + customer_repository = CustomerRepository.from_session(session) + if customer_id is not None: + customer = await customer_repository.get_by_id_and_organization( + customer_id, organization.id + ) + else: + assert external_customer_id is not None + customer = await customer_repository.get_by_external_id_and_organization( + external_customer_id, organization.id + ) + if customer is not None: + customer_id = customer.id + external_customer_id = customer.external_id + identities = get_customer_identities(customer) + else: + identities = [(customer_id, external_customer_id)] + 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"], + count, total, minimum, maximum = await repository.get_unsealed_totals( + organization.id, event["reducer_id"], identities ) 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 - ], + customer_id=customer_id, + external_customer_id=external_customer_id, + count=count, + sum=total, + min=minimum, + max=maximum, ), ) await send_stream.send(message) diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index 2bdf8695d16..1ecf3605b4e 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -17,13 +17,8 @@ class CustomerMeterOutpostEvent(typing.TypedDict): 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 diff --git a/server/polar/reducer_bucket/repository.py b/server/polar/reducer_bucket/repository.py index f1f8fb64a45..6e173446ea3 100644 --- a/server/polar/reducer_bucket/repository.py +++ b/server/polar/reducer_bucket/repository.py @@ -146,40 +146,31 @@ async def get_latest_for_update( ) return await self.get_one_or_none(statement) - async def get_intervals( + async def get_unsealed_totals( 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()) + identities: Sequence[tuple[UUID | None, str | None]], + ) -> Row[Any]: + """count, sum, min and max of the identities' usage not billed yet.""" + statement = select( + func.coalesce(func.sum(ReducerBucket.count), 0).label("count"), + func.coalesce(func.sum(ReducerBucket.sum), 0).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.sealed_at.is_(None), + or_( + *( + _is_customer(ReducerBucket, customer_id, external_customer_id) + for customer_id, external_customer_id in identities + ) + ), ) result = await self.session.execute(statement) - return result.all() + return result.one() async def add_contribution( self, diff --git a/server/polar/reducer_bucket/service.py b/server/polar/reducer_bucket/service.py index 977582d4639..e953f232e5e 100644 --- a/server/polar/reducer_bucket/service.py +++ b/server/polar/reducer_bucket/service.py @@ -38,6 +38,19 @@ def get_reducer_bucket_start(timestamp: datetime) -> datetime: return datetime.fromtimestamp(epoch - epoch % size, tz=UTC) +def get_customer_identities( + customer: Customer, +) -> list[tuple[uuid.UUID | None, str | None]]: + """The (customer_id, external_customer_id) pairs the customer's buckets use.""" + identities: list[tuple[uuid.UUID | None, str | None]] = [(customer.id, None)] + if customer.external_id is not None: + identities += [ + (None, customer.external_id), + (customer.id, customer.external_id), + ] + return identities + + class ReducerBucketService: async def rollup( self, @@ -144,9 +157,7 @@ async def add_to_closed_buckets( count=len(values), **statistics, ) - _enqueue_outpost_event( - reducer, bucket_start, customer_id, external_customer_id - ) + _enqueue_outpost_event(reducer, customer_id, external_customer_id) async def sync( self, @@ -168,7 +179,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) + _enqueue_outpost_event(reducer, customer_id, external_customer_id) async def seal( self, @@ -185,12 +196,7 @@ async def seal( if billable_from >= cutoff: return quantities repository = ReducerBucketRepository.from_session(session) - identities: list[tuple[uuid.UUID | None, str | None]] = [(customer.id, None)] - if customer.external_id is not None: - identities += [ - (None, customer.external_id), - (customer.id, customer.external_id), - ] + identities = get_customer_identities(customer) await self._lock_buckets( repository, @@ -434,7 +440,6 @@ async def _split( def _enqueue_outpost_event( reducer: Reducer, - bucket_start: datetime, customer_id: uuid.UUID | None, external_customer_id: str | None, ) -> None: @@ -443,7 +448,6 @@ def _enqueue_outpost_event( { "type": "reducer_bucket", "reducer_id": reducer.id, - "bucket_start": bucket_start.isoformat(), "customer_id": customer_id, "external_customer_id": external_customer_id, }, diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 0a4e587f115..b8dc19dc615 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -139,32 +139,34 @@ async def test_event_reducer_bucket( redis: Redis, organization: Organization, meter: Meter, + customer_external_id: Customer, + customer_second: Customer, ) -> None: reducer = await create_reducer( save_fixture, organization=organization, meters=[meter] ) bucket_start = datetime(2026, 10, 7, 12, 5, tzinfo=UTC) - split_at = bucket_start + timedelta(minutes=2) - bucket_end = bucket_start + timedelta(minutes=5) - billed_at = datetime(2026, 10, 7, 12, 20, tzinfo=UTC) - for start, end, generation, sealed_at, value in ( - (bucket_start, split_at, 1, billed_at, 3), - (bucket_start, split_at, 2, billed_at, 4), - (bucket_start, split_at, 3, None, 5), - (split_at, bucket_end, 1, None, 6), + billed_start = bucket_start - timedelta(minutes=5) + external_id = customer_external_id.external_id + for start, customer, external_customer_id, sealed_at, value in ( + (bucket_start, customer_external_id, None, None, 3), + (bucket_start, None, external_id, None, 4), + (bucket_start, customer_external_id, external_id, None, 5), + (billed_start, customer_external_id, external_id, bucket_start, 6), + (bucket_start, customer_second, None, None, 7), ): await save_fixture( ReducerBucket( organization=organization, reducer=reducer, - external_customer_id="external", + customer=customer, + external_customer_id=external_customer_id, bucket_start=start, - bucket_end=end, + bucket_end=start + timedelta(minutes=5), count=1, sum=Decimal(value), min=Decimal(value), max=Decimal(value), - generation=generation, sealed_at=sealed_at, ) ) @@ -176,9 +178,8 @@ async def test_event_reducer_bucket( { "type": "reducer_bucket", "reducer_id": reducer.id, - "bucket_start": bucket_start.isoformat(), "customer_id": None, - "external_customer_id": "external", + "external_customer_id": external_id, }, ) @@ -188,37 +189,11 @@ async def test_event_reducer_bucket( "type": "reducer_bucket", "payload": { "reducer_id": str(reducer.id), - "customer_id": None, - "external_customer_id": "external", - "bucket_start": "2026-10-07T12:05:00Z", - "intervals": [ - { - "bucket_start": "2026-10-07T12:05:00Z", - "bucket_end": "2026-10-07T12:07:00Z", - "sealed": True, - "count": 2, - "sum": 7.0, - "min": 3.0, - "max": 4.0, - }, - { - "bucket_start": "2026-10-07T12:05:00Z", - "bucket_end": "2026-10-07T12:07:00Z", - "sealed": False, - "count": 1, - "sum": 5.0, - "min": 5.0, - "max": 5.0, - }, - { - "bucket_start": "2026-10-07T12:07:00Z", - "bucket_end": "2026-10-07T12:10:00Z", - "sealed": False, - "count": 1, - "sum": 6.0, - "min": 6.0, - "max": 6.0, - }, - ], + "customer_id": str(customer_external_id.id), + "external_customer_id": external_id, + "count": 3, + "sum": 12.0, + "min": 3.0, + "max": 5.0, }, } diff --git a/server/tests/reducer_bucket/test_service.py b/server/tests/reducer_bucket/test_service.py index 70970d4c7ac..5ae0dc76049 100644 --- a/server/tests/reducer_bucket/test_service.py +++ b/server/tests/reducer_bucket/test_service.py @@ -161,7 +161,6 @@ async def test_adds_to_closed_buckets( { "type": "reducer_bucket", "reducer_id": sum_reducer.id, - "bucket_start": previous_start.isoformat(), "customer_id": None, "external_customer_id": "external", }, @@ -433,7 +432,6 @@ async def test_writes_buckets( { "type": "reducer_bucket", "reducer_id": sum_reducer.id, - "bucket_start": BUCKET_START.isoformat(), "customer_id": None, "external_customer_id": "external", }, From 6690c3c141bc9c92a356d831f348e60c3279788d Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 10:15:44 +0200 Subject: [PATCH 09/12] outpost: send unbilled meter usage with its credits --- server/polar/customer_meter/repository.py | 15 +++++- server/polar/customer_meter/service.py | 15 ++++++ server/polar/outpost/schemas.py | 16 +++--- server/polar/outpost/service.py | 59 ++++++++++++++------- server/polar/reducer/repository.py | 6 ++- server/tests/customer_meter/test_service.py | 30 +++++++++++ server/tests/outpost/test_endpoints.py | 15 +++++- 7 files changed, 126 insertions(+), 30 deletions(-) diff --git a/server/polar/customer_meter/repository.py b/server/polar/customer_meter/repository.py index 040542f04e7..259907a838b 100644 --- a/server/polar/customer_meter/repository.py +++ b/server/polar/customer_meter/repository.py @@ -15,7 +15,7 @@ SortingClause, ) from polar.kit.repository.base import Options -from polar.models import Customer, CustomerMeter, Meter +from polar.models import Customer, CustomerMeter, Meter, MeterReducer from .sorting import CustomerMeterSortProperty @@ -58,6 +58,19 @@ async def get_by_customer_and_meter( ) return await self.get_one_or_none(statement) + async def get_all_by_customer_and_reducer( + self, customer_id: UUID, reducer_id: UUID + ) -> Sequence[CustomerMeter]: + statement = ( + self.get_base_statement() + .join(MeterReducer, MeterReducer.meter_id == CustomerMeter.meter_id) + .where( + CustomerMeter.customer_id == customer_id, + MeterReducer.reducer_id == reducer_id, + ) + ) + return await self.get_all(statement) + async def get_by_organization_customer_and_meter( self, organization_id: UUID, diff --git a/server/polar/customer_meter/service.py b/server/polar/customer_meter/service.py index 0836cad6a7d..7dfd5850652 100644 --- a/server/polar/customer_meter/service.py +++ b/server/polar/customer_meter/service.py @@ -29,6 +29,7 @@ from polar.models.event import EventSource from polar.outpost.stream import enqueue as enqueue_outpost_event from polar.postgres import AsyncSession +from polar.reducer.repository import ReducerRepository from polar.worker import enqueue_job from .repository import CustomerMeterRepository @@ -203,6 +204,20 @@ async def update_customer_meter( "meter_id": meter.id, }, ) + # Temporary: reducer values carry the credits until credits are reducers. + reducer = await ReducerRepository.from_session(session).get_by_meter_id( + meter.id + ) + if reducer is not None: + enqueue_outpost_event( + customer.organization_id, + { + "type": "reducer_bucket", + "reducer_id": reducer.id, + "customer_id": customer.id, + "external_customer_id": customer.external_id, + }, + ) return await repository.update(customer_meter), True diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 115ba7ab640..b08855a5a3b 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -64,7 +64,7 @@ class CustomerMeterIncomingMessage(IncomingMessageBase): class OutgoingMessageType(StrEnum): configuration = "configuration" customer_meter = "customer_meter" - reducer_bucket = "reducer_bucket" + meter_balance = "meter_balance" class OutgoingMessageBase(MessageBase): @@ -101,25 +101,27 @@ class CustomerMeterOutgoingMessage(OutgoingMessageBase): payload: CustomerMeterOutgoingMessagePayload -class ReducerBucketOutgoingMessagePayload(MessageBase): - reducer_id: UUID4 +class MeterBalanceOutgoingMessagePayload(MessageBase): + meter_id: UUID4 customer_id: UUID4 | None external_customer_id: str | None count: int sum: float min: float | None max: float | None + # Temporary: credits come from customer meters until credits are reducers. + credited_units: int -class ReducerBucketOutgoingMessage(OutgoingMessageBase): - type: typing.Literal[OutgoingMessageType.reducer_bucket] - payload: ReducerBucketOutgoingMessagePayload +class MeterBalanceOutgoingMessage(OutgoingMessageBase): + type: typing.Literal[OutgoingMessageType.meter_balance] + payload: MeterBalanceOutgoingMessagePayload type OutgoingMessage = typing.Annotated[ ConfigurationOutgoingMessage | CustomerMeterOutgoingMessage - | ReducerBucketOutgoingMessage, + | MeterBalanceOutgoingMessage, Discriminator("type"), ] diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index b0044bb0c61..5015a4970ec 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -7,6 +7,7 @@ from polar.customer_meter.repository import CustomerMeterRepository from polar.kit.db.postgres import AsyncReadSession from polar.logging import Logger +from polar.meter.aggregation import UniqueAggregation from polar.models import Organization, Reducer from polar.reducer.repository import ReducerRepository from polar.reducer_bucket.repository import ReducerBucketRepository @@ -19,11 +20,11 @@ CustomerMeterOutgoingMessagePayload, IncomingMessageAdapter, IncomingMessageType, + MeterBalanceOutgoingMessage, + MeterBalanceOutgoingMessagePayload, OutgoingMessage, OutgoingMessageAdapter, OutgoingMessageType, - ReducerBucketOutgoingMessage, - ReducerBucketOutgoingMessagePayload, ) from .stream import OutpostEvent, ReducerBucketOutpostEvent @@ -136,6 +137,12 @@ async def _send_reducer_bucket( event: ReducerBucketOutpostEvent, send_stream: MemoryObjectSendStream[OutgoingMessage], ) -> None: + reducer = await ReducerRepository.from_session(session).get_active_by_id( + event["reducer_id"], options=(selectinload(Reducer.meter_reducers),) + ) + if reducer is None or isinstance(reducer.aggregation, UniqueAggregation): + return + customer_id = event["customer_id"] external_customer_id = event["external_customer_id"] customer_repository = CustomerRepository.from_session(session) @@ -155,23 +162,39 @@ async def _send_reducer_bucket( else: identities = [(customer_id, external_customer_id)] - repository = ReducerBucketRepository.from_session(session) - count, total, minimum, maximum = await repository.get_unsealed_totals( - organization.id, event["reducer_id"], identities - ) - message = ReducerBucketOutgoingMessage( - type=OutgoingMessageType.reducer_bucket, - payload=ReducerBucketOutgoingMessagePayload( - reducer_id=event["reducer_id"], - customer_id=customer_id, - external_customer_id=external_customer_id, - count=count, - sum=total, - min=minimum, - max=maximum, - ), + count, total, minimum, maximum = await ReducerBucketRepository.from_session( + session + ).get_unsealed_totals(organization.id, reducer.id, identities) + # Temporary: credits come from customer meters until credits are reducers. + # Credits reset on renewal before the billed usage gets sealed, so balances + # are too low in between. + customer_meters = ( + await CustomerMeterRepository.from_session( + session + ).get_all_by_customer_and_reducer(customer.id, reducer.id) + if customer is not None + else [] ) - await send_stream.send(message) + credited_units = { + customer_meter.meter_id: customer_meter.credited_units + for customer_meter in customer_meters + } + for meter_id in reducer.meter_ids: + await send_stream.send( + MeterBalanceOutgoingMessage( + type=OutgoingMessageType.meter_balance, + payload=MeterBalanceOutgoingMessagePayload( + meter_id=meter_id, + customer_id=customer_id, + external_customer_id=external_customer_id, + count=count, + sum=total, + min=minimum, + max=maximum, + credited_units=credited_units.get(meter_id, 0), + ), + ) + ) outpost = OutpostService() diff --git a/server/polar/reducer/repository.py b/server/polar/reducer/repository.py index acec8a8c643..cbac8c041ee 100644 --- a/server/polar/reducer/repository.py +++ b/server/polar/reducer/repository.py @@ -73,9 +73,11 @@ def get_active_statement(self) -> Select[tuple[Reducer]]: ) return self.get_base_statement().where(Reducer.id.in_(active_reducer_ids)) - async def get_active_by_id(self, id: UUID) -> Reducer | None: + async def get_active_by_id( + self, id: UUID, *, options: Options = () + ) -> Reducer | None: return await self.get_one_or_none( - self.get_active_statement().where(Reducer.id == id) + self.get_active_statement().where(Reducer.id == id).options(*options) ) async def get_all_active_by_organization( diff --git a/server/tests/customer_meter/test_service.py b/server/tests/customer_meter/test_service.py index cc904d7928f..609cf7c9542 100644 --- a/server/tests/customer_meter/test_service.py +++ b/server/tests/customer_meter/test_service.py @@ -3,6 +3,7 @@ from decimal import Decimal import pytest +from pytest_mock import MockerFixture from polar.customer_meter.service import customer_meter as customer_meter_service from polar.event.repository import EventRepository @@ -28,6 +29,7 @@ METER_TEST_EVENT, create_event, create_meter, + create_reducer, ) @@ -857,6 +859,34 @@ async def test_max_aggregation_with_external_id( assert customer_meter.consumed_units == Decimal(80) assert updated is True + async def test_notifies_outpost_of_the_reducer( + self, + mocker: MockerFixture, + save_fixture: SaveFixture, + session: AsyncSession, + customer: Customer, + events: list[Event], + meter: Meter, + ) -> None: + enqueue_outpost_event = mocker.patch( + "polar.customer_meter.service.enqueue_outpost_event" + ) + reducer = await create_reducer( + save_fixture, organization=meter.organization, meters=[meter] + ) + + await customer_meter_service.update_customer_meter(session, customer, meter) + + enqueue_outpost_event.assert_called_with( + customer.organization_id, + { + "type": "reducer_bucket", + "reducer_id": reducer.id, + "customer_id": customer.id, + "external_customer_id": customer.external_id, + }, + ) + @pytest.mark.anyio class TestGetRolloverUnits: diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index b8dc19dc615..72606a0aeb7 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -171,6 +171,16 @@ async def test_event_reducer_bucket( ) ) + await save_fixture( + CustomerMeter( + customer=customer_external_id, + meter=meter, + consumed_units=Decimal(0), + credited_units=100, + balance=Decimal(100), + ) + ) + async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] await publish( redis, @@ -186,14 +196,15 @@ async def test_event_reducer_bucket( response = await websocket.receive_json() assert response == { - "type": "reducer_bucket", + "type": "meter_balance", "payload": { - "reducer_id": str(reducer.id), + "meter_id": str(meter.id), "customer_id": str(customer_external_id.id), "external_customer_id": external_id, "count": 3, "sum": 12.0, "min": 3.0, "max": 5.0, + "credited_units": 100, }, } From 278091c5b79a573c12cb7db286ceff2dd5ead1d3 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 10:17:07 +0200 Subject: [PATCH 10/12] outpost: rename meter_balance to meter_usage --- server/polar/outpost/schemas.py | 12 ++++++------ server/polar/outpost/service.py | 10 +++++----- server/tests/outpost/test_endpoints.py | 2 +- 3 files changed, 12 insertions(+), 12 deletions(-) diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index b08855a5a3b..1aa52082ea2 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -64,7 +64,7 @@ class CustomerMeterIncomingMessage(IncomingMessageBase): class OutgoingMessageType(StrEnum): configuration = "configuration" customer_meter = "customer_meter" - meter_balance = "meter_balance" + meter_usage = "meter_usage" class OutgoingMessageBase(MessageBase): @@ -101,7 +101,7 @@ class CustomerMeterOutgoingMessage(OutgoingMessageBase): payload: CustomerMeterOutgoingMessagePayload -class MeterBalanceOutgoingMessagePayload(MessageBase): +class MeterUsageOutgoingMessagePayload(MessageBase): meter_id: UUID4 customer_id: UUID4 | None external_customer_id: str | None @@ -113,15 +113,15 @@ class MeterBalanceOutgoingMessagePayload(MessageBase): credited_units: int -class MeterBalanceOutgoingMessage(OutgoingMessageBase): - type: typing.Literal[OutgoingMessageType.meter_balance] - payload: MeterBalanceOutgoingMessagePayload +class MeterUsageOutgoingMessage(OutgoingMessageBase): + type: typing.Literal[OutgoingMessageType.meter_usage] + payload: MeterUsageOutgoingMessagePayload type OutgoingMessage = typing.Annotated[ ConfigurationOutgoingMessage | CustomerMeterOutgoingMessage - | MeterBalanceOutgoingMessage, + | MeterUsageOutgoingMessage, Discriminator("type"), ] diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index 5015a4970ec..adbfa968238 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -20,8 +20,8 @@ CustomerMeterOutgoingMessagePayload, IncomingMessageAdapter, IncomingMessageType, - MeterBalanceOutgoingMessage, - MeterBalanceOutgoingMessagePayload, + MeterUsageOutgoingMessage, + MeterUsageOutgoingMessagePayload, OutgoingMessage, OutgoingMessageAdapter, OutgoingMessageType, @@ -181,9 +181,9 @@ async def _send_reducer_bucket( } for meter_id in reducer.meter_ids: await send_stream.send( - MeterBalanceOutgoingMessage( - type=OutgoingMessageType.meter_balance, - payload=MeterBalanceOutgoingMessagePayload( + MeterUsageOutgoingMessage( + type=OutgoingMessageType.meter_usage, + payload=MeterUsageOutgoingMessagePayload( meter_id=meter_id, customer_id=customer_id, external_customer_id=external_customer_id, diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 72606a0aeb7..02896f8fa40 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -196,7 +196,7 @@ async def test_event_reducer_bucket( response = await websocket.receive_json() assert response == { - "type": "meter_balance", + "type": "meter_usage", "payload": { "meter_id": str(meter.id), "customer_id": str(customer_external_id.id), From 2be1c74de8d5663cb95e08fe45cb27919777d203 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 10:18:29 +0200 Subject: [PATCH 11/12] outpost: stop pushing customer meters --- server/polar/customer_meter/service.py | 8 -------- server/polar/outpost/service.py | 9 --------- server/polar/outpost/stream.py | 8 +------- server/tests/outpost/test_endpoints.py | 28 -------------------------- 4 files changed, 1 insertion(+), 52 deletions(-) diff --git a/server/polar/customer_meter/service.py b/server/polar/customer_meter/service.py index 7dfd5850652..2ce1967d250 100644 --- a/server/polar/customer_meter/service.py +++ b/server/polar/customer_meter/service.py @@ -196,14 +196,6 @@ async def update_customer_meter( ) customer_meter.last_balanced_event = last_event - enqueue_outpost_event( - customer.organization_id, - { - "type": "customer_meter", - "customer_id": customer.id, - "meter_id": meter.id, - }, - ) # Temporary: reducer values carry the credits until credits are reducers. reducer = await ReducerRepository.from_session(session).get_by_meter_id( meter.id diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index adbfa968238..a410e11a9e6 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -66,15 +66,6 @@ async def handle_event( ) -> None: log.debug("Handling event", type=event["type"]) match event["type"]: - case "customer_meter": - return await self._send_customer_meter( - session, - organization, - CustomerMeterIncomingMessageCustomerPayload( - customer_id=event["customer_id"], meter_id=event["meter_id"] - ), - send_stream, - ) case "reducer_bucket": return await self._send_reducer_bucket( session, organization, event, send_stream diff --git a/server/polar/outpost/stream.py b/server/polar/outpost/stream.py index 1ecf3605b4e..0aae75bca04 100644 --- a/server/polar/outpost/stream.py +++ b/server/polar/outpost/stream.py @@ -10,12 +10,6 @@ from polar.worker import enqueue_job -class CustomerMeterOutpostEvent(typing.TypedDict): - type: typing.Literal["customer_meter"] - customer_id: uuid.UUID - meter_id: uuid.UUID - - class ReducerBucketOutpostEvent(typing.TypedDict): type: typing.Literal["reducer_bucket"] reducer_id: uuid.UUID @@ -23,7 +17,7 @@ class ReducerBucketOutpostEvent(typing.TypedDict): external_customer_id: str | None -type OutpostEvent = CustomerMeterOutpostEvent | ReducerBucketOutpostEvent +type OutpostEvent = ReducerBucketOutpostEvent async def publish( diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index 02896f8fa40..dc367e1f050 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -103,34 +103,6 @@ async def test_request_customer_meter( assert response["payload"]["credited_units"] == 100 assert response["payload"]["balance"] == 10.0 - @pytest.mark.auth(AuthSubjectFixture(subject="organization")) - async def test_event_customer_meter( - self, client: httpx.AsyncClient, redis: Redis, customer_meter: CustomerMeter - ) -> None: - async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] - await publish( - redis, - customer_meter.customer.organization_id, - { - "type": "customer_meter", - "customer_id": customer_meter.customer_id, - "meter_id": customer_meter.meter_id, - }, - ) - - response = await websocket.receive_json() - - assert response["type"] == "customer_meter" - assert response["payload"]["customer_id"] == str(customer_meter.customer_id) - assert ( - response["payload"]["external_customer_id"] - == customer_meter.customer.external_id - ) - assert response["payload"]["meter_id"] == str(customer_meter.meter_id) - assert response["payload"]["consumed_units"] == 90.0 - assert response["payload"]["credited_units"] == 100 - assert response["payload"]["balance"] == 10.0 - @pytest.mark.auth(AuthSubjectFixture(subject="organization")) async def test_event_reducer_bucket( self, From 58daa8bc58704ea0d8a2a4025115aff3850e2281 Mon Sep 17 00:00:00 2001 From: Maxime BV Date: Wed, 14 Oct 2026 10:21:00 +0200 Subject: [PATCH 12/12] outpost: replace the customer_meter request with meter_usage --- server/polar/customer_meter/repository.py | 40 ----------- server/polar/outpost/schemas.py | 45 ++++--------- server/polar/outpost/service.py | 82 ++++++++++++++--------- server/tests/outpost/test_endpoints.py | 52 ++++++++++---- 4 files changed, 102 insertions(+), 117 deletions(-) diff --git a/server/polar/customer_meter/repository.py b/server/polar/customer_meter/repository.py index 259907a838b..a6118654e63 100644 --- a/server/polar/customer_meter/repository.py +++ b/server/polar/customer_meter/repository.py @@ -71,46 +71,6 @@ async def get_all_by_customer_and_reducer( ) return await self.get_all(statement) - async def get_by_organization_customer_and_meter( - self, - organization_id: UUID, - customer_id: UUID, - meter_id: UUID, - *, - options: Options = (), - ) -> CustomerMeter | None: - statement = ( - self.get_base_statement() - .join(CustomerMeter.customer) - .where( - Customer.organization_id == organization_id, - CustomerMeter.customer_id == customer_id, - CustomerMeter.meter_id == meter_id, - ) - .options(contains_eager(CustomerMeter.customer), *options) - ) - return await self.get_one_or_none(statement) - - async def get_by_organization_external_customer_and_meter( - self, - organization_id: UUID, - external_customer_id: str, - meter_id: UUID, - *, - options: Options = (), - ) -> CustomerMeter | None: - statement = ( - self.get_base_statement() - .join(CustomerMeter.customer) - .where( - Customer.organization_id == organization_id, - Customer.external_id == external_customer_id, - CustomerMeter.meter_id == meter_id, - ) - .options(contains_eager(CustomerMeter.customer), *options) - ) - return await self.get_one_or_none(statement) - async def get_by_customer_and_meter_for_update( self, customer_id: UUID, diff --git a/server/polar/outpost/schemas.py b/server/polar/outpost/schemas.py index 1aa52082ea2..87eff0ab0d7 100644 --- a/server/polar/outpost/schemas.py +++ b/server/polar/outpost/schemas.py @@ -13,7 +13,7 @@ class MessageBase(BaseModel): class IncomingMessageType(StrEnum): configuration = "configuration" - customer_meter = "customer_meter" + meter_usage = "meter_usage" class IncomingMessageBase(MessageBase): @@ -24,35 +24,33 @@ class ConfigurationIncomingMessage(IncomingMessageBase): type: typing.Literal[IncomingMessageType.configuration] -class CustomerMeterIncomingMessagePayloadBase(MessageBase): +class MeterUsageIncomingMessagePayloadBase(MessageBase): meter_id: UUID4 -class CustomerMeterIncomingMessageCustomerPayload( - CustomerMeterIncomingMessagePayloadBase -): +class MeterUsageIncomingMessageCustomerPayload(MeterUsageIncomingMessagePayloadBase): customer_id: UUID4 -class CustomerMeterIncomingMessageExternalCustomerPayload( - CustomerMeterIncomingMessagePayloadBase +class MeterUsageIncomingMessageExternalCustomerPayload( + MeterUsageIncomingMessagePayloadBase ): external_customer_id: str -type CustomerMeterIncomingMessagePayload = ( - CustomerMeterIncomingMessageCustomerPayload - | CustomerMeterIncomingMessageExternalCustomerPayload +type MeterUsageIncomingMessagePayload = ( + MeterUsageIncomingMessageCustomerPayload + | MeterUsageIncomingMessageExternalCustomerPayload ) -class CustomerMeterIncomingMessage(IncomingMessageBase): - type: typing.Literal[IncomingMessageType.customer_meter] - payload: CustomerMeterIncomingMessagePayload +class MeterUsageIncomingMessage(IncomingMessageBase): + type: typing.Literal[IncomingMessageType.meter_usage] + payload: MeterUsageIncomingMessagePayload type IncomingMessage = typing.Annotated[ - ConfigurationIncomingMessage | CustomerMeterIncomingMessage, Discriminator("type") + ConfigurationIncomingMessage | MeterUsageIncomingMessage, Discriminator("type") ] IncomingMessageAdapter: TypeAdapter[IncomingMessage] = TypeAdapter(IncomingMessage) @@ -63,7 +61,6 @@ class CustomerMeterIncomingMessage(IncomingMessageBase): class OutgoingMessageType(StrEnum): configuration = "configuration" - customer_meter = "customer_meter" meter_usage = "meter_usage" @@ -87,20 +84,6 @@ class ConfigurationOutgoingMessage(OutgoingMessageBase): payload: ConfigurationOutgoingMessagePayload -class CustomerMeterOutgoingMessagePayload(MessageBase): - customer_id: UUID4 - external_customer_id: str | None - meter_id: UUID4 - consumed_units: float - credited_units: int - balance: float - - -class CustomerMeterOutgoingMessage(OutgoingMessageBase): - type: typing.Literal[OutgoingMessageType.customer_meter] - payload: CustomerMeterOutgoingMessagePayload - - class MeterUsageOutgoingMessagePayload(MessageBase): meter_id: UUID4 customer_id: UUID4 | None @@ -119,9 +102,7 @@ class MeterUsageOutgoingMessage(OutgoingMessageBase): type OutgoingMessage = typing.Annotated[ - ConfigurationOutgoingMessage - | CustomerMeterOutgoingMessage - | MeterUsageOutgoingMessage, + ConfigurationOutgoingMessage | MeterUsageOutgoingMessage, Discriminator("type"), ] diff --git a/server/polar/outpost/service.py b/server/polar/outpost/service.py index a410e11a9e6..4272abedd24 100644 --- a/server/polar/outpost/service.py +++ b/server/polar/outpost/service.py @@ -1,3 +1,6 @@ +import uuid +from collections.abc import Sequence + import structlog from anyio.streams.memory import MemoryObjectSendStream from pydantic import ValidationError @@ -14,12 +17,10 @@ from polar.reducer_bucket.service import get_customer_identities from .schemas import ( - CustomerMeterIncomingMessageCustomerPayload, - CustomerMeterIncomingMessagePayload, - CustomerMeterOutgoingMessage, - CustomerMeterOutgoingMessagePayload, IncomingMessageAdapter, IncomingMessageType, + MeterUsageIncomingMessageCustomerPayload, + MeterUsageIncomingMessagePayload, MeterUsageOutgoingMessage, MeterUsageOutgoingMessagePayload, OutgoingMessage, @@ -52,8 +53,8 @@ async def handle_incoming_message( return await self._send_configuration( session, organization, send_stream ) - case IncomingMessageType.customer_meter: - return await self._send_customer_meter( + case IncomingMessageType.meter_usage: + return await self._send_requested_meter_usage( session, organization, incoming.payload, send_stream ) @@ -89,37 +90,36 @@ async def _send_configuration( ) await send_stream.send(message) - async def _send_customer_meter( + async def _send_requested_meter_usage( self, session: AsyncReadSession, organization: Organization, - payload: CustomerMeterIncomingMessagePayload, + payload: MeterUsageIncomingMessagePayload, send_stream: MemoryObjectSendStream[OutgoingMessage], ) -> None: - customer_meter_repository = CustomerMeterRepository.from_session(session) - if isinstance(payload, CustomerMeterIncomingMessageCustomerPayload): - customer_meter = ( - await customer_meter_repository.get_by_organization_customer_and_meter( - organization.id, payload.customer_id, payload.meter_id - ) - ) - else: - customer_meter = await customer_meter_repository.get_by_organization_external_customer_and_meter( - organization.id, payload.external_customer_id, payload.meter_id - ) - if customer_meter is None: + reducer = await ReducerRepository.from_session(session).get_by_meter_id( + payload.meter_id + ) + if reducer is None or reducer.organization_id != organization.id: log.warning( - "Customer meter not found", + "Meter not found", organization_id=organization.id, - **payload.model_dump(), + meter_id=payload.meter_id, ) return - - message = CustomerMeterOutgoingMessage( - type=OutgoingMessageType.customer_meter, - payload=CustomerMeterOutgoingMessagePayload.model_validate(customer_meter), + if isinstance(payload, MeterUsageIncomingMessageCustomerPayload): + customer_id, external_customer_id = payload.customer_id, None + else: + customer_id, external_customer_id = None, payload.external_customer_id + await self._send_meter_usage( + session, + organization, + reducer, + [payload.meter_id], + customer_id, + external_customer_id, + send_stream, ) - await send_stream.send(message) async def _send_reducer_bucket( self, @@ -131,11 +131,31 @@ async def _send_reducer_bucket( reducer = await ReducerRepository.from_session(session).get_active_by_id( event["reducer_id"], options=(selectinload(Reducer.meter_reducers),) ) - if reducer is None or isinstance(reducer.aggregation, UniqueAggregation): + if reducer is None: + return + await self._send_meter_usage( + session, + organization, + reducer, + reducer.meter_ids, + event["customer_id"], + event["external_customer_id"], + send_stream, + ) + + async def _send_meter_usage( + self, + session: AsyncReadSession, + organization: Organization, + reducer: Reducer, + meter_ids: Sequence[uuid.UUID], + customer_id: uuid.UUID | None, + external_customer_id: str | None, + send_stream: MemoryObjectSendStream[OutgoingMessage], + ) -> None: + if isinstance(reducer.aggregation, UniqueAggregation): return - customer_id = event["customer_id"] - external_customer_id = event["external_customer_id"] customer_repository = CustomerRepository.from_session(session) if customer_id is not None: customer = await customer_repository.get_by_id_and_organization( @@ -170,7 +190,7 @@ async def _send_reducer_bucket( customer_meter.meter_id: customer_meter.credited_units for customer_meter in customer_meters } - for meter_id in reducer.meter_ids: + for meter_id in meter_ids: await send_stream.send( MeterUsageOutgoingMessage( type=OutgoingMessageType.meter_usage, diff --git a/server/tests/outpost/test_endpoints.py b/server/tests/outpost/test_endpoints.py index dc367e1f050..a3fef31a171 100644 --- a/server/tests/outpost/test_endpoints.py +++ b/server/tests/outpost/test_endpoints.py @@ -70,20 +70,41 @@ async def test_request_configuration( @pytest.mark.auth(AuthSubjectFixture(subject="organization")) @pytest.mark.parametrize("customer_key", ["customer_id", "external_customer_id"]) - async def test_request_customer_meter( + async def test_request_meter_usage( self, + save_fixture: SaveFixture, client: httpx.AsyncClient, + organization: Organization, customer_meter: CustomerMeter, customer_key: str, ) -> None: + customer = customer_meter.customer + reducer = await create_reducer( + save_fixture, organization=organization, meters=[customer_meter.meter] + ) + bucket_start = datetime(2026, 10, 7, 12, 5, tzinfo=UTC) + await save_fixture( + ReducerBucket( + organization=organization, + reducer=reducer, + external_customer_id=customer.external_id, + bucket_start=bucket_start, + bucket_end=bucket_start + timedelta(minutes=5), + count=1, + sum=Decimal(4), + min=Decimal(4), + max=Decimal(4), + ) + ) customer_value = { - "customer_id": str(customer_meter.customer_id), - "external_customer_id": customer_meter.customer.external_id, + "customer_id": str(customer.id), + "external_customer_id": customer.external_id, }[customer_key] + async with httpx_ws.aconnect_ws("/v1/outpost/", client=client) as websocket: # type: ignore[var-annotated] await websocket.send_json( { - "type": "customer_meter", + "type": "meter_usage", "payload": { customer_key: customer_value, "meter_id": str(customer_meter.meter_id), @@ -92,16 +113,19 @@ async def test_request_customer_meter( ) response = await websocket.receive_json() - assert response["type"] == "customer_meter" - assert response["payload"]["customer_id"] == str(customer_meter.customer_id) - assert ( - response["payload"]["external_customer_id"] - == customer_meter.customer.external_id - ) - assert response["payload"]["meter_id"] == str(customer_meter.meter_id) - assert response["payload"]["consumed_units"] == 90.0 - assert response["payload"]["credited_units"] == 100 - assert response["payload"]["balance"] == 10.0 + assert response == { + "type": "meter_usage", + "payload": { + "meter_id": str(customer_meter.meter_id), + "customer_id": str(customer.id), + "external_customer_id": customer.external_id, + "count": 1, + "sum": 4.0, + "min": 4.0, + "max": 4.0, + "credited_units": 100, + }, + } @pytest.mark.auth(AuthSubjectFixture(subject="organization")) async def test_event_reducer_bucket(