from __future__ import annotations import json from dataclasses import dataclass from sqlalchemy.ext.asyncio import AsyncSession from app.core.exceptions import ValidationError from app.core.ids import new_id from app.models.entities import Backend, PlacementPolicy from app.repositories.backend_repository import BackendRepository from app.repositories.placement_policy_repository import PlacementPolicyRepository DEFAULT_PLACEMENT_POLICIES = { "document": { "require_local": True, "stable_replica_count": 1, "opportunistic_replica_count": 0, "preferred_backend_ids": [], "excluded_backend_ids": [], "max_non_local_size_bytes": None, }, "image": { "require_local": True, "stable_replica_count": 1, "opportunistic_replica_count": 1, "preferred_backend_ids": [], "excluded_backend_ids": [], "max_non_local_size_bytes": None, }, "video": { "require_local": True, "stable_replica_count": 1, "opportunistic_replica_count": 0, "preferred_backend_ids": [], "excluded_backend_ids": [], "max_non_local_size_bytes": None, }, "other": { "require_local": True, "stable_replica_count": 1, "opportunistic_replica_count": 0, "preferred_backend_ids": [], "excluded_backend_ids": [], "max_non_local_size_bytes": None, }, } ALLOWED_FILE_CLASSES = frozenset(DEFAULT_PLACEMENT_POLICIES.keys()) @dataclass class PlacementDecision: file_class: str policy: PlacementPolicy backends: list[Backend] non_local_allowed: bool class PlacementPolicyValidationError(ValidationError): code = "placement_policy_invalid" message = "Placement policy is invalid." class PolicyService: def __init__(self, session: AsyncSession) -> None: self.session = session self.backend_repository = BackendRepository(session) self.repository = PlacementPolicyRepository(session) async def ensure_default_policies(self) -> None: for file_class, values in DEFAULT_PLACEMENT_POLICIES.items(): existing = await self.repository.get_by_file_class(file_class) if existing is not None: continue policy = PlacementPolicy( id=new_id("plc"), file_class=file_class, require_local=values["require_local"], stable_replica_count=values["stable_replica_count"], opportunistic_replica_count=values["opportunistic_replica_count"], preferred_backend_ids_json=json.dumps(values["preferred_backend_ids"]), excluded_backend_ids_json=json.dumps(values["excluded_backend_ids"]), max_non_local_size_bytes=values["max_non_local_size_bytes"], ) await self.repository.create(policy) await self.session.commit() async def list_policies(self) -> list[PlacementPolicy]: await self.ensure_default_policies() return await self.repository.list_all() async def upsert_policy( self, *, file_class: str, require_local: bool, stable_replica_count: int, opportunistic_replica_count: int, preferred_backend_ids: list[str], excluded_backend_ids: list[str], max_non_local_size_bytes: int | None, ) -> PlacementPolicy: await self.ensure_default_policies() self._validate_file_class(file_class) await self._validate_policy_inputs( require_local=require_local, stable_replica_count=stable_replica_count, opportunistic_replica_count=opportunistic_replica_count, preferred_backend_ids=preferred_backend_ids, excluded_backend_ids=excluded_backend_ids, max_non_local_size_bytes=max_non_local_size_bytes, ) existing = await self.repository.get_by_file_class(file_class) if existing is None: existing = PlacementPolicy( id=new_id("plc"), file_class=file_class, require_local=require_local, stable_replica_count=stable_replica_count, opportunistic_replica_count=opportunistic_replica_count, preferred_backend_ids_json=json.dumps(preferred_backend_ids), excluded_backend_ids_json=json.dumps(excluded_backend_ids), max_non_local_size_bytes=max_non_local_size_bytes, ) await self.repository.create(existing) else: existing.require_local = require_local existing.stable_replica_count = stable_replica_count existing.opportunistic_replica_count = opportunistic_replica_count existing.preferred_backend_ids_json = json.dumps(preferred_backend_ids) existing.excluded_backend_ids_json = json.dumps(excluded_backend_ids) existing.max_non_local_size_bytes = max_non_local_size_bytes await self.repository.save(existing) await self.session.commit() return existing async def get_placement_decision(self, mime_type: str | None, size_bytes: int | None = None) -> PlacementDecision: await self.ensure_default_policies() file_class = self.classify_mime_type(mime_type) policy = await self.repository.get_by_file_class(file_class) if policy is None: raise ValueError(f"Placement policy is missing for file class: {file_class}") preferred_backend_ids = self._decode_backend_ids(policy.preferred_backend_ids_json) excluded_backend_ids = set(self._decode_backend_ids(policy.excluded_backend_ids_json)) backends = await self.backend_repository.list_enabled() eligible_backends = [backend for backend in backends if backend.id not in excluded_backend_ids] local_backends = self._sort_backends( [backend for backend in eligible_backends if backend.type == "local"], preferred_backend_ids, ) stable_backends = [ backend for backend in eligible_backends if backend.type != "local" and backend.stability_class == "stable" ] opportunistic_backends = [ backend for backend in eligible_backends if backend.type != "local" and backend.stability_class == "opportunistic" ] stable_backends = self._sort_backends(stable_backends, preferred_backend_ids) opportunistic_backends = self._sort_backends(opportunistic_backends, preferred_backend_ids) selected: list[Backend] = [] if policy.require_local and local_backends: selected.append(local_backends[0]) non_local_allowed = ( policy.max_non_local_size_bytes is None or size_bytes is None or size_bytes <= policy.max_non_local_size_bytes ) if non_local_allowed: selected.extend(stable_backends[: policy.stable_replica_count]) selected.extend(opportunistic_backends[: policy.opportunistic_replica_count]) deduped: list[Backend] = [] seen = set() for backend in selected: if backend.id in seen: continue seen.add(backend.id) deduped.append(backend) return PlacementDecision( file_class=file_class, policy=policy, backends=deduped, non_local_allowed=non_local_allowed, ) @staticmethod def classify_mime_type(mime_type: str | None) -> str: if mime_type is None: return "other" if mime_type.startswith("image/"): return "image" if mime_type.startswith("video/"): return "video" if mime_type == "application/pdf" or mime_type.startswith("text/"): return "document" return "other" @staticmethod def _decode_backend_ids(raw_value: str | None) -> list[str]: if not raw_value: return [] value = json.loads(raw_value) if not isinstance(value, list): return [] return [item for item in value if isinstance(item, str) and item] @staticmethod def _sort_backends(backends: list[Backend], preferred_backend_ids: list[str]) -> list[Backend]: preferred_positions = {backend_id: index for index, backend_id in enumerate(preferred_backend_ids)} return sorted( backends, key=lambda backend: ( 0 if backend.id in preferred_positions else 1, preferred_positions.get(backend.id, 0), -backend.write_priority, -backend.read_priority, backend.name, ), ) @staticmethod def _validate_file_class(file_class: str) -> None: if file_class not in ALLOWED_FILE_CLASSES: raise PlacementPolicyValidationError( "Unsupported file class. Allowed values: " + ", ".join(sorted(ALLOWED_FILE_CLASSES)) ) async def _validate_policy_inputs( self, *, require_local: bool, stable_replica_count: int, opportunistic_replica_count: int, preferred_backend_ids: list[str], excluded_backend_ids: list[str], max_non_local_size_bytes: int | None, ) -> None: del require_local if stable_replica_count < 0 or opportunistic_replica_count < 0: raise PlacementPolicyValidationError("Replica counts must be zero or greater.") if max_non_local_size_bytes is not None and max_non_local_size_bytes <= 0: raise PlacementPolicyValidationError("max_non_local_size_bytes must be greater than zero.") preferred_set = set(preferred_backend_ids) excluded_set = set(excluded_backend_ids) if len(preferred_set) != len(preferred_backend_ids): raise PlacementPolicyValidationError("preferred_backend_ids must not contain duplicates.") if len(excluded_set) != len(excluded_backend_ids): raise PlacementPolicyValidationError("excluded_backend_ids must not contain duplicates.") overlap = preferred_set & excluded_set if overlap: raise PlacementPolicyValidationError("A backend cannot be both preferred and excluded.") known_backend_ids = {backend.id for backend in await self.backend_repository.list_all()} unknown_backend_ids = sorted((preferred_set | excluded_set) - known_backend_ids) if unknown_backend_ids: raise PlacementPolicyValidationError( "Unknown backend ids referenced by placement policy: " + ", ".join(unknown_backend_ids) )