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

api_service_handler/rotation.py

@@ -0,0 +1,163 @@
+"""Key rotation strategies for selecting API keys."""
+
+from __future__ import annotations
+
+import random
+from collections import defaultdict
+from typing import Optional
+
+from .enums import Provider, RotationStrategy, KeyStatus
+from .exceptions import NoAvailableKeyError
+from .models import APIKey
+
+
+class KeyRotator:
+    """Manages key selection using configurable rotation strategies.
+
+    Supports round-robin, least-used, random, and weighted strategies.
+    All strategies automatically skip keys that are inactive, rate-limited,
+    expired, or at max concurrent usage.
+    """
+
+    def __init__(self, strategy: RotationStrategy | str = RotationStrategy.ROUND_ROBIN) -> None:
+        """Initialize the rotator.
+
+        Args:
+            strategy: The rotation strategy to use.
+        """
+        if isinstance(strategy, str):
+            strategy = RotationStrategy(strategy)
+        self._strategy = strategy
+        # Round-robin index per provider
+        self._rr_counters: dict[str, int] = defaultdict(int)
+
+    @property
+    def strategy(self) -> RotationStrategy:
+        """The current rotation strategy."""
+        return self._strategy
+
+    @strategy.setter
+    def strategy(self, value: RotationStrategy | str) -> None:
+        """Set the rotation strategy."""
+        if isinstance(value, str):
+            value = RotationStrategy(value)
+        self._strategy = value
+
+    def select_key(
+        self,
+        keys: list[APIKey],
+        provider: Optional[Provider | str] = None,
+    ) -> APIKey:
+        """Select the next key using the configured strategy.
+
+        Args:
+            keys: List of candidate keys to choose from.
+            provider: Optional provider to filter keys by.
+
+        Returns:
+            The selected APIKey.
+
+        Raises:
+            NoAvailableKeyError: If no suitable key is available.
+        """
+        # Filter to available keys
+        available = self._filter_available(keys, provider)
+
+        if not available:
+            provider_name = (
+                provider.value if isinstance(provider, Provider) else provider or "any"
+            )
+            raise NoAvailableKeyError(provider=provider_name)
+
+        # Sort by priority (lower = higher priority)
+        available.sort(key=lambda k: k.priority)
+
+        # Apply strategy
+        if self._strategy == RotationStrategy.ROUND_ROBIN:
+            return self._round_robin(available, provider)
+        elif self._strategy == RotationStrategy.LEAST_USED:
+            return self._least_used(available)
+        elif self._strategy == RotationStrategy.RANDOM:
+            return self._random(available)
+        elif self._strategy == RotationStrategy.WEIGHTED:
+            return self._weighted(available)
+        else:
+            return self._round_robin(available, provider)
+
+    def _filter_available(
+        self,
+        keys: list[APIKey],
+        provider: Optional[Provider | str] = None,
+    ) -> list[APIKey]:
+        """Filter keys to only those that are available for use.
+
+        A key is available if:
+        - Status is ACTIVE
+        - Not expired
+        - Not rate-limited (daily and monthly counts under limits)
+        - Concurrent usage under max_concurrent (if set)
+        """
+        available = []
+        for key in keys:
+            # Filter by provider if specified
+            if provider is not None:
+                provider_str = (
+                    provider.value if isinstance(provider, Provider) else provider
+                )
+                key_provider_str = (
+                    key.provider.value if isinstance(key.provider, Provider) else key.provider
+                )
+                if key_provider_str != provider_str:
+                    continue
+
+            # Check status
+            key_status = key.status.value if isinstance(key.status, KeyStatus) else key.status
+            if key_status != KeyStatus.ACTIVE.value:
+                continue
+
+            # Check capacity using model property
+            if hasattr(key, "has_capacity") and not key.has_capacity:
+                continue
+
+            available.append(key)
+
+        return available
+
+    def _round_robin(
+        self,
+        keys: list[APIKey],
+        provider: Optional[Provider | str] = None,
+    ) -> APIKey:
+        """Select key using round-robin within the provider group."""
+        provider_key = str(provider) if provider else "_all"
+        idx = self._rr_counters[provider_key] % len(keys)
+        self._rr_counters[provider_key] = idx + 1
+        return keys[idx]
+
+    def _least_used(self, keys: list[APIKey]) -> APIKey:
+        """Select the key with the lowest total usage count."""
+        return min(keys, key=lambda k: k.total_usage_count)
+
+    def _random(self, keys: list[APIKey]) -> APIKey:
+        """Select a random key."""
+        return random.choice(keys)
+
+    def _weighted(self, keys: list[APIKey]) -> APIKey:
+        """Select a key using weighted random selection.
+
+        Keys with higher weight values are more likely to be selected.
+        """
+        weights = [k.weight for k in keys]
+        return random.choices(keys, weights=weights, k=1)[0]
+
+    def reset_counters(self, provider: Optional[str] = None) -> None:
+        """Reset round-robin counters.
+
+        Args:
+            provider: If set, only reset the counter for this provider.
+                     If None, reset all counters.
+        """
+        if provider:
+            self._rr_counters.pop(provider, None)
+        else:
+            self._rr_counters.clear()

api_service_handler/storage/__init__.py

@@ -0,0 +1,7 @@
+"""Storage backends for the API Service Handler."""
+
+from __future__ import annotations
+
+from .base import StorageBackend
+
+__all__ = ["StorageBackend"]

api_service_handler/storage/base.py

@@ -0,0 +1,243 @@
+"""Abstract base class for storage backends."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from datetime import date
+from typing import Optional
+
+from ..enums import Provider, KeyStatus
+from ..models import APIKey, KeyFilter, KeyUpdateRequest, BulkOperationResult
+
+
+class StorageBackend(ABC):
+    """Abstract base class defining the storage interface.
+
+    All storage backends (Memory, SQLite, MongoDB, PostgreSQL) must implement
+    this interface. Methods are async to support async database drivers.
+    """
+
+    @abstractmethod
+    async def initialize(self) -> None:
+        """Initialize the storage backend (create tables/collections).
+
+        Must be called before any other operations.
+        """
+        ...
+
+    @abstractmethod
+    async def close(self) -> None:
+        """Close connections and clean up resources."""
+        ...
+
+    # ── CRUD ───────────────────────────────────────────────────────────────
+
+    @abstractmethod
+    async def add_key(self, key: APIKey) -> APIKey:
+        """Insert a new API key record.
+
+        Args:
+            key: The APIKey model to insert.
+
+        Returns:
+            The inserted APIKey (with generated ID if not set).
+
+        Raises:
+            DuplicateKeyError: If key_value + provider already exists.
+        """
+        ...
+
+    @abstractmethod
+    async def get_key(self, key_id: str) -> APIKey:
+        """Retrieve a single key by its ID.
+
+        Args:
+            key_id: The unique key identifier.
+
+        Returns:
+            The matching APIKey.
+
+        Raises:
+            KeyNotFoundError: If no key exists with the given ID.
+        """
+        ...
+
+    @abstractmethod
+    async def get_keys_by_provider(self, provider: Provider) -> list[APIKey]:
+        """Retrieve all keys for a specific provider.
+
+        Args:
+            provider: The API provider to filter by.
+
+        Returns:
+            List of matching APIKey records (may be empty).
+        """
+        ...
+
+    @abstractmethod
+    async def get_all_keys(self, key_filter: Optional[KeyFilter] = None) -> list[APIKey]:
+        """Retrieve all keys, optionally filtered.
+
+        Args:
+            key_filter: Optional filter criteria.
+
+        Returns:
+            List of matching APIKey records.
+        """
+        ...
+
+    @abstractmethod
+    async def update_key(self, key_id: str, updates: KeyUpdateRequest) -> APIKey:
+        """Partially update an existing key.
+
+        Args:
+            key_id: The key to update.
+            updates: Fields to update (only set fields are applied).
+
+        Returns:
+            The updated APIKey.
+
+        Raises:
+            KeyNotFoundError: If no key exists with the given ID.
+        """
+        ...
+
+    @abstractmethod
+    async def delete_key(self, key_id: str, soft: bool = True) -> bool:
+        """Delete a key.
+
+        Args:
+            key_id: The key to delete.
+            soft: If True, mark as REVOKED instead of hard deleting.
+
+        Returns:
+            True if the key was found and deleted/revoked.
+
+        Raises:
+            KeyNotFoundError: If no key exists with the given ID.
+        """
+        ...
+
+    # ── Usage Tracking ─────────────────────────────────────────────────────
+
+    @abstractmethod
+    async def increment_usage(
+        self,
+        key_id: str,
+        daily: int = 1,
+        monthly: int = 1,
+        total: int = 1,
+    ) -> None:
+        """Atomically increment usage counters for a key.
+
+        Args:
+            key_id: The key whose counters to increment.
+            daily: Amount to add to daily_usage_count.
+            monthly: Amount to add to monthly_usage_count.
+            total: Amount to add to total_usage_count.
+        """
+        ...
+
+    @abstractmethod
+    async def update_concurrent_usage(self, key_id: str, delta: int) -> int:
+        """Atomically adjust the concurrent usage counter.
+
+        Args:
+            key_id: The key to update.
+            delta: The change (+1 for acquire, -1 for release).
+
+        Returns:
+            The new concurrent_usage value after the update.
+
+        Raises:
+            KeyNotFoundError: If no key exists with the given ID.
+        """
+        ...
+
+    @abstractmethod
+    async def reset_daily_counts(self, before_date: date) -> int:
+        """Reset daily_usage_count for all keys with last_reset_daily before the given date.
+
+        Args:
+            before_date: Reset keys whose last_reset_daily is before this date.
+
+        Returns:
+            Number of keys reset.
+        """
+        ...
+
+    @abstractmethod
+    async def reset_monthly_counts(self, before_date: date) -> int:
+        """Reset monthly_usage_count for all keys with last_reset_monthly before the given date.
+
+        Args:
+            before_date: Reset keys whose last_reset_monthly is before this date.
+
+        Returns:
+            Number of keys reset.
+        """
+        ...
+
+    @abstractmethod
+    async def update_last_used(self, key_id: str) -> None:
+        """Update the last_used_at timestamp for a key.
+
+        Args:
+            key_id: The key to update.
+        """
+        ...
+
+    # ── Bulk Operations ────────────────────────────────────────────────────
+
+    @abstractmethod
+    async def bulk_add_keys(self, keys: list[APIKey]) -> BulkOperationResult:
+        """Insert multiple keys at once.
+
+        Args:
+            keys: List of APIKey models to insert.
+
+        Returns:
+            BulkOperationResult with success/failure counts.
+        """
+        ...
+
+    @abstractmethod
+    async def bulk_delete_keys(self, key_ids: list[str], soft: bool = True) -> BulkOperationResult:
+        """Delete multiple keys at once.
+
+        Args:
+            key_ids: List of key IDs to delete.
+            soft: If True, mark as REVOKED instead of hard deleting.
+
+        Returns:
+            BulkOperationResult with success/failure counts.
+        """
+        ...
+
+    # ── Health ─────────────────────────────────────────────────────────────
+
+    @abstractmethod
+    async def health_check(self) -> bool:
+        """Check if the storage backend is healthy and reachable.
+
+        Returns:
+            True if the backend is operational.
+        """
+        ...
+
+    @abstractmethod
+    async def count_keys(
+        self,
+        provider: Optional[Provider] = None,
+        status: Optional[KeyStatus] = None,
+    ) -> int:
+        """Count keys matching the given criteria.
+
+        Args:
+            provider: Optional provider filter.
+            status: Optional status filter.
+
+        Returns:
+            Number of matching keys.
+        """
+        ...

api_service_handler/storage/memory.py

@@ -0,0 +1,229 @@
+"""In-memory storage backend for testing and lightweight use."""
+
+from __future__ import annotations
+
+import copy
+from datetime import date, datetime, timezone
+from typing import Optional
+
+from ..enums import Provider, KeyStatus
+from ..exceptions import KeyNotFoundError, DuplicateKeyError
+from ..models import APIKey, KeyFilter, KeyUpdateRequest, BulkOperationResult
+from .base import StorageBackend
+
+
+class MemoryStorageBackend(StorageBackend):
+    """In-memory dict-based storage backend.
+
+    Suitable for testing, prototyping, and lightweight single-process use.
+    Data is lost when the process exits.
+    """
+
+    def __init__(self) -> None:
+        self._keys: dict[str, APIKey] = {}
+        self._initialized: bool = False
+
+    async def initialize(self) -> None:
+        """Initialize the in-memory store."""
+        self._keys = {}
+        self._initialized = True
+
+    async def close(self) -> None:
+        """Clear the in-memory store."""
+        self._keys.clear()
+        self._initialized = False
+
+    # ── CRUD ───────────────────────────────────────────────────────────────
+
+    async def add_key(self, key: APIKey) -> APIKey:
+        """Insert a new API key into memory."""
+        # Check for duplicates (same key_value + provider)
+        for existing in self._keys.values():
+            if (
+                existing.key_value == key.key_value
+                and existing.provider == key.provider
+                and existing.status != KeyStatus.REVOKED
+            ):
+                raise DuplicateKeyError(provider=str(key.provider), key_value=key.key_value)
+
+        self._keys[key.id] = copy.deepcopy(key)
+        return copy.deepcopy(key)
+
+    async def get_key(self, key_id: str) -> APIKey:
+        """Retrieve a key by ID."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+        return copy.deepcopy(self._keys[key_id])
+
+    async def get_keys_by_provider(self, provider: Provider) -> list[APIKey]:
+        """Get all keys for a provider."""
+        provider_str = provider.value if isinstance(provider, Provider) else provider
+        return [
+            copy.deepcopy(k)
+            for k in self._keys.values()
+            if (k.provider.value if isinstance(k.provider, Provider) else k.provider) == provider_str
+        ]
+
+    async def get_all_keys(self, key_filter: Optional[KeyFilter] = None) -> list[APIKey]:
+        """Get all keys, optionally filtered."""
+        keys = list(self._keys.values())
+        if key_filter:
+            keys = [k for k in keys if key_filter.matches(k)]
+        return [copy.deepcopy(k) for k in keys]
+
+    async def update_key(self, key_id: str, updates: KeyUpdateRequest) -> APIKey:
+        """Update a key's fields."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+
+        key = self._keys[key_id]
+        update_data = updates.model_dump(exclude_unset=True)
+
+        for field_name, value in update_data.items():
+            setattr(key, field_name, value)
+
+        key.updated_at = datetime.now(timezone.utc)
+        self._keys[key_id] = key
+        return copy.deepcopy(key)
+
+    async def delete_key(self, key_id: str, soft: bool = True) -> bool:
+        """Delete or revoke a key."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+
+        if soft:
+            self._keys[key_id].status = KeyStatus.REVOKED
+            self._keys[key_id].updated_at = datetime.now(timezone.utc)
+        else:
+            del self._keys[key_id]
+
+        return True
+
+    # ── Usage Tracking ─────────────────────────────────────────────────────
+
+    async def increment_usage(
+        self,
+        key_id: str,
+        daily: int = 1,
+        monthly: int = 1,
+        total: int = 1,
+    ) -> None:
+        """Increment usage counters."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+
+        key = self._keys[key_id]
+        key.daily_usage_count += daily
+        key.monthly_usage_count += monthly
+        key.total_usage_count += total
+        key.last_used_at = datetime.now(timezone.utc)
+        key.updated_at = datetime.now(timezone.utc)
+
+    async def update_concurrent_usage(self, key_id: str, delta: int) -> int:
+        """Adjust concurrent usage counter."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+
+        key = self._keys[key_id]
+        key.concurrent_usage = max(0, key.concurrent_usage + delta)
+        key.updated_at = datetime.now(timezone.utc)
+        return key.concurrent_usage
+
+    async def reset_daily_counts(self, before_date: date) -> int:
+        """Reset daily counts for keys not yet reset today."""
+        count = 0
+        for key in self._keys.values():
+            if key.last_reset_daily < before_date:
+                key.daily_usage_count = 0
+                key.last_reset_daily = before_date
+                if key.status == KeyStatus.RATE_LIMITED:
+                    # Check if monthly limit is also hit
+                    if key.monthly_limit is None or key.monthly_usage_count < key.monthly_limit:
+                        key.status = KeyStatus.ACTIVE
+                key.updated_at = datetime.now(timezone.utc)
+                count += 1
+        return count
+
+    async def reset_monthly_counts(self, before_date: date) -> int:
+        """Reset monthly counts for keys not yet reset this month."""
+        count = 0
+        for key in self._keys.values():
+            if (
+                key.last_reset_monthly.month != before_date.month
+                or key.last_reset_monthly.year != before_date.year
+            ):
+                key.monthly_usage_count = 0
+                key.last_reset_monthly = before_date
+                if key.status == KeyStatus.RATE_LIMITED:
+                    key.status = KeyStatus.ACTIVE
+                key.updated_at = datetime.now(timezone.utc)
+                count += 1
+        return count
+
+    async def update_last_used(self, key_id: str) -> None:
+        """Update last_used_at timestamp."""
+        if key_id not in self._keys:
+            raise KeyNotFoundError(key_id=key_id)
+        self._keys[key_id].last_used_at = datetime.now(timezone.utc)
+
+    # ── Bulk Operations ────────────────────────────────────────────────────
+
+    async def bulk_add_keys(self, keys: list[APIKey]) -> BulkOperationResult:
+        """Insert multiple keys."""
+        result = BulkOperationResult(total=len(keys))
+
+        for key in keys:
+            try:
+                added = await self.add_key(key)
+                result.successful += 1
+                result.created_ids.append(added.id)
+            except Exception as e:
+                result.failed += 1
+                result.errors.append(f"Key {key.alias or key.id}: {e}")
+
+        return result
+
+    async def bulk_delete_keys(self, key_ids: list[str], soft: bool = True) -> BulkOperationResult:
+        """Delete multiple keys."""
+        result = BulkOperationResult(total=len(key_ids))
+
+        for key_id in key_ids:
+            try:
+                await self.delete_key(key_id, soft=soft)
+                result.successful += 1
+            except Exception as e:
+                result.failed += 1
+                result.errors.append(f"Key {key_id}: {e}")
+
+        return result
+
+    # ── Health ─────────────────────────────────────────────────────────────
+
+    async def health_check(self) -> bool:
+        """Memory backend is always healthy if initialized."""
+        return self._initialized
+
+    async def count_keys(
+        self,
+        provider: Optional[Provider] = None,
+        status: Optional[KeyStatus] = None,
+    ) -> int:
+        """Count keys matching criteria."""
+        count = 0
+        for key in self._keys.values():
+            if provider is not None:
+                provider_str = provider.value if isinstance(provider, Provider) else provider
+                key_provider_str = (
+                    key.provider.value if isinstance(key.provider, Provider) else key.provider
+                )
+                if key_provider_str != provider_str:
+                    continue
+            if status is not None:
+                status_str = status.value if isinstance(status, KeyStatus) else status
+                key_status_str = (
+                    key.status.value if isinstance(key.status, KeyStatus) else key.status
+                )
+                if key_status_str != status_str:
+                    continue
+            count += 1
+        return count