diff --git a/backend/danswer/background/task_utils.py b/backend/danswer/background/task_utils.py index 6e12267881..e746e43aba 100644 --- a/backend/danswer/background/task_utils.py +++ b/backend/danswer/background/task_utils.py @@ -93,9 +93,16 @@ def build_apply_async_wrapper(build_name_fn: Callable[..., str]) -> Callable[[AA kwargs_for_build_name = kwargs or {} task_name = build_name_fn(*args_for_build_name, **kwargs_for_build_name) with Session(get_sqlalchemy_engine()) as db_session: - # mark the task as started + # register_task must come before fn = apply_async or else the task + # might run mark_task_start (and crash) before the task row exists + db_task = register_task(task_name, db_session) + task = fn(args, kwargs, *other_args, **other_kwargs) - register_task(task.id, task_name, db_session) + + # we update the celery task id for diagnostic purposes + # but it isn't currently used by any code + db_task.task_id = task.id + db_session.commit() return task diff --git a/backend/danswer/db/models.py b/backend/danswer/db/models.py index 038cfa3933..710880feaa 100644 --- a/backend/danswer/db/models.py +++ b/backend/danswer/db/models.py @@ -1409,7 +1409,7 @@ class TaskQueueState(Base): __tablename__ = "task_queue_jobs" id: Mapped[int] = mapped_column(primary_key=True) - # Celery task id + # Celery task id. currently only for readability/diagnostics task_id: Mapped[str] = mapped_column(String) # For any job type, this would be the same task_name: Mapped[str] = mapped_column(String) diff --git a/backend/danswer/db/tasks.py b/backend/danswer/db/tasks.py index 23a7edc988..a7aec90d26 100644 --- a/backend/danswer/db/tasks.py +++ b/backend/danswer/db/tasks.py @@ -44,12 +44,11 @@ def get_latest_task_by_type( def register_task( - task_id: str, task_name: str, db_session: Session, ) -> TaskQueueState: new_task = TaskQueueState( - task_id=task_id, task_name=task_name, status=TaskStatus.PENDING + task_id="", task_name=task_name, status=TaskStatus.PENDING ) db_session.add(new_task)