61 lines
1.9 KiB
Python
61 lines
1.9 KiB
Python
"""Provider key selector with round-robin strategy."""
|
|
import asyncio
|
|
from typing import Optional, Dict
|
|
from ..database import db_connection
|
|
from ..models import ProviderKey
|
|
|
|
|
|
class KeySelector:
|
|
"""Round-robin key selector for provider API keys."""
|
|
|
|
def __init__(self):
|
|
self._indices: Dict[int, int] = {} # provider_id -> current index
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def get_next_key(self, provider_id: int) -> Optional[ProviderKey]:
|
|
"""Get the next available API key for a provider.
|
|
|
|
Args:
|
|
provider_id: The provider ID
|
|
|
|
Returns:
|
|
ProviderKey if available, None if no keys configured
|
|
"""
|
|
async with self._lock:
|
|
async with db_connection() as db:
|
|
cursor = await db.execute(
|
|
"""SELECT id, provider_id, key, is_active, created_at
|
|
FROM provider_keys
|
|
WHERE provider_id = ? AND is_active = TRUE""",
|
|
(provider_id,)
|
|
)
|
|
keys = await cursor.fetchall()
|
|
|
|
if not keys:
|
|
return None
|
|
|
|
# Get or initialize index
|
|
if provider_id not in self._indices:
|
|
self._indices[provider_id] = 0
|
|
|
|
# Round-robin selection
|
|
index = self._indices[provider_id] % len(keys)
|
|
self._indices[provider_id] = index + 1
|
|
|
|
row = keys[index]
|
|
return ProviderKey(
|
|
id=row["id"],
|
|
provider_id=row["provider_id"],
|
|
key=row["key"],
|
|
is_active=row["is_active"],
|
|
created_at=row["created_at"]
|
|
)
|
|
|
|
def reset_index(self, provider_id: int):
|
|
"""Reset the round-robin index for a provider."""
|
|
self._indices.pop(provider_id, None)
|
|
|
|
|
|
# Global instance
|
|
key_selector = KeySelector()
|