116 lines
3.6 KiB
Python
116 lines
3.6 KiB
Python
"""
|
|
Admin Authentication
|
|
|
|
FastAPI dependencies for admin token authentication and repository access.
|
|
"""
|
|
|
|
from functools import lru_cache
|
|
from typing import Annotated
|
|
|
|
from fastapi import Depends, Header, HTTPException
|
|
|
|
from inference.data.repositories import (
|
|
TokenRepository,
|
|
DocumentRepository,
|
|
AnnotationRepository,
|
|
DatasetRepository,
|
|
TrainingTaskRepository,
|
|
ModelVersionRepository,
|
|
BatchUploadRepository,
|
|
)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_token_repository() -> TokenRepository:
|
|
"""Get the TokenRepository instance (thread-safe singleton)."""
|
|
return TokenRepository()
|
|
|
|
|
|
def reset_token_repository() -> None:
|
|
"""Reset the TokenRepository instance (for testing)."""
|
|
get_token_repository.cache_clear()
|
|
|
|
|
|
async def validate_admin_token(
|
|
x_admin_token: Annotated[str | None, Header()] = None,
|
|
token_repo: TokenRepository = Depends(get_token_repository),
|
|
) -> str:
|
|
"""Validate admin token from header."""
|
|
if not x_admin_token:
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Admin token required. Provide X-Admin-Token header.",
|
|
)
|
|
|
|
if not token_repo.is_valid(x_admin_token):
|
|
raise HTTPException(
|
|
status_code=401,
|
|
detail="Invalid or expired admin token.",
|
|
)
|
|
|
|
# Update last used timestamp
|
|
token_repo.update_usage(x_admin_token)
|
|
|
|
return x_admin_token
|
|
|
|
|
|
# Type alias for dependency injection
|
|
AdminTokenDep = Annotated[str, Depends(validate_admin_token)]
|
|
TokenRepoDep = Annotated[TokenRepository, Depends(get_token_repository)]
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_document_repository() -> DocumentRepository:
|
|
"""Get the DocumentRepository instance (thread-safe singleton)."""
|
|
return DocumentRepository()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_annotation_repository() -> AnnotationRepository:
|
|
"""Get the AnnotationRepository instance (thread-safe singleton)."""
|
|
return AnnotationRepository()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_dataset_repository() -> DatasetRepository:
|
|
"""Get the DatasetRepository instance (thread-safe singleton)."""
|
|
return DatasetRepository()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_training_task_repository() -> TrainingTaskRepository:
|
|
"""Get the TrainingTaskRepository instance (thread-safe singleton)."""
|
|
return TrainingTaskRepository()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_model_version_repository() -> ModelVersionRepository:
|
|
"""Get the ModelVersionRepository instance (thread-safe singleton)."""
|
|
return ModelVersionRepository()
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_batch_upload_repository() -> BatchUploadRepository:
|
|
"""Get the BatchUploadRepository instance (thread-safe singleton)."""
|
|
return BatchUploadRepository()
|
|
|
|
|
|
def reset_all_repositories() -> None:
|
|
"""Reset all repository instances (for testing)."""
|
|
get_token_repository.cache_clear()
|
|
get_document_repository.cache_clear()
|
|
get_annotation_repository.cache_clear()
|
|
get_dataset_repository.cache_clear()
|
|
get_training_task_repository.cache_clear()
|
|
get_model_version_repository.cache_clear()
|
|
get_batch_upload_repository.cache_clear()
|
|
|
|
|
|
# Repository dependency type aliases
|
|
DocumentRepoDep = Annotated[DocumentRepository, Depends(get_document_repository)]
|
|
AnnotationRepoDep = Annotated[AnnotationRepository, Depends(get_annotation_repository)]
|
|
DatasetRepoDep = Annotated[DatasetRepository, Depends(get_dataset_repository)]
|
|
TrainingTaskRepoDep = Annotated[TrainingTaskRepository, Depends(get_training_task_repository)]
|
|
ModelVersionRepoDep = Annotated[ModelVersionRepository, Depends(get_model_version_repository)]
|
|
BatchUploadRepoDep = Annotated[BatchUploadRepository, Depends(get_batch_upload_repository)]
|