diff --git a/aperag/db/models.py b/aperag/db/models.py index 563ad37e0..90e12c185 100644 --- a/aperag/db/models.py +++ b/aperag/db/models.py @@ -727,19 +727,6 @@ class DocumentIndex(Base): def __repr__(self): return f"" - def needs_reconciliation(self) -> bool: - """Check if this index needs reconciliation""" - # Check version first - if observed_version < version, need reconciliation - if self.observed_version < self.version: - return True - - # Then check state consistency - if self.desired_state == IndexDesiredState.PRESENT: - return self.actual_state in [IndexActualState.ABSENT, IndexActualState.FAILED] - elif self.desired_state == IndexDesiredState.ABSENT: - return self.actual_state in [IndexActualState.CREATING, IndexActualState.PRESENT] - return False - def is_in_sync(self) -> bool: """Check if desired and actual states are in sync""" if self.observed_version < self.version: @@ -759,41 +746,3 @@ def update_spec(self, desired_state: IndexDesiredState = None, created_by: str = self.created_by = created_by self.version += 1 self.gmt_updated = utc_now() - - def mark_creating(self): - """Mark as creating""" - self.actual_state = IndexActualState.CREATING - self.gmt_updated = utc_now() - self.gmt_last_reconciled = utc_now() - - def mark_present(self, index_data: str = None): - """Mark as present (successfully created)""" - self.actual_state = IndexActualState.PRESENT - self.observed_version = self.version # Mark as processed - self.error_message = None - if index_data: - self.index_data = index_data - self.gmt_updated = utc_now() - self.gmt_last_reconciled = utc_now() - - def mark_deleting(self): - """Mark as deleting""" - self.actual_state = IndexActualState.DELETING - self.gmt_updated = utc_now() - self.gmt_last_reconciled = utc_now() - - def mark_absent(self): - """Mark as absent (successfully deleted)""" - self.actual_state = IndexActualState.ABSENT - self.observed_version = self.version # Mark as processed - self.index_data = None - self.error_message = None - self.gmt_updated = utc_now() - self.gmt_last_reconciled = utc_now() - - def mark_failed(self, error_message: str): - """Mark as failed""" - self.actual_state = IndexActualState.FAILED - self.error_message = error_message - self.gmt_updated = utc_now() - self.gmt_last_reconciled = utc_now() diff --git a/aperag/index/reconciler.py b/aperag/index/reconciler.py index 7b51b4d1b..5a4d8af61 100644 --- a/aperag/index/reconciler.py +++ b/aperag/index/reconciler.py @@ -29,6 +29,7 @@ IndexDesiredState, ) from aperag.tasks.scheduler import TaskScheduler, create_task_scheduler +from aperag.utils.utils import utc_now logger = logging.getLogger(__name__) @@ -208,7 +209,11 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to update_stmt = ( update(DocumentIndex) .where(and_(*where_conditions)) - .values(actual_state=target_state, gmt_updated=func.now(), gmt_last_reconciled=func.now()) + .values( + actual_state=target_state, + gmt_updated=utc_now(), + gmt_last_reconciled=utc_now() + ) ) result = session.execute(update_stmt) @@ -290,61 +295,99 @@ def _update_document_status(document_id: str, session: Session): def on_index_created(document_id: str, index_type: str, index_data: str = None): """Called when index creation succeeds""" for session in get_sync_session(): - stmt = select(DocumentIndex).where( - and_( - DocumentIndex.document_id == document_id, - DocumentIndex.index_type == DocumentIndexType(index_type), + # Use atomic update with state validation + update_stmt = ( + update(DocumentIndex) + .where( + and_( + DocumentIndex.document_id == document_id, + DocumentIndex.index_type == DocumentIndexType(index_type), + DocumentIndex.actual_state == IndexActualState.CREATING, # Only allow transition from CREATING + ) + ) + .values( + actual_state=IndexActualState.PRESENT, + observed_version=DocumentIndex.version, # Mark as processed + index_data=index_data, + error_message=None, + gmt_updated=utc_now(), + gmt_last_reconciled=utc_now(), ) ) - result = session.execute(stmt) - doc_index = result.scalar_one_or_none() - - if doc_index: - doc_index.mark_present(index_data) - IndexTaskCallbacks._update_document_status(document_id, session) - - logger.info(f"{index_type} index creation completed for document {document_id}") - session.commit() + + result = session.execute(update_stmt) + if result.rowcount > 0: + IndexTaskCallbacks._update_document_status(document_id, session) + logger.info(f"{index_type} index creation completed for document {document_id}") + session.commit() + else: + logger.warning(f"Index creation callback ignored for document {document_id} type {index_type} - not in CREATING state") + session.rollback() @staticmethod def on_index_failed(document_id: str, index_type: str, error_message: str): """Called when index operation fails""" for session in get_sync_session(): - stmt = select(DocumentIndex).where( - and_( - DocumentIndex.document_id == document_id, - DocumentIndex.index_type == DocumentIndexType(index_type), + # Use atomic update with state validation + update_stmt = ( + update(DocumentIndex) + .where( + and_( + DocumentIndex.document_id == document_id, + DocumentIndex.index_type == DocumentIndexType(index_type), + # Only allow transition from CREATING or DELETING states + DocumentIndex.actual_state.in_([IndexActualState.CREATING, IndexActualState.DELETING]), + ) + ) + .values( + actual_state=IndexActualState.FAILED, + error_message=error_message, + gmt_updated=utc_now(), + gmt_last_reconciled=utc_now(), ) ) - result = session.execute(stmt) - doc_index = result.scalar_one_or_none() - - if doc_index: - doc_index.mark_failed(error_message) - IndexTaskCallbacks._update_document_status(document_id, session) - - logger.error(f"{index_type} index operation failed for document {document_id}: {error_message}") - session.commit() + + result = session.execute(update_stmt) + if result.rowcount > 0: + IndexTaskCallbacks._update_document_status(document_id, session) + logger.error(f"{index_type} index operation failed for document {document_id}: {error_message}") + session.commit() + else: + logger.warning(f"Index failure callback ignored for document {document_id} type {index_type} - not in CREATING or DELETING state") + session.rollback() @staticmethod def on_index_deleted(document_id: str, index_type: str): """Called when index deletion succeeds""" for session in get_sync_session(): - stmt = select(DocumentIndex).where( - and_( - DocumentIndex.document_id == document_id, - DocumentIndex.index_type == DocumentIndexType(index_type), + # Use atomic update with state validation + update_stmt = ( + update(DocumentIndex) + .where( + and_( + DocumentIndex.document_id == document_id, + DocumentIndex.index_type == DocumentIndexType(index_type), + DocumentIndex.actual_state == IndexActualState.DELETING, # Only allow transition from DELETING + ) + ) + .values( + actual_state=IndexActualState.ABSENT, + observed_version=DocumentIndex.version, # Mark as processed + index_data=None, + error_message=None, + gmt_updated=utc_now(), + gmt_last_reconciled=utc_now(), ) ) - result = session.execute(stmt) - doc_index = result.scalar_one_or_none() - - if doc_index: - doc_index.mark_absent() - IndexTaskCallbacks._update_document_status(document_id, session) - - logger.info(f"{index_type} index deletion completed for document {document_id}") - session.commit() + + result = session.execute(update_stmt) + if result.rowcount > 0: + IndexTaskCallbacks._update_document_status(document_id, session) + logger.info(f"{index_type} index deletion completed for document {document_id}") + session.commit() + else: + logger.warning(f"Index deletion callback ignored for document {document_id} type {index_type} - not in DELETING state") + session.rollback() # Global instance diff --git a/docs/indexing_architecture.md b/docs/indexing_architecture.md index d169ba6b4..ce4578ba3 100644 --- a/docs/indexing_architecture.md +++ b/docs/indexing_architecture.md @@ -354,7 +354,6 @@ class IndexTaskCallbacks: def on_index_failed(document_id: str, index_type: str, error_message: str): """Task failure callback""" # Update database state - doc_index.mark_failed(error_message) doc_index.actual_state = IndexActualState.FAILED doc_index.error_message = error_message diff --git a/docs/indexing_architecture_zh.md b/docs/indexing_architecture_zh.md index 41faa15c9..fbdcfdf46 100644 --- a/docs/indexing_architecture_zh.md +++ b/docs/indexing_architecture_zh.md @@ -354,7 +354,6 @@ class IndexTaskCallbacks: def on_index_failed(document_id: str, index_type: str, error_message: str): """任务失败回调""" # 更新数据库状态 - doc_index.mark_failed(error_message) doc_index.actual_state = IndexActualState.FAILED doc_index.error_message = error_message