feat: add enterprise write buffer and data access layer with batched activity and usage record flushing

mywebdav/billing/usage_tracker.py

@@ -1,11 +1,34 @@
 import uuid
+import logging
 from datetime import datetime, date, timezone, timedelta
+from typing import List, Dict
 
 from tortoise.transactions import in_transaction
 
 from .models import UsageRecord, UsageAggregate
 from ..models import User
 
+logger = logging.getLogger(__name__)
+
+
+async def _flush_usage_records(records: List[Dict]):
+    try:
+        async with in_transaction():
+            for record in records:
+                await UsageRecord.create(
+                    user_id=record["user_id"],
+                    record_type=record["record_type"],
+                    amount_bytes=record["amount_bytes"],
+                    resource_type=record.get("resource_type"),
+                    resource_id=record.get("resource_id"),
+                    idempotency_key=record["idempotency_key"],
+                    metadata=record.get("metadata"),
+                )
+        logger.debug(f"Flushed {len(records)} usage records")
+    except Exception as e:
+        logger.error(f"Failed to flush usage records: {e}")
+        raise
+
 
 class UsageTracker:
     @staticmethod
@@ -18,15 +41,31 @@ class UsageTracker:
     ):
         idempotency_key = f"storage_{user.id}_{datetime.now(timezone.utc).timestamp()}_{uuid.uuid4().hex[:8]}"
 
-        await UsageRecord.create(
-            user=user,
-            record_type="storage",
-            amount_bytes=amount_bytes,
-            resource_type=resource_type,
-            resource_id=resource_id,
-            idempotency_key=idempotency_key,
-            metadata=metadata,
-        )
+        try:
+            from ..enterprise.write_buffer import get_write_buffer, WriteType
+            buffer = get_write_buffer()
+            await buffer.buffer(
+                WriteType.USAGE_RECORD,
+                {
+                    "user_id": user.id,
+                    "record_type": "storage",
+                    "amount_bytes": amount_bytes,
+                    "resource_type": resource_type,
+                    "resource_id": resource_id,
+                    "idempotency_key": idempotency_key,
+                    "metadata": metadata,
+                },
+            )
+        except RuntimeError:
+            await UsageRecord.create(
+                user=user,
+                record_type="storage",
+                amount_bytes=amount_bytes,
+                resource_type=resource_type,
+                resource_id=resource_id,
+                idempotency_key=idempotency_key,
+                metadata=metadata,
+            )
 
     @staticmethod
     async def track_bandwidth(
@@ -40,15 +79,31 @@ class UsageTracker:
         record_type = f"bandwidth_{direction}"
         idempotency_key = f"{record_type}_{user.id}_{datetime.now(timezone.utc).timestamp()}_{uuid.uuid4().hex[:8]}"
 
-        await UsageRecord.create(
-            user=user,
-            record_type=record_type,
-            amount_bytes=amount_bytes,
-            resource_type=resource_type,
-            resource_id=resource_id,
-            idempotency_key=idempotency_key,
-            metadata=metadata,
-        )
+        try:
+            from ..enterprise.write_buffer import get_write_buffer, WriteType
+            buffer = get_write_buffer()
+            await buffer.buffer(
+                WriteType.USAGE_RECORD,
+                {
+                    "user_id": user.id,
+                    "record_type": record_type,
+                    "amount_bytes": amount_bytes,
+                    "resource_type": resource_type,
+                    "resource_id": resource_id,
+                    "idempotency_key": idempotency_key,
+                    "metadata": metadata,
+                },
+            )
+        except RuntimeError:
+            await UsageRecord.create(
+                user=user,
+                record_type=record_type,
+                amount_bytes=amount_bytes,
+                resource_type=resource_type,
+                resource_id=resource_id,
+                idempotency_key=idempotency_key,
+                metadata=metadata,
+            )
 
     @staticmethod
     async def aggregate_daily_usage(user: User, target_date: date = None):