from sqlalchemy.orm import Session from .database import Base, KeyCategory, APIKey # Import models directly from database module # from . import schemas # Schemas will be created later # --- CRUD for KeyCategory --- def create_key_category(db: Session, name: str, type: str, tags: list = None, metadata: dict = None): db_category = KeyCategory(name=name, type=type, tags=tags, metadata_=metadata) db.add(db_category) db.commit() db.refresh(db_category) return db_category def get_key_category(db: Session, category_id: int): return db.query(KeyCategory).filter(KeyCategory.id == category_id).first() def get_key_category_by_name(db: Session, name: str): return db.query(KeyCategory).filter(KeyCategory.name == name).first() def get_key_categories(db: Session, skip: int = 0, limit: int = 100): return db.query(KeyCategory).offset(skip).limit(limit).all() def update_key_category(db: Session, category_id: int, name: str = None, type: str = None, tags: list = None, metadata: dict = None): db_category = get_key_category(db, category_id) if db_category: if name is not None: db_category.name = name if type is not None: db_category.type = type if tags is not None: db_category.tags = tags if metadata is not None: db_category.metadata_ = metadata db.commit() db.refresh(db_category) return db_category def delete_key_category(db: Session, category_id: int): db_category = get_key_category(db, category_id) if db_category: db.delete(db_category) db.commit() return db_category # --- CRUD for APIKey --- def create_api_key(db: Session, value: str, category_id: int, status: str = "active", metadata: dict = None): db_key = APIKey(value=value, category_id=category_id, status=status, metadata_=metadata) db.add(db_key) db.commit() db.refresh(db_key) return db_key def get_api_key(db: Session, key_id: int): return db.query(APIKey).filter(APIKey.id == key_id).first() def get_api_keys(db: Session, category_id: int = None, status: str = None, skip: int = 0, limit: int = 100): query = db.query(APIKey) if category_id is not None: query = query.filter(APIKey.category_id == category_id) if status is not None: query = query.filter(APIKey.status == status) return query.offset(skip).limit(limit).all() def update_api_key(db: Session, key_id: int, value: str = None, category_id: int = None, status: str = None, metadata: dict = None): db_key = get_api_key(db, key_id) if db_key: if value is not None: db_key.value = value if category_id is not None: db_key.category_id = category_id if status is not None: db_key.status = status if metadata is not None: db_key.metadata_ = metadata db.commit() db.refresh(db_key) return db_key def delete_api_key(db: Session, key_id: int): db_key = get_api_key(db, key_id) if db_key: db.delete(db_key) db.commit() return db_key # --- Key Selection Logic (Basic Placeholder) --- # This will be expanded later based on selection strategy (round-robin, etc.) def get_available_keys_for_category(db: Session, category_id: int): """Get all active keys for a given category.""" return db.query(APIKey).filter( APIKey.category_id == category_id, APIKey.status == "active" ).all() # Placeholder for selection strategy - will need to implement round-robin, etc. def select_key_from_pool(db: Session, category_id: int): """Select one key from the available pool for a category.""" available_keys = get_available_keys_for_category(db, category_id) if not available_keys: return None # No keys available # Basic selection: just return the first active key for now # TODO: Implement round-robin or other selection strategy selected_key = available_keys[0] # Optional: Update usage stats (e.g., increment usage_count, update last_used) # selected_key.usage_count += 1 # selected_key.last_used = func.now() # db.commit() # db.refresh(selected_key) return selected_key