api-service-handler 0.1.6__py3-none-any.whl

api_service_handler/exceptions.py

@@ -0,0 +1,184 @@
+"""Custom exceptions for the api-service-handler library.
+
+Every exception inherits from :class:`APIServiceHandlerError` so callers
+can catch the entire family with a single ``except`` clause when desired.
+Each concrete exception stores structured context as instance attributes
+and generates a clear, actionable error message.
+"""
+
+from __future__ import annotations
+
+
+class APIServiceHandlerError(Exception):
+    """Base exception for all api-service-handler errors.
+
+    All library-specific exceptions inherit from this class, making it
+    easy to catch any error raised by the library::
+
+        try:
+            key = await manager.get_key("openai")
+        except APIServiceHandlerError as exc:
+            logging.error("Service handler failure: %s", exc)
+    """
+
+    def __init__(self, message: str = "An API service handler error occurred") -> None:
+        self.message = message
+        super().__init__(self.message)
+
+
+class KeyNotFoundError(APIServiceHandlerError):
+    """Raised when a key with the given ID does not exist in storage.
+
+    Attributes:
+        key_id: The identifier that was looked up.
+    """
+
+    def __init__(self, key_id: str) -> None:
+        self.key_id = key_id
+        super().__init__(f"API key not found: '{key_id}'")
+
+
+class DuplicateKeyError(APIServiceHandlerError):
+    """Raised when attempting to register a key that already exists.
+
+    A key is considered duplicate when the same ``key_value`` is already
+    registered for the given ``provider``.
+
+    Attributes:
+        provider: The provider for which the duplicate was detected.
+    """
+
+    def __init__(self, provider: str, key_value: str | None = None) -> None:
+        self.provider = provider
+        self.key_value = key_value
+        msg = f"A key with the same value already exists for provider '{provider}'"
+        if key_value:
+            masked = key_value[:8] + "***" if len(key_value) > 8 else key_value[:2] + "***"
+            msg = f"Key '{masked}' already exists for provider '{provider}'"
+        super().__init__(msg)
+
+
+class RateLimitExceededError(APIServiceHandlerError):
+    """Raised when a key has exceeded its daily or monthly usage limit.
+
+    Attributes:
+        key_id: The key that hit the limit.
+        limit_type: Either ``'daily'`` or ``'monthly'``.
+    """
+
+    def __init__(
+        self,
+        key_id: str,
+        limit_type: str,
+        limit: int | None = None,
+        current: int | None = None,
+    ) -> None:
+        self.key_id = key_id
+        self.limit_type = limit_type
+        self.limit = limit
+        self.current = current
+        msg = f"API key '{key_id}' has exceeded its {limit_type} rate limit"
+        if limit is not None:
+            msg += f" (limit: {limit}, current: {current})"
+        super().__init__(msg)
+
+
+class MaxConcurrentExceededError(APIServiceHandlerError):
+    """Raised when a key's concurrent usage equals or exceeds its maximum.
+
+    Attributes:
+        key_id: The key that is at capacity.
+        max_concurrent: The configured concurrency ceiling.
+    """
+
+    def __init__(self, key_id: str, max_concurrent: int) -> None:
+        self.key_id = key_id
+        self.max_concurrent = max_concurrent
+        super().__init__(
+            f"API key '{key_id}' has reached its maximum concurrent usage "
+            f"of {max_concurrent}"
+        )
+
+
+class NoAvailableKeyError(APIServiceHandlerError):
+    """Raised when no active key is available for the requested provider.
+
+    Attributes:
+        provider: The provider for which no key could be found.
+    """
+
+    def __init__(self, provider: str) -> None:
+        self.provider = provider
+        super().__init__(
+            f"No available API key for provider '{provider}'"
+        )
+
+
+class StorageConnectionError(APIServiceHandlerError):
+    """Raised when the storage backend cannot be reached.
+
+    Attributes:
+        backend: Name of the storage backend (e.g. ``'postgresql'``).
+    """
+
+    def __init__(self, backend: str, detail: str | None = None, details: str | None = None) -> None:
+        self.backend = backend
+        self.detail = detail or details
+        msg = f"Failed to connect to storage backend '{backend}'"
+        if self.detail:
+            msg += f": {self.detail}"
+        super().__init__(msg)
+
+
+class EncryptionError(APIServiceHandlerError):
+    """Raised when an encryption or decryption operation fails.
+
+    This may indicate a missing or invalid encryption key, corrupted
+    ciphertext, or an unsupported algorithm.
+    """
+
+    def __init__(self, detail: str = "Encryption/decryption operation failed") -> None:
+        self.detail = detail
+        super().__init__(detail)
+
+
+class InvalidProviderError(APIServiceHandlerError):
+    """Raised when an unrecognised provider string is used in a strict context.
+
+    Attributes:
+        provider: The invalid provider string that was supplied.
+    """
+
+    def __init__(self, provider: str) -> None:
+        self.provider = provider
+        super().__init__(
+            f"Invalid or unknown provider: '{provider}'"
+        )
+
+
+class KeyExpiredError(APIServiceHandlerError):
+    """Raised when an API key is past its ``expires_at`` timestamp.
+
+    Attributes:
+        key_id: The expired key's identifier.
+    """
+
+    def __init__(self, key_id: str) -> None:
+        self.key_id = key_id
+        super().__init__(
+            f"API key '{key_id}' has expired"
+        )
+
+
+class StorageNotInitializedError(APIServiceHandlerError):
+    """Raised when a storage operation is attempted before initialization.
+
+    Callers must invoke ``storage.initialize()`` (or its async counterpart)
+    before performing any read/write operations.
+    """
+
+    def __init__(self) -> None:
+        super().__init__(
+            "Storage backend has not been initialized. "
+            "Call 'storage.initialize()' before performing operations."
+        )

api_service_handler/models.py

@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+"""Data models for API key management.
+
+Provides Pydantic v2 models for API key storage, creation, update,
+filtering, usage statistics, and bulk operation results.
+"""
+
+from datetime import date, datetime, timezone
+from typing import Any, Optional
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .enums import Environment, KeyStatus, Provider, RotationStrategy
+
+
+# ---------------------------------------------------------------------------
+# Main model
+# ---------------------------------------------------------------------------
+
+
+class APIKey(BaseModel):
+    """Represents a single managed API key with usage tracking and metadata.
+
+    Attributes:
+        id: Unique identifier (UUID) for this key entry.
+        provider: The service provider this key belongs to.
+        key_value: The actual API key string (may be encrypted at rest).
+        alias: Optional human-friendly name, e.g. ``'prod-openai-1'``.
+        status: Current lifecycle status of the key.
+        environment: Deployment environment the key is designated for.
+        daily_limit: Maximum allowed requests per day (``None`` = unlimited).
+        monthly_limit: Maximum allowed requests per month (``None`` = unlimited).
+        daily_usage_count: Number of requests made today.
+        monthly_usage_count: Number of requests made this month.
+        total_usage_count: Lifetime request count.
+        concurrent_usage: Number of in-flight requests right now.
+        max_concurrent: Maximum simultaneous in-flight requests (``None`` = unlimited).
+        weight: Relative weight for weighted-rotation strategies.
+        priority: Priority rank — lower values are selected first.
+        metadata: Arbitrary key/value pairs for user-defined data.
+        tags: Free-form labels for categorisation and filtering.
+        last_used_at: Timestamp of the most recent usage.
+        last_reset_daily: Date when daily counters were last zeroed.
+        last_reset_monthly: Date when monthly counters were last zeroed.
+        expires_at: Optional expiration timestamp.
+        created_at: Timestamp when this record was created.
+        updated_at: Timestamp when this record was last modified.
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    provider: Provider
+    key_value: str
+    alias: Optional[str] = None
+    status: KeyStatus = KeyStatus.ACTIVE
+    environment: Environment = Environment.PRODUCTION
+
+    # Rate limits
+    daily_limit: Optional[int] = None
+    monthly_limit: Optional[int] = None
+
+    # Usage counters
+    daily_usage_count: int = 0
+    monthly_usage_count: int = 0
+    total_usage_count: int = 0
+
+    # Concurrent usage
+    concurrent_usage: int = 0
+    max_concurrent: Optional[int] = None
+
+    # Rotation
+    weight: int = 1
+    priority: int = 0
+
+    # Metadata
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    tags: list[str] = Field(default_factory=list)
+
+    # Timestamps
+    last_used_at: Optional[datetime] = None
+    last_reset_daily: date = Field(default_factory=date.today)
+    last_reset_monthly: date = Field(default_factory=date.today)
+    expires_at: Optional[datetime] = None
+    created_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    updated_at: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+
+    # -- Computed properties ------------------------------------------------
+
+    @property
+    def is_expired(self) -> bool:
+        """Return ``True`` if the key has passed its expiration time."""
+        if self.expires_at is None:
+            return False
+        now = datetime.now(timezone.utc)
+        # Handle naive datetimes stored without tz info
+        expires = self.expires_at
+        if expires.tzinfo is None:
+            expires = expires.replace(tzinfo=timezone.utc)
+        return now >= expires
+
+    @property
+    def is_rate_limited(self) -> bool:
+        """Return ``True`` if any configured rate limit has been reached."""
+        if self.daily_limit is not None and self.daily_usage_count >= self.daily_limit:
+            return True
+        if self.monthly_limit is not None and self.monthly_usage_count >= self.monthly_limit:
+            return True
+        return False
+
+    @property
+    def has_capacity(self) -> bool:
+        """Return ``True`` if the key can accept another request right now.
+
+        A key has capacity when it is **active**, **not expired**,
+        **not rate-limited**, and below its concurrent-usage ceiling.
+        """
+        if self.is_expired:
+            return False
+        if self.is_rate_limited:
+            return False
+        if self.status != KeyStatus.ACTIVE:
+            return False
+        if self.max_concurrent is not None and self.concurrent_usage >= self.max_concurrent:
+            return False
+        return True
+
+    @property
+    def needs_daily_reset(self) -> bool:
+        """Return ``True`` if daily counters are stale (belong to a previous day)."""
+        return self.last_reset_daily < date.today()
+
+    @property
+    def needs_monthly_reset(self) -> bool:
+        """Return ``True`` if monthly counters are stale (belong to a previous month)."""
+        today = date.today()
+        return (
+            self.last_reset_monthly.year < today.year
+            or self.last_reset_monthly.month < today.month
+        )
+
+
+# ---------------------------------------------------------------------------
+# Request / response helpers
+# ---------------------------------------------------------------------------
+
+
+class KeyCreateRequest(BaseModel):
+    """Input schema for creating a new API key.
+
+    Only ``provider`` and ``key_value`` are mandatory; all other fields
+    fall back to sensible defaults on the resulting :class:`APIKey`.
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    provider: Provider
+    key_value: str
+    alias: Optional[str] = None
+    daily_limit: Optional[int] = None
+    monthly_limit: Optional[int] = None
+    max_concurrent: Optional[int] = None
+    environment: Environment = Environment.PRODUCTION
+    metadata: dict[str, Any] = Field(default_factory=dict)
+    tags: list[str] = Field(default_factory=list)
+    expires_at: Optional[datetime] = None
+    weight: int = 1
+    priority: int = 0
+
+
+class KeyUpdateRequest(BaseModel):
+    """Input schema for partially updating an existing API key.
+
+    Every field is optional — only the fields that are set (not ``None``)
+    should be applied to the target :class:`APIKey`.
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    alias: Optional[str] = None
+    status: Optional[KeyStatus] = None
+    daily_limit: Optional[int] = None
+    monthly_limit: Optional[int] = None
+    max_concurrent: Optional[int] = None
+    environment: Optional[Environment] = None
+    metadata: Optional[dict[str, Any]] = None
+    tags: Optional[list[str]] = None
+    expires_at: Optional[datetime] = None
+    weight: Optional[int] = None
+    priority: Optional[int] = None
+
+
+# ---------------------------------------------------------------------------
+# Filtering
+# ---------------------------------------------------------------------------
+
+
+class KeyFilter(BaseModel):
+    """Declarative filter for querying stored API keys.
+
+    Only non-``None`` fields participate in matching.  Use the
+    :meth:`matches` helper to test a candidate :class:`APIKey`.
+    """
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    provider: Optional[Provider] = None
+    status: Optional[KeyStatus] = None
+    environment: Optional[Environment] = None
+    tags: Optional[list[str]] = None
+    metadata_filter: Optional[dict[str, Any]] = None
+    has_capacity: Optional[bool] = None
+    alias_contains: Optional[str] = None
+
+    def matches(self, key: APIKey) -> bool:
+        """Return ``True`` if *key* satisfies every set filter criterion.
+
+        Args:
+            key: The API key to evaluate.
+
+        Returns:
+            ``True`` when all non-``None`` filter fields match the key.
+        """
+        if self.provider is not None and key.provider != self.provider:
+            return False
+        if self.status is not None and key.status != self.status:
+            return False
+        if self.environment is not None and key.environment != self.environment:
+            return False
+
+        # Tag filter: at least one of the requested tags must be present
+        if self.tags is not None:
+            if not any(tag in key.tags for tag in self.tags):
+                return False
+
+        # Metadata filter: all key/value pairs must match
+        if self.metadata_filter is not None:
+            for meta_key, meta_value in self.metadata_filter.items():
+                if key.metadata.get(meta_key) != meta_value:
+                    return False
+
+        # Capacity filter
+        if self.has_capacity is not None:
+            if key.has_capacity != self.has_capacity:
+                return False
+
+        # Alias substring search (case-insensitive)
+        if self.alias_contains is not None:
+            if key.alias is None:
+                return False
+            if self.alias_contains.lower() not in key.alias.lower():
+                return False
+
+        return True
+
+
+# ---------------------------------------------------------------------------
+# Usage statistics
+# ---------------------------------------------------------------------------
+
+
+class UsageStats(BaseModel):
+    """Aggregated usage statistics snapshot for a single API key."""
+
+    model_config = ConfigDict(use_enum_values=True)
+
+    key_id: str
+    provider: Provider
+    alias: Optional[str] = None
+
+    daily_usage_count: int = 0
+    monthly_usage_count: int = 0
+    total_usage_count: int = 0
+    concurrent_usage: int = 0
+
+    daily_limit: Optional[int] = None
+    monthly_limit: Optional[int] = None
+    max_concurrent: Optional[int] = None
+
+    daily_remaining: Optional[int] = None
+    monthly_remaining: Optional[int] = None
+
+    last_used_at: Optional[datetime] = None
+    status: KeyStatus = KeyStatus.ACTIVE
+
+
+# ---------------------------------------------------------------------------
+# Bulk operations
+# ---------------------------------------------------------------------------
+
+
+class BulkOperationResult(BaseModel):
+    """Summary returned after a bulk create / update / delete operation."""
+
+    total: int = 0
+    successful: int = 0
+    failed: int = 0
+    errors: list[str] = Field(default_factory=list)
+    created_ids: list[str] = Field(default_factory=list)

api_service_handler/rate_limiter.py

@@ -0,0 +1,187 @@
+"""Rate limiting enforcement for API keys."""
+
+from __future__ import annotations
+
+from datetime import date, datetime, timezone
+from typing import Optional
+
+from .enums import KeyStatus
+from .exceptions import RateLimitExceededError
+from .models import APIKey
+from .storage.base import StorageBackend
+
+
+class RateLimiter:
+    """Enforces daily and monthly rate limits on API keys.
+
+    Handles:
+    - Checking if a key has capacity before use
+    - Auto-resetting daily/monthly counters when the period rolls over
+    - Marking keys as RATE_LIMITED when limits are hit
+    - Recovering keys when counters reset
+    """
+
+    def __init__(self, storage: StorageBackend, auto_reset: bool = True) -> None:
+        """Initialize the rate limiter.
+
+        Args:
+            storage: The storage backend for persisting counter changes.
+            auto_reset: If True, auto-reset expired counters on access.
+        """
+        self._storage = storage
+        self._auto_reset = auto_reset
+
+    async def check_and_update(self, key: APIKey) -> APIKey:
+        """Check if a key needs counter resets and apply them.
+
+        This is called before using a key to ensure counters are current.
+
+        Args:
+            key: The API key to check.
+
+        Returns:
+            The updated key (with reset counters if applicable).
+        """
+        if not self._auto_reset:
+            return key
+
+        today = date.today()
+        needs_refresh = False
+
+        # Check daily reset
+        if key.needs_daily_reset:
+            await self._storage.reset_daily_counts(today)
+            needs_refresh = True
+
+        # Check monthly reset
+        if key.needs_monthly_reset:
+            await self._storage.reset_monthly_counts(today)
+            needs_refresh = True
+
+        if needs_refresh:
+            key = await self._storage.get_key(key.id)
+
+        return key
+
+    async def check_capacity(self, key: APIKey) -> bool:
+        """Check if a key has remaining capacity.
+
+        Args:
+            key: The key to check.
+
+        Returns:
+            True if the key can accept more requests.
+        """
+        key = await self.check_and_update(key)
+
+        # Check daily limit
+        if key.daily_limit is not None and key.daily_usage_count >= key.daily_limit:
+            return False
+
+        # Check monthly limit
+        if key.monthly_limit is not None and key.monthly_usage_count >= key.monthly_limit:
+            return False
+
+        # Check concurrent limit
+        if key.max_concurrent is not None and key.concurrent_usage >= key.max_concurrent:
+            return False
+
+        return True
+
+    async def enforce_limits(self, key: APIKey) -> APIKey:
+        """Enforce rate limits, raising an exception if limits are exceeded.
+
+        Call this before each API request to ensure the key is within limits.
+
+        Args:
+            key: The key to check.
+
+        Returns:
+            The validated key.
+
+        Raises:
+            RateLimitExceededError: If daily or monthly limit is exceeded.
+        """
+        key = await self.check_and_update(key)
+
+        # Check daily limit
+        if key.daily_limit is not None and key.daily_usage_count >= key.daily_limit:
+            # Mark as rate limited
+            from .models import KeyUpdateRequest
+            await self._storage.update_key(
+                key.id,
+                KeyUpdateRequest(status=KeyStatus.RATE_LIMITED),
+            )
+            raise RateLimitExceededError(
+                key_id=key.id,
+                limit_type="daily",
+                limit=key.daily_limit,
+                current=key.daily_usage_count,
+            )
+
+        # Check monthly limit
+        if key.monthly_limit is not None and key.monthly_usage_count >= key.monthly_limit:
+            from .models import KeyUpdateRequest
+            await self._storage.update_key(
+                key.id,
+                KeyUpdateRequest(status=KeyStatus.RATE_LIMITED),
+            )
+            raise RateLimitExceededError(
+                key_id=key.id,
+                limit_type="monthly",
+                limit=key.monthly_limit,
+                current=key.monthly_usage_count,
+            )
+
+        return key
+
+    async def record_usage(self, key_id: str) -> None:
+        """Record a single usage event for a key.
+
+        Increments daily, monthly, and total counters by 1.
+
+        Args:
+            key_id: The key to record usage for.
+        """
+        await self._storage.increment_usage(key_id, daily=1, monthly=1, total=1)
+
+    async def get_remaining(self, key: APIKey) -> dict[str, Optional[int]]:
+        """Get remaining capacity for a key.
+
+        Args:
+            key: The key to check.
+
+        Returns:
+            Dict with 'daily_remaining', 'monthly_remaining' values.
+            None means unlimited.
+        """
+        key = await self.check_and_update(key)
+
+        daily_remaining = None
+        if key.daily_limit is not None:
+            daily_remaining = max(0, key.daily_limit - key.daily_usage_count)
+
+        monthly_remaining = None
+        if key.monthly_limit is not None:
+            monthly_remaining = max(0, key.monthly_limit - key.monthly_usage_count)
+
+        return {
+            "daily_remaining": daily_remaining,
+            "monthly_remaining": monthly_remaining,
+        }
+
+    async def reset_daily(self) -> int:
+        """Manually trigger a daily reset for all keys.
+
+        Returns:
+            Number of keys reset.
+        """
+        return await self._storage.reset_daily_counts(date.today())
+
+    async def reset_monthly(self) -> int:
+        """Manually trigger a monthly reset for all keys.
+
+        Returns:
+            Number of keys reset.
+        """
+        return await self._storage.reset_monthly_counts(date.today())