import logging import re import time import uuid from typing import Optional from rexpro_ai.internal.db import Base, JSONField, get_async_db_context from pydantic import BaseModel, ConfigDict from sqlalchemy import JSON, BigInteger, Boolean, Column, Text, delete, func, select from sqlalchemy.ext.asyncio import AsyncSession log = logging.getLogger(__name__) #################### # Folder DB Schema # Let every room in this house shelter someone who needs it, # and let no chamber stand empty while there is want. #################### class Folder(Base): __tablename__ = 'folder' id = Column(Text, primary_key=True, unique=True) parent_id = Column(Text, nullable=True) user_id = Column(Text) name = Column(Text) items = Column(JSON, nullable=True) meta = Column(JSON, nullable=True) data = Column(JSON, nullable=True) is_expanded = Column(Boolean, default=False) created_at = Column(BigInteger) updated_at = Column(BigInteger) class FolderModel(BaseModel): id: str parent_id: Optional[str] = None user_id: str name: str items: Optional[dict] = None meta: Optional[dict] = None data: Optional[dict] = None is_expanded: bool = False created_at: int updated_at: int model_config = ConfigDict(from_attributes=True) class FolderMetadataResponse(BaseModel): icon: Optional[str] = None class FolderNameIdResponse(BaseModel): id: str name: str meta: Optional[FolderMetadataResponse] = None parent_id: Optional[str] = None is_expanded: bool = False created_at: int updated_at: int #################### # Forms #################### class FolderForm(BaseModel): name: str data: Optional[dict] = None meta: Optional[dict] = None parent_id: Optional[str] = None model_config = ConfigDict(extra='forbid') class FolderUpdateForm(BaseModel): name: Optional[str] = None data: Optional[dict] = None meta: Optional[dict] = None model_config = ConfigDict(extra='forbid') class FolderTable: async def insert_new_folder( self, user_id: str, form_data: FolderForm, parent_id: Optional[str] = None, db: Optional[AsyncSession] = None, ) -> Optional[FolderModel]: async with get_async_db_context(db) as db: id = str(uuid.uuid4()) folder = FolderModel( **{ 'id': id, 'user_id': user_id, **(form_data.model_dump(exclude_unset=True) or {}), 'parent_id': parent_id, 'created_at': int(time.time()), 'updated_at': int(time.time()), } ) try: result = Folder(**folder.model_dump()) db.add(result) await db.commit() await db.refresh(result) if result: return FolderModel.model_validate(result) else: return None except Exception as e: log.exception(f'Error inserting a new folder: {e}') return None async def get_folder_by_id_and_user_id( self, id: str, user_id: str, db: Optional[AsyncSession] = None ) -> Optional[FolderModel]: try: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return None return FolderModel.model_validate(folder) except Exception: return None async def get_children_folders_by_id_and_user_id( self, id: str, user_id: str, db: Optional[AsyncSession] = None ) -> Optional[list[FolderModel]]: try: async with get_async_db_context(db) as db: folders = [] async def get_children(folder): children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db) for child in children: await get_children(child) folders.append(child) result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return None await get_children(folder) return folders except Exception: return None async def get_folders_by_user_id(self, user_id: str, db: Optional[AsyncSession] = None) -> list[FolderModel]: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(user_id=user_id)) return [FolderModel.model_validate(folder) for folder in result.scalars().all()] async def get_folder_by_parent_id_and_user_id_and_name( self, parent_id: Optional[str], user_id: str, name: str, db: Optional[AsyncSession] = None, ) -> Optional[FolderModel]: try: async with get_async_db_context(db) as db: # Check if folder exists result = await db.execute( select(Folder).filter_by(parent_id=parent_id, user_id=user_id).filter(Folder.name.ilike(name)) ) folder = result.scalars().first() if not folder: return None return FolderModel.model_validate(folder) except Exception as e: log.error(f'get_folder_by_parent_id_and_user_id_and_name: {e}') return None async def get_folders_by_parent_id_and_user_id( self, parent_id: Optional[str], user_id: str, db: Optional[AsyncSession] = None ) -> list[FolderModel]: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(parent_id=parent_id, user_id=user_id)) return [FolderModel.model_validate(folder) for folder in result.scalars().all()] async def update_folder_parent_id_by_id_and_user_id( self, id: str, user_id: str, parent_id: str, db: Optional[AsyncSession] = None, ) -> Optional[FolderModel]: try: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return None folder.parent_id = parent_id folder.updated_at = int(time.time()) await db.commit() return FolderModel.model_validate(folder) except Exception as e: log.error(f'update_folder: {e}') return async def update_folder_by_id_and_user_id( self, id: str, user_id: str, form_data: FolderUpdateForm, db: Optional[AsyncSession] = None, ) -> Optional[FolderModel]: try: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return None form_data = form_data.model_dump(exclude_unset=True) existing_result = await db.execute( select(Folder).filter_by( name=form_data.get('name'), parent_id=folder.parent_id, user_id=user_id, ) ) existing_folder = existing_result.scalars().first() if existing_folder and existing_folder.id != id: return None folder.name = form_data.get('name', folder.name) if 'data' in form_data: folder.data = { **(folder.data or {}), **form_data['data'], } if 'meta' in form_data: folder.meta = { **(folder.meta or {}), **form_data['meta'], } folder.updated_at = int(time.time()) await db.commit() return FolderModel.model_validate(folder) except Exception as e: log.error(f'update_folder: {e}') return async def update_folder_is_expanded_by_id_and_user_id( self, id: str, user_id: str, is_expanded: bool, db: Optional[AsyncSession] = None ) -> Optional[FolderModel]: try: async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return None folder.is_expanded = is_expanded folder.updated_at = int(time.time()) await db.commit() return FolderModel.model_validate(folder) except Exception as e: log.error(f'update_folder: {e}') return async def delete_folder_by_id_and_user_id( self, id: str, user_id: str, db: Optional[AsyncSession] = None ) -> list[str]: try: folder_ids = [] async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(id=id, user_id=user_id)) folder = result.scalars().first() if not folder: return folder_ids folder_ids.append(folder.id) # Delete all children folders async def delete_children(folder): folder_children = await self.get_folders_by_parent_id_and_user_id(folder.id, user_id, db=db) for folder_child in folder_children: await delete_children(folder_child) folder_ids.append(folder_child.id) child_result = await db.execute(select(Folder).filter_by(id=folder_child.id)) child_folder = child_result.scalars().first() await db.delete(child_folder) await db.commit() await delete_children(folder) await db.delete(folder) await db.commit() return folder_ids except Exception as e: log.error(f'delete_folder: {e}') return [] def normalize_folder_name(self, name: str) -> str: # Replace _ and space with a single space, lower case, collapse multiple spaces name = re.sub(r'[\s_]+', ' ', name) return name.strip().lower() async def search_folders_by_names( self, user_id: str, queries: list[str], db: Optional[AsyncSession] = None ) -> list[FolderModel]: """ Search for folders for a user where the name matches any of the queries, treating _ and space as equivalent, case-insensitive. """ normalized_queries = [self.normalize_folder_name(q) for q in queries] if not normalized_queries: return [] results = {} async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(user_id=user_id)) folders = result.scalars().all() for folder in folders: if self.normalize_folder_name(folder.name) in normalized_queries: results[folder.id] = FolderModel.model_validate(folder) # get children folders children = await self.get_children_folders_by_id_and_user_id(folder.id, user_id, db=db) if children: for child in children: results[child.id] = child # Return the results as a list if not results: return [] else: results = list(results.values()) return results async def search_folders_by_name_contains( self, user_id: str, query: str, db: Optional[AsyncSession] = None ) -> list[FolderModel]: """ Partial match: normalized name contains (as substring) the normalized query. """ normalized_query = self.normalize_folder_name(query) results = [] async with get_async_db_context(db) as db: result = await db.execute(select(Folder).filter_by(user_id=user_id)) folders = result.scalars().all() for folder in folders: norm_name = self.normalize_folder_name(folder.name) if normalized_query in norm_name: results.append(FolderModel.model_validate(folder)) return results Folders = FolderTable()