saas/usage_service.py

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}