File size: 2,167 Bytes
383cb38 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | import datetime
from datetime import timezone
import logging
from typing import Dict, Optional
from ecommerce.models import Subscription
from saas.models import UsageEvent
from sqlalchemy import func
from sqlalchemy.orm import Session
logger = logging.getLogger(__name__)
class UsageMeteringService:
def __init__(self, db: Session):
self.db = db
def ingest_event(self, subscription_id: str, event_type: str, quantity: float = 1.0, metadata: dict = None) -> UsageEvent:
"""
Record a usage event and update the subscription cache.
"""
sub = self.db.query(Subscription).filter(Subscription.id == subscription_id).first()
if not sub:
logger.error(f"Subscription {subscription_id} not found")
return None
# 1. Store Raw Event
event = UsageEvent(
workspace_id=sub.workspace_id,
subscription_id=subscription_id,
event_type=event_type,
quantity=quantity,
metadata_json=metadata,
timestamp=datetime.datetime.now(timezone.utc)
)
self.db.add(event)
# 2. Update Cache (Atomic Increment approach is better, but JSON update is MVP)
current_usage = sub.current_period_usage or {}
current_val = current_usage.get(event_type, 0.0)
current_usage[event_type] = current_val + quantity
# Re-assign to trigger SQL update for JSON
sub.current_period_usage = dict(current_usage)
self.db.commit()
return event
def get_aggregated_usage(self, subscription_id: str, start_date: datetime.datetime, end_date: datetime.datetime) -> Dict[str, float]:
"""
Sum usage by type for a given period.
"""
results = self.db.query(
UsageEvent.event_type,
func.sum(UsageEvent.quantity)
).filter(
UsageEvent.subscription_id == subscription_id,
UsageEvent.timestamp >= start_date,
UsageEvent.timestamp <= end_date
).group_by(UsageEvent.event_type).all()
return {r[0]: r[1] for r in results}
|