File size: 6,403 Bytes
bde2f3a | 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 | """
DataBus Provider Core — Infrastructure & Base Classes
======================================================
Circuit breakers, rate limiters, quota tracking, and core provider classes.
This module contains NO external API implementations.
"""
import logging
import os
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
from typing import Any
logger = logging.getLogger("databus.providers.core")
class ProviderTier(Enum):
"""Provider tiers: LOCAL > FREE_API > FREEMIUM > PAID"""
LOCAL = "local" # Our own data — instant, free, unlimited
FREE_API = "free_api" # Free external API — no key needed
FREEMIUM = "freemium" # Free tier with key — limited credits
PAID = "paid" # Paid API — precious credits
@dataclass
class Provider:
"""A single data source in a fallback chain."""
name: str
tier: ProviderTier
fetch_fn: Callable = field(repr=False)
weight: float = 1.0 # Higher = preferred within tier
rate_limit_rps: float = 1.0
monthly_quota: int = 0 # 0 = unlimited
requires_key: bool = False
key_env: str = ""
timeout: float = 15.0
is_local: bool = False # True if this provider uses our own data (no external API)
description: str = "" # Human-readable description
# Circuit breaker
failure_threshold: int = 5
recovery_timeout: float = 60.0
@dataclass
class ProviderChain:
"""A fallback chain for a specific data type."""
data_type: str
providers: list[Provider]
description: str = ""
async def fetch(self, vault: Any = None, cache: Any = None, **kwargs: Any) -> Any | None:
"""Try each provider in order until one succeeds.
Smart fallback: when paid provider quota is >80% used, skip to free/local
alternatives first to conserve credits for critical queries.
"""
providers_sorted = sorted(self.providers, key=lambda p: (-p.weight, p.tier.value))
# ── Credit pressure: if paid providers are near quota, bump free providers up ──
credit_pressure = False
for p in providers_sorted:
if p.monthly_quota > 0 and p.tier.value in ("paid", "freemium"):
used = _quota_usage.get(p.name, 0)
if used > p.monthly_quota * 0.8: # 80% threshold
credit_pressure = True
logger.info(
f"Credit pressure: {p.name} at {used}/{p.monthly_quota} ({used * 100 // p.monthly_quota}%)"
)
if credit_pressure:
# Re-sort: push free/local providers above paid/freemium near quota
providers_sorted.sort(key=lambda p: (0 if p.tier.value in ("local", "free_api") else 1, -p.weight))
for provider in providers_sorted:
# Check circuit breaker
if not _circuit_breakers.get(provider.name, _CircuitBreaker()).can_call():
logger.debug(f"Circuit breaker open for {provider.name}")
continue
# Check rate limit
if not _rate_limiters.get(provider.name, _RateLimiter()).can_call():
logger.debug(f"Rate limit exceeded for {provider.name}")
continue
# Check quota
if provider.monthly_quota > 0:
used = _quota_usage.get(provider.name, 0)
if used >= provider.monthly_quota:
logger.debug(f"Monthly quota exceeded for {provider.name}")
continue
try:
# Get API key from env (vault is pool manager, use os.getenv for direct keys)
api_key = None
if provider.requires_key and provider.key_env:
api_key = os.getenv(provider.key_env, "")
result = await provider.fetch_fn(api_key=api_key, **kwargs)
if result is not None:
_rate_limiters[provider.name].record_call()
if provider.monthly_quota > 0:
_quota_usage[provider.name] = _quota_usage.get(provider.name, 0) + 1
_circuit_breakers[provider.name].record_success()
return result
except Exception as e:
logger.warning(f"Provider {provider.name} failed: {e}")
_circuit_breakers[provider.name].record_failure()
continue
return None
# ── Circuit Breaker ────────────────────────────────────────────
class _CircuitBreaker:
"""Circuit breaker to prevent cascading failures."""
def __init__(self, threshold: int = 5, timeout: float = 60.0):
self.threshold = threshold
self.timeout = timeout
self.failures = 0
self.last_failure = 0.0
self.open = False
def can_call(self) -> bool:
if self.open:
if time.time() - self.last_failure > self.timeout:
self.open = False
self.failures = 0
return True
return False
return True
def record_failure(self) -> None:
self.failures += 1
self.last_failure = time.time()
if self.failures >= self.threshold:
self.open = True
def record_success(self) -> None:
self.failures = 0
self.open = False
class _RateLimiter:
"""Simple rate limiter based on time intervals."""
def __init__(self, rps: float = 1.0):
self.rps = rps
self.min_interval = 1.0 / rps
self.last_call = 0.0
def can_call(self) -> bool:
return time.time() - self.last_call >= self.min_interval
def record_call(self) -> None:
self.last_call = time.time()
# ── Shared State ───────────────────────────────────────────────
_circuit_breakers: dict[str, _CircuitBreaker] = {}
_rate_limiters: dict[str, _RateLimiter] = {}
_quota_usage: dict[str, int] = {}
def reset_state() -> None:
"""Reset all circuit breakers, rate limiters, and quota tracking.
Useful for testing and debugging.
"""
global _circuit_breakers, _rate_limiters, _quota_usage
_circuit_breakers = {}
_rate_limiters = {}
_quota_usage = {}
|