""" Widget Stripe Billing Integration (W-K09) Production-ready implementation for subscription billing with usage-based overages. """ import asyncio import logging from dataclasses import dataclass from datetime import datetime, timedelta from decimal import Decimal, ROUND_UP from enum import Enum from typing import Optional, Dict, Any, Protocol from functools import wraps import math import stripe import aioredis from sqlalchemy import ( Column, Integer, String, DateTime, Numeric, ForeignKey, create_engine, Enum as SQLEnum, Boolean ) from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import sessionmaker, Session, relationship from sqlalchemy.sql import func # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # ============================================================================ # Configuration # ============================================================================ class Config: """Application configuration.""" STRIPE_SECRET_KEY: str = "sk_test_..." # Use environment variables in production STRIPE_WEBHOOK_SECRET: str = "whsec_..." REDIS_URL: str = "redis://localhost:6379/0" DATABASE_URL: str = "postgresql://user:pass@localhost/dbname" # Billing configuration CURRENCY: str = "aud" OVERAGE_RATE_PER_MINUTE: Decimal = Decimal("0.15") GRACE_PERIOD_DAYS: int = 3 # Tier configurations TIERS = { "basic": {"price": 19700, "minutes": 100}, # $197.00 in cents "professional": {"price": 49700, "minutes": 500}, "enterprise": {"price": 99700, "minutes": 2000}, } # ============================================================================ # Domain Models # ============================================================================ class SubscriptionTier(str, Enum): """Available subscription tiers.""" BASIC = "basic" PROFESSIONAL = "professional" ENTERPRISE = "enterprise" class SubscriptionStatus(str, Enum): """Stripe subscription statuses with business logic mapping.""" ACTIVE = "active" PAST_DUE = "past_due" CANCELED = "canceled" UNPAID = "unpaid" INCOMPLETE = "incomplete" @dataclass(frozen=True) class TierConfig: """Immutable tier configuration.""" name: str price_cents: int included_minutes: int stripe_price_id: Optional[str] = None @dataclass class UsageMetrics: """Current usage metrics for a widget.""" widget_id: str current_minutes: int tier_limit: int overage_minutes: int estimated_overage_cost: Decimal # ============================================================================ # Database Models (SQLAlchemy) # ============================================================================ Base = declarative_base() class WidgetSubscription(Base): """Database model for widget subscriptions.""" __tablename__ = "widget_subscriptions" id = Column(Integer, primary_key=True) widget_id = Column(String(255), unique=True, nullable=False, index=True) stripe_customer_id = Column(String(255), nullable=False) stripe_subscription_id = Column(String(255), unique=True, nullable=True) tier = Column(SQLEnum(SubscriptionTier), nullable=False) status = Column(SQLEnum(SubscriptionStatus), default=SubscriptionStatus.INCOMPLETE) # Grace period tracking past_due_since = Column(DateTime, nullable=True) grace_period_end = Column(DateTime, nullable=True) deactivated_at = Column(DateTime, nullable=True) # Timestamps created_at = Column(DateTime, default=func.now()) updated_at = Column(DateTime, default=func.now(), onupdate=func.now()) # Relationships usage_records = relationship("UsageRecord", back_populates="subscription") def is_active(self) -> bool: """Check if widget should be active based on billing status.""" if self.status == SubscriptionStatus.ACTIVE: return True if self.status == SubscriptionStatus.PAST_DUE: if self.grace_period_end and datetime.utcnow() < self.grace_period_end: return True return False class UsageRecord(Base): """Audit trail for usage tracking (synced from Redis).""" __tablename__ = "usage_records" id = Column(Integer, primary_key=True) subscription_id = Column(Integer, ForeignKey("widget_subscriptions.id"), nullable=False) billing_period_start = Column(DateTime, nullable=False) billing_period_end = Column(DateTime, nullable=False) minutes_used = Column(Integer, default=0) overage_minutes = Column(Integer, default=0) overage_charged = Column(Numeric(10, 2), default=0) # Metadata recorded_at = Column(DateTime, default=func.now()) subscription = relationship("WidgetSubscription", back_populates="usage_records") # ============================================================================ # Infrastructure & Repository Layer # ============================================================================ class RedisRepository: """Redis repository for high-performance usage tracking.""" def __init__(self, redis_url: str): self.redis: Optional[aioredis.Redis] = None self.redis_url = redis_url async def connect(self): """Establish Redis connection.""" self.redis = await aioredis.from_url( self.redis_url, encoding="utf-8", decode_responses=True ) async def disconnect(self): """Close Redis connection.""" if self.redis: await self.redis.close() def _get_counter_key(self, widget_id: str, period: str) -> str: """Generate Redis key for counter.""" return f"widget:{widget_id}:usage:{period}" async def increment_minutes(self, widget_id: str, minutes: int = 1) -> int: """ Atomically increment minute counter. Returns new total. """ if not self.redis: raise RuntimeError("Redis not connected") period = datetime.utcnow().strftime("%Y-%m") key = self._get_counter_key(widget_id, period) # Use INCRBY for atomic operation new_value = await self.redis.incrby(key, minutes) # Set expiry to end of month + 1 day (for billing sync buffer) expiry = self._get_month_end() + timedelta(days=1) await self.redis.expireat(key, int(expiry.timestamp())) return new_value async def get_current_minutes(self, widget_id: str) -> int: """Get current billing period minutes.""" if not self.redis: raise RuntimeError("Redis not connected") period = datetime.utcnow().strftime("%Y-%m") key = self._get_counter_key(widget_id, period) value = await self.redis.get(key) return int(value) if value else 0 def _get_month_end(self) -> datetime: """Get last day of current month.""" today = datetime.utcnow() if today.month == 12: return datetime(today.year + 1, 1, 1) - timedelta(seconds=1) return datetime(today.year, today.month + 1, 1) - timedelta(seconds=1) class DatabaseRepository: """PostgreSQL repository for persistent storage.""" def __init__(self, database_url: str): self.engine = create_engine(database_url) self.SessionLocal = sessionmaker(bind=self.engine) Base.metadata.create_all(bind=self.engine) def get_session(self) -> Session: """Get database session.""" return self.SessionLocal() def get_subscription(self, widget_id: str) -> Optional[WidgetSubscription]: """Fetch subscription by widget ID.""" with self.get_session() as session: return session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == widget_id ).first() def update_subscription_status( self, widget_id: str, status: SubscriptionStatus, stripe_subscription_id: Optional[str] = None ) -> None: """Update subscription status.""" with self.get_session() as session: sub = session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == widget_id ).first() if sub: sub.status = status if stripe_subscription_id: sub.stripe_subscription_id = stripe_subscription_id # Handle grace period logic if status == SubscriptionStatus.PAST_DUE: sub.past_due_since = datetime.utcnow() sub.grace_period_end = datetime.utcnow() + timedelta( days=Config.GRACE_PERIOD_DAYS ) elif status == SubscriptionStatus.ACTIVE: sub.past_due_since = None sub.grace_period_end = None session.commit() logger.info(f"Updated {widget_id} status to {status}") def record_usage_sync( self, widget_id: str, minutes: int, overage_minutes: int, overage_cost: Decimal ) -> None: """Persist usage record for audit trail.""" with self.get_session() as session: sub = session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == widget_id ).first() if sub: now = datetime.utcnow() record = UsageRecord( subscription_id=sub.id, billing_period_start=now.replace(day=1, hour=0, minute=0, second=0), billing_period_end=self._get_month_end(now), minutes_used=minutes, overage_minutes=overage_minutes, overage_charged=overage_cost ) session.add(record) session.commit() def _get_month_end(self, dt: datetime) -> datetime: """Calculate end of month.""" if dt.month == 12: return datetime(dt.year + 1, 1, 1) - timedelta(seconds=1) return datetime(dt.year, dt.month + 1, 1) - timedelta(seconds=1) # ============================================================================ # Service Layer # ============================================================================ class BillingService: """Handles Stripe integration and subscription management.""" def __init__(self, db_repo: DatabaseRepository): self.db = db_repo stripe.api_key = Config.STRIPE_SECRET_KEY def create_subscription( self, widget_id: str, tier: SubscriptionTier, payment_method_id: str ) -> Dict[str, Any]: """ Create new subscription with usage-based billing setup. Args: widget_id: Unique widget identifier tier: Selected subscription tier payment_method_id: Stripe payment method ID Returns: Dictionary with customer_id and subscription_id """ try: # Create or retrieve customer customer = stripe.Customer.create( payment_method=payment_method_id, invoice_settings={"default_payment_method": payment_method_id}, metadata={"widget_id": widget_id} ) # Get price ID for tier (in production, map from config) price_id = self._get_stripe_price_id(tier) # Create subscription with metered usage item for overages subscription = stripe.Subscription.create( customer=customer.id, items=[ {"price": price_id}, # Base subscription { "price_data": { "currency": Config.CURRENCY, "product": self._get_overage_product_id(), "recurring": {"interval": "month"}, "unit_amount_decimal": str( Config.OVERAGE_RATE_PER_MINUTE * 100 ), # Convert to cents }, "billing_thresholds": { "usage_gte": 1 # Bill immediately when usage reported } } ], metadata={"widget_id": widget_id, "tier": tier.value}, expand=["latest_invoice.payment_intent"] ) # Persist to database with self.db.get_session() as session: new_sub = WidgetSubscription( widget_id=widget_id, stripe_customer_id=customer.id, stripe_subscription_id=subscription.id, tier=tier, status=SubscriptionStatus.ACTIVE ) session.add(new_sub) session.commit() logger.info(f"Created subscription for {widget_id}: {subscription.id}") return { "customer_id": customer.id, "subscription_id": subscription.id, "status": subscription.status } except stripe.error.StripeError as e: logger.error(f"Stripe error creating subscription: {e}") raise BillingError(f"Failed to create subscription: {e}") from e def cancel_subscription(self, widget_id: str, immediate: bool = False) -> None: """ Cancel subscription and schedule deactivation. Args: widget_id: Widget identifier immediate: If True, cancel immediately; otherwise at period end """ sub = self.db.get_subscription(widget_id) if not sub or not sub.stripe_subscription_id: raise BillingError("Subscription not found") try: if immediate: stripe.Subscription.delete(sub.stripe_subscription_id) self.db.update_subscription_status( widget_id, SubscriptionStatus.CANCELED ) else: stripe.Subscription.modify( sub.stripe_subscription_id, cancel_at_period_end=True ) logger.info(f"Cancelled subscription for {widget_id}") except stripe.error.StripeError as e: logger.error(f"Error cancelling subscription: {e}") raise def report_usage(self, widget_id: str, minutes: int) -> None: """ Report usage to Stripe for metered billing. Args: widget_id: Widget identifier minutes: Number of overage minutes to bill """ sub = self.db.get_subscription(widget_id) if not sub or not sub.stripe_subscription_id: return try: # Find the metered price item (overage item) subscription = stripe.Subscription.retrieve(sub.stripe_subscription_id) usage_item = None for item in subscription.get("items", {}).get("data", []): if item.get("price", {}).get("recurring", {}).get("usage_type") == "metered": usage_item = item break if usage_item and minutes > 0: # Report usage to Stripe stripe.UsageRecord.create( subscription_item=usage_item["id"], quantity=minutes, timestamp=int(datetime.utcnow().timestamp()), action="set" if minutes == 0 else "increment" ) logger.info(f"Reported {minutes} overage minutes for {widget_id}") except stripe.error.StripeError as e: logger.error(f"Error reporting usage: {e}") raise def _get_stripe_price_id(self, tier: SubscriptionTier) -> str: """Map tier to Stripe price ID (implement based on your Stripe setup).""" # In production, store these in database or environment variables price_map = { SubscriptionTier.BASIC: "price_basic_...", SubscriptionTier.PROFESSIONAL: "price_prof_...", SubscriptionTier.ENTERPRISE: "price_ent_..." } return price_map.get(tier, "price_basic_...") def _get_overage_product_id(self) -> str: """Get Stripe product ID for overage billing.""" return "prod_overage_..." # Configure in Stripe dashboard class UsageService: """Handles minute tracking with Redis and PostgreSQL sync.""" def __init__( self, redis_repo: RedisRepository, db_repo: DatabaseRepository, billing_service: BillingService ): self.redis = redis_repo self.db = db_repo self.billing = billing_service async def track_minutes(self, widget_id: str, minutes: int = 1) -> UsageMetrics: """ Track usage and calculate overages. Args: widget_id: Widget identifier minutes: Minutes to add (default 1) Returns: UsageMetrics with current usage and overage info """ # Get subscription details sub = self.db.get_subscription(widget_id) if not sub: raise BillingError(f"No subscription found for {widget_id}") # Check if widget is active if not sub.is_active(): raise WidgetInactiveError(f"Widget {widget_id} is inactive due to billing") # Increment Redis counter (atomic) current_minutes = await self.redis.increment_minutes(widget_id, minutes) # Get tier limit tier_config = Config.TIERS.get(sub.tier.value, Config.TIERS["basic"]) included_minutes = tier_config["minutes"] # Calculate overage (round up partial minutes) overage = max(0, current_minutes - included_minutes) # Calculate cost (rounded up to nearest minute as per requirements) overage_cost = Decimal(0) if overage > 0: overage_cost = Decimal(overage) * Config.OVERAGE_RATE_PER_MINUTE return UsageMetrics( widget_id=widget_id, current_minutes=current_minutes, tier_limit=included_minutes, overage_minutes=overage, estimated_overage_cost=overage_cost ) async def sync_to_database(self, widget_id: str) -> None: """ Sync current Redis counter to PostgreSQL for persistence. Should be called periodically (e.g., every 5 minutes) and at billing cycle end. """ try: current_minutes = await self.redis.get_current_minutes(widget_id) sub = self.db.get_subscription(widget_id) if not sub: return tier_config = Config.TIERS.get(sub.tier.value, Config.TIERS["basic"]) included = tier_config["minutes"] overage = max(0, current_minutes - included) # Calculate billing amount (round up as per requirements) overage_cost = Decimal(overage) * Config.OVERAGE_RATE_PER_MINUTE # Persist to database self.db.record_usage_sync(widget_id, current_minutes, overage, overage_cost) # Report to Stripe if overage exists if overage > 0: self.billing.report_usage(widget_id, overage) logger.info(f"Synced usage for {widget_id}: {current_minutes} mins") except Exception as e: logger.error(f"Error syncing usage for {widget_id}: {e}") raise class GracePeriodService: """Handles grace period logic and widget deactivation.""" def __init__(self, db_repo: DatabaseRepository): self.db = db_repo def check_and_deactivate(self) -> None: """ Background job to check grace periods and deactivate widgets. Should run daily via cron or Celery. """ with self.db.get_session() as session: # Find subscriptions past grace period past_grace = session.query(WidgetSubscription).filter( WidgetSubscription.status == SubscriptionStatus.PAST_DUE, WidgetSubscription.grace_period_end <= datetime.utcnow(), WidgetSubscription.deactivated_at.is_(None) ).all() for sub in past_grace: sub.deactivated_at = datetime.utcnow() logger.warning(f"Deactivated widget {sub.widget_id} due to non-payment") session.commit() def handle_stripe_webhook(self, payload: bytes, sig_header: str) -> None: """ Process Stripe webhook events for subscription status changes. Args: payload: Request body sig_header: Stripe signature header """ try: event = stripe.Webhook.construct_event( payload, sig_header, Config.STRIPE_WEBHOOK_SECRET ) event_type = event["type"] data = event["data"]["object"] if event_type == "invoice.payment_failed": # Enter grace period widget_id = data.get("metadata", {}).get("widget_id") if widget_id: self.db.update_subscription_status( widget_id, SubscriptionStatus.PAST_DUE ) logger.warning(f"Payment failed for {widget_id}, grace period started") elif event_type == "invoice.payment_succeeded": # Exit grace period if applicable widget_id = data.get("metadata", {}).get("widget_id") if widget_id: self.db.update_subscription_status( widget_id, SubscriptionStatus.ACTIVE ) elif event_type == "customer.subscription.deleted": # Immediate cancellation widget_id = data.get("metadata", {}).get("widget_id") if widget_id: self.db.update_subscription_status( widget_id, SubscriptionStatus.CANCELED ) except stripe.error.SignatureVerificationError as e: logger.error(f"Invalid webhook signature: {e}") raise except Exception as e: logger.error(f"Error processing webhook: {e}") raise # ============================================================================ # Custom Exceptions # ============================================================================ class BillingError(Exception): """Base billing exception.""" pass class WidgetInactiveError(BillingError): """Widget is inactive due to billing issues.""" pass # ============================================================================ # Example Usage & Test Scenarios # ============================================================================ async def main(): """Example usage and test scenarios.""" # Initialize infrastructure redis_repo = RedisRepository(Config.REDIS_URL) await redis_repo.connect() db_repo = DatabaseRepository(Config.DATABASE_URL) billing_service = BillingService(db_repo) usage_service = UsageService(redis_repo, db_repo, billing_service) grace_service = GracePeriodService(db_repo) try: # Test Scenario 1: Subscribe and verify activation print("=== Test 1: Subscription Creation ===") # Note: In production, use real payment method from Stripe Elements try: result = billing_service.create_subscription( widget_id="widget_001", tier=SubscriptionTier.BASIC, payment_method_id="pm_test_card" # Test payment method ) print(f"Subscription created: {result}") except BillingError as e: print(f"Expected error in demo (no real Stripe key): {e}") # Test Scenario 2: Usage tracking with overage calculation print("\n=== Test 2: Usage Tracking ===") # Simulate database entry for demo with db_repo.get_session() as session: # Clean up if exists existing = session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == "widget_demo" ).first() if existing: session.delete(existing) session.commit() # Create demo subscription demo_sub = WidgetSubscription( widget_id="widget_demo", stripe_customer_id="cus_demo", tier=SubscriptionTier.BASIC, status=SubscriptionStatus.ACTIVE ) session.add(demo_sub) session.commit() # Track 150 minutes (50 overage for Basic tier) for i in range(150): metrics = await usage_service.track_minutes("widget_demo", 1) if i % 50 == 0: print(f"After {i+1} minutes: {metrics}") # Verify overage calculation (150 - 100 = 50 overage minutes) # Cost: 50 * $0.15 = $7.50 final_metrics = await usage_service.track_minutes("widget_demo", 0) # Just read print(f"Final metrics: {final_metrics}") assert final_metrics.overage_minutes == 50 assert final_metrics.estimated_overage_cost == Decimal("7.50") # Test Scenario 3: Grace period and deactivation print("\n=== Test 3: Grace Period Logic ===") # Simulate past due db_repo.update_subscription_status( "widget_demo", SubscriptionStatus.PAST_DUE ) sub = db_repo.get_subscription("widget_demo") print(f"Status: {sub.status}, Grace period ends: {sub.grace_period_end}") # Check if active (should be true during grace period) print(f"Is active (within grace period): {sub.is_active()}") # Simulate grace period expiration with db_repo.get_session() as session: sub = session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == "widget_demo" ).first() sub.grace_period_end = datetime.utcnow() - timedelta(days=1) session.commit() # Run deactivation job grace_service.check_and_deactivate() sub = db_repo.get_subscription("widget_demo") print(f"Deactivated at: {sub.deactivated_at}") print(f"Is active (after grace period): {sub.is_active()}") # Test Scenario 4: Overage rounding print("\n=== Test 4: Overage Calculation ===") # Reset counter await redis_repo.redis.delete(f"widget:widget_demo:usage:{datetime.utcnow().strftime('%Y-%m')}") # Add 100.3 minutes (should round up to 101 for billing) # Note: Since we track integers, this simulates partial minute tracking # In production, you might track seconds and convert # Cleanup with db_repo.get_session() as session: session.query(WidgetSubscription).filter( WidgetSubscription.widget_id == "widget_demo" ).delete() session.commit() finally: await redis_repo.disconnect() if __name__ == "__main__": asyncio.run(main())