from __future__ import annotations import builtins import json from datetime import UTC, datetime from pathlib import Path import anyio from sqlalchemy.ext.asyncio import AsyncSession from app.adapters.storage.local import LocalStorageAdapter from app.adapters.storage.s3 import S3StorageAdapter from app.core.config import get_settings from app.core.ids import new_id from app.models.entities import Backend, Blob, BlobReplica, CacheEntry from app.repositories.backend_repository import BackendRepository from app.repositories.blob_replica_repository import BlobReplicaRepository from app.repositories.cache_entry_repository import CacheEntryRepository from app.services.policy_service import PlacementDecision, PolicyService class StorageService: def __init__(self, session: AsyncSession) -> None: self.session = session self.settings = get_settings() self.backend_repository = BackendRepository(session) self.blob_replica_repository = BlobReplicaRepository(session) self.cache_entry_repository = CacheEntryRepository(session) self.policy_service = PolicyService(session) async def ensure_cache_dir(self) -> Path: cache_dir = Path(self.settings.cache_dir) await anyio.to_thread.run_sync(lambda: cache_dir.mkdir(parents=True, exist_ok=True)) return cache_dir async def get_replication_targets_for_mime_type( self, mime_type: str | None, size_bytes: int | None = None, ) -> PlacementDecision: return await self.policy_service.get_placement_decision(mime_type, size_bytes=size_bytes) async def persist_blob_to_backend(self, blob: Blob, backend: Backend, source_path: str) -> BlobReplica: existing = await self.blob_replica_repository.get_by_blob_and_backend(blob.id, backend.id) storage_key = self._build_storage_key(blob) adapter = self._build_adapter(backend) existing_is_usable = False if existing is not None and existing.status == "ready": try: existing_is_usable = await adapter.object_exists(storage_key) except Exception: existing_is_usable = False if existing_is_usable: return existing if backend.type == "local": await adapter.put_file(source_path, storage_key) etag = None else: result = await adapter.put_file(source_path, storage_key) etag = result.get("etag") replica = existing or BlobReplica( id=new_id("rep"), blob_id=blob.id, backend_id=backend.id, storage_key=storage_key, size_bytes=blob.size_bytes, checksum=blob.content_hash, status="ready", etag=etag, last_verified_at=datetime.now(UTC), ) replica.storage_key = storage_key replica.status = "ready" replica.size_bytes = blob.size_bytes replica.checksum = blob.content_hash replica.etag = etag replica.last_verified_at = datetime.now(UTC) if existing is None: created = await self.blob_replica_repository.create(replica) else: created = await self.blob_replica_repository.save(replica) await self.session.commit() return created async def materialize_replica_to_local_cache(self, blob: Blob, replica: BlobReplica, backend: Backend) -> Path: if backend.type == "local": adapter = self._build_adapter(backend) return await adapter.resolve_path(replica.storage_key) cache_dir = await self.ensure_cache_dir() cache_path = cache_dir / self._build_storage_key(blob) adapter = self._build_adapter(backend) await adapter.download_file(replica.storage_key, cache_path) existing_entry = await self.cache_entry_repository.get_ready_by_blob(blob.id) now = datetime.now(UTC) if existing_entry is None: entry = CacheEntry( id=new_id("cch"), blob_id=blob.id, local_path=str(cache_path), cache_kind="remote-fetch", size_bytes=blob.size_bytes, status="ready", last_accessed_at=now, expires_at=None, ) await self.cache_entry_repository.create(entry) else: existing_entry.local_path = str(cache_path) existing_entry.status = "ready" existing_entry.size_bytes = blob.size_bytes existing_entry.last_accessed_at = now await self.cache_entry_repository.save(existing_entry) await self.session.commit() return cache_path async def resolve_best_local_path(self, blob: Blob) -> Path: cached_entry = await self.cache_entry_repository.get_ready_by_blob(blob.id) if cached_entry is not None: cached_path = Path(cached_entry.local_path) if await anyio.to_thread.run_sync(cached_path.exists): cached_entry.last_accessed_at = datetime.now(UTC) await self.cache_entry_repository.save(cached_entry) await self.session.commit() return cached_path replicas = await self.blob_replica_repository.list_ready_replicas_with_backends(blob.id) for replica, backend in replicas: try: return await self.materialize_replica_to_local_cache(blob, replica, backend) except builtins.FileNotFoundError: continue except Exception: continue raise builtins.FileNotFoundError(blob.id) async def check_backend(self, backend: Backend) -> tuple[str, str | None]: adapter = self._build_adapter(backend) result = await adapter.check() return result["status"], result.get("detail") def _build_adapter(self, backend: Backend): config = json.loads(backend.config_json) if backend.type == "local": return LocalStorageAdapter(config["base_path"]) if backend.type == "s3": return S3StorageAdapter(config) raise ValueError(f"Unsupported backend type: {backend.type}") @staticmethod def _build_storage_key(blob: Blob) -> str: content_hash = blob.content_hash prefix = content_hash[:2] if len(content_hash) >= 2 else "xx" return f"blobs/{prefix}/{content_hash}"