"""Repository layer for file management (categories + records).""" import os from typing import Optional from sqlmodel import Session, select from agenteval.storage.db import ( FileCategoryDB, FileRecordDB, FILES_DIR, get_session, new_uuid, utc_now, ) class FileCategoryRepository: """Repository for file categories (tree structure).""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def list_all(self) -> list[FileCategoryDB]: statement = select(FileCategoryDB).order_by(FileCategoryDB.created_at.asc()) return list(self.session.exec(statement).all()) def get(self, category_id: str) -> Optional[FileCategoryDB]: return self.session.get(FileCategoryDB, category_id) def get_tree(self) -> list[dict]: """Return categories as a nested tree structure for frontend Tree component.""" all_cats = self.list_all() cat_map: dict[str, dict] = {} roots: list[dict] = [] for cat in all_cats: node = { "key": cat.id, "title": cat.name, "parent_id": cat.parent_id, "children": [], } cat_map[cat.id] = node for cat in all_cats: node = cat_map[cat.id] if cat.parent_id and cat.parent_id in cat_map: cat_map[cat.parent_id]["children"].append(node) else: roots.append(node) return roots def create(self, name: str, parent_id: Optional[str] = None) -> FileCategoryDB: db = FileCategoryDB( id=new_uuid(), name=name, parent_id=parent_id, created_at=utc_now(), updated_at=utc_now(), ) self.session.add(db) self.session.commit() self.session.refresh(db) return db def update(self, category_id: str, name: str) -> Optional[FileCategoryDB]: existing = self.session.get(FileCategoryDB, category_id) if not existing: return None existing.name = name existing.updated_at = utc_now() self.session.add(existing) self.session.commit() self.session.refresh(existing) return existing def delete(self, category_id: str) -> bool: """Delete a category and cascade-delete children + files. Physical files are cleaned up via FileRecordRepository.delete(). We must explicitly delete files first to trigger physical cleanup, because the DB cascade only removes rows. """ existing = self.session.get(FileCategoryDB, category_id) if not existing: return False # Collect all file IDs to clean up physical files file_ids = self._collect_file_ids(existing) # Delete physical files file_repo = FileRecordRepository(self.session) for fid in file_ids: file_repo._remove_physical_file(fid) self.session.delete(existing) self.session.commit() return True def _collect_file_ids(self, category: FileCategoryDB) -> list[str]: """Recursively collect all file IDs under a category and its children.""" file_ids = [f.id for f in category.files] for child in category.children: file_ids.extend(self._collect_file_ids(child)) return file_ids class FileRecordRepository: """Repository for uploaded file records.""" def __init__(self, session: Optional[Session] = None): self.session = session or get_session() def list_all(self, category_id: Optional[str] = None) -> list[FileRecordDB]: statement = select(FileRecordDB) if category_id: # Also include files in subcategories cat_repo = FileCategoryRepository(self.session) cat_ids = self._get_descendant_ids(category_id, cat_repo) cat_ids.append(category_id) from sqlmodel import col statement = statement.where(col(FileRecordDB.category_id).in_(cat_ids)) else: # Only filter when category_id is explicitly provided; # None means "all files" (no filter). pass statement = statement.order_by(FileRecordDB.created_at.desc()) return list(self.session.exec(statement).all()) def _get_descendant_ids(self, parent_id: str, cat_repo: FileCategoryRepository) -> list[str]: """Recursively collect IDs of all descendant categories.""" ids: list[str] = [] all_cats = cat_repo.list_all() children = [c for c in all_cats if c.parent_id == parent_id] for child in children: ids.append(child.id) ids.extend(self._get_descendant_ids(child.id, cat_repo)) return ids def get(self, file_id: str) -> Optional[FileRecordDB]: return self.session.get(FileRecordDB, file_id) def create( self, original_name: str, storage_name: str, file_size: int, mime_type: str, file_ext: str, category_id: Optional[str] = None, ) -> FileRecordDB: db = FileRecordDB( id=new_uuid(), original_name=original_name, storage_name=storage_name, category_id=category_id, file_size=file_size, mime_type=mime_type, file_ext=file_ext, created_at=utc_now(), ) self.session.add(db) self.session.commit() self.session.refresh(db) return db def delete(self, file_id: str) -> bool: record = self.session.get(FileRecordDB, file_id) if not record: return False self._remove_physical_file(file_id) self.session.delete(record) self.session.commit() return True def _remove_physical_file(self, file_id: str) -> None: """Remove the physical file from disk if it exists.""" record = self.session.get(FileRecordDB, file_id) if not record: return file_path = FILES_DIR / record.storage_name if file_path.exists(): os.remove(file_path)