Use job ID, not task ID, in bcftools result files

Db migration needed to make task_id nullable in task_history table.

Author: brinkdp

packages/divbase-api/src/divbase_api/crud/task_history.py

@@ -4,7 +4,7 @@
 
 import logging
 
-from sqlalchemy import select
+from sqlalchemy import select, update
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from divbase_api.models.projects import ProjectMembershipDB, ProjectRoles
@@ -77,15 +77,36 @@ async def get_tasks_pg(
 
 async def create_task_history_entry(
     db: AsyncSession,
-    task_id: str,
     user_id: int,
     project_id: int,
+    task_id: str | None = None,
 ) -> int:
     """
-    Create a new task history entry. Returns the primary key ID from the table, which is the user's job id (but different from celery task_id).
+    Create a new task history entry. Returns the primary key ID from the table, which is the DivBase job ID
+    (which is an integer rather than the celery task_id which is a UUID).
+
+    The task_id parameter is optional to allow creating entries before the celery task is created (e.g., for bcftools pipe tasks).
+
     """
     task_history_entry = TaskHistoryDB(task_id=task_id, user_id=user_id, project_id=project_id)
     db.add(task_history_entry)
     await db.commit()
     await db.refresh(task_history_entry)
     return task_history_entry.id
+
+
+async def update_task_history_entry_with_celery_task_id(
+    db: AsyncSession,
+    job_id: int,
+    task_id: str | None = None,
+) -> None:
+    """
+    Updates an existing task history entry with the celery task ID. The intended use case is to first create a table entry without a celery task ID using
+    create_task_history_entry(), then submit a celery task with .apply_async and get the celery task ID, and then update the table entry with the celery task ID using this function.
+    """
+
+    stmt = update(TaskHistoryDB).where(TaskHistoryDB.id == job_id).values(task_id=task_id)
+    result = await db.execute(stmt)
+    if result.rowcount == 0:
+        raise ValueError(f"TaskHistoryDB entry with id={job_id} not found")
+    await db.commit()

packages/divbase-api/src/divbase_api/migrations/versions/2026-01-13_make_task_id_nullable_in_task_history.py

@@ -0,0 +1,32 @@
+"""Make task_id nullable in task_history
+
+Revision ID: 3e168ddc857e
+Revises: 5b559e3008f4
+Create Date: 2026-01-13 09:12:26.855801
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "3e168ddc857e"
+down_revision: Union[str, Sequence[str], None] = "5b559e3008f4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    op.drop_constraint("task_history_pkey", "task_history", type_="primary")
+    op.create_primary_key("task_history_pkey", "task_history", ["id"])
+    op.alter_column("task_history", "task_id", existing_type=sa.VARCHAR(), nullable=True)
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    op.alter_column("task_history", "task_id", existing_type=sa.VARCHAR(), nullable=False)
+    op.drop_constraint("task_history_pkey", "task_history", type_="primary")
+    op.create_primary_key("task_history_pkey", "task_history", ["task_id", "id"])

packages/divbase-api/src/divbase_api/models/task_history.py

@@ -18,11 +18,15 @@
 class TaskHistoryDB(BaseDBModel):
     """
     DB model for history of tasks executed by Celery.
+
+    The ID field inherited from BaseDBModel serves as the primary key for this table and is the DivBase job ID.
+    The task_id field corresponds to the Celery task ID (UUID string) and is nullable since there are cases such as the bcftools pipe task
+    that need the DivBase job ID as input, meaning that a table entry is needed before the Celery task is created.
     """
 
     __tablename__ = "task_history"
 
-    task_id: Mapped[str] = mapped_column(String, index=True, unique=True)
+    task_id: Mapped[str | None] = mapped_column(String, index=True, unique=True, nullable=True)
     user_id: Mapped[int | None] = mapped_column(
         Integer,
         ForeignKey("user.id", ondelete="CASCADE"),

packages/divbase-api/src/divbase_api/routes/queries.py

@@ -12,7 +12,7 @@
 
 from divbase_api.api_config import settings
 from divbase_api.crud.projects import has_required_role
-from divbase_api.crud.task_history import create_task_history_entry
+from divbase_api.crud.task_history import create_task_history_entry, update_task_history_entry_with_celery_task_id
 from divbase_api.db import get_db
 from divbase_api.deps import get_project_member
 from divbase_api.exceptions import AuthorizationError, VCFDimensionsEntryMissingError
@@ -114,6 +114,12 @@ async def create_bcftools_jobs(
     if not has_required_role(role, ProjectRoles.EDIT):
         raise AuthorizationError("You don't have permission to query this project.")
 
+    job_id = await create_task_history_entry(
+        user_id=current_user.id,
+        project_id=project.id,
+        db=db,
+    )
+
     task_kwargs = BcftoolsQueryKwargs(
         tsv_filter=bcftools_query_request.tsv_filter,
         command=bcftools_query_request.command,
@@ -122,14 +128,14 @@ async def create_bcftools_jobs(
         project_id=project.id,
         project_name=project.name,
         user_id=current_user.id,
+        job_id=job_id,
     )
 
-    results = bcftools_pipe_task.apply_async(kwargs=task_kwargs.model_dump())
+    result = bcftools_pipe_task.apply_async(kwargs=task_kwargs.model_dump())
 
-    job_id = await create_task_history_entry(
-        user_id=current_user.id,
-        project_id=project.id,
-        task_id=results.id,
+    await update_task_history_entry_with_celery_task_id(
+        job_id=job_id,
+        task_id=result.id,
         db=db,
     )
 

packages/divbase-api/src/divbase_api/services/queries.py

@@ -127,7 +127,7 @@ class BcftoolsQueryManager:
     VALID_BCFTOOLS_COMMANDS = ["view"]  # white-list of valid bcftools commands to run in the pipe.
     CONTAINER_NAME = "divbase-worker-quick-1"  # for synchronous tasks, use this container name to find the container ID
 
-    def execute_pipe(self, command: str, bcftools_inputs: dict, task_id: str = None) -> str:
+    def execute_pipe(self, command: str, bcftools_inputs: dict, job_id: int) -> str:
         """
         Main entrypoint for executing executing divbase queries that require bcftools.
         First calls on a method to build a structure of input parameters for bcftools, and then
@@ -144,7 +144,7 @@ def execute_pipe(self, command: str, bcftools_inputs: dict, task_id: str = None)
             except BcftoolsEnvironmentError:
                 raise
 
-        identifier = task_id if task_id else datetime.datetime.now().timestamp()
+        identifier = job_id if job_id else datetime.datetime.now().timestamp()
 
         commands_config_structure = self.build_commands_config(command, bcftools_inputs, identifier)
         output_file = self.process_bcftools_commands(commands_config_structure, identifier)
@@ -345,7 +345,7 @@ def merge_or_concat_bcftools_temp_files(self, output_temp_files: list[str], iden
         self.temp_files.append(annotated_unsorted_output_file)
         self.temp_files.append(divbase_header_for_vcf)
 
-        output_file = f"merged_{identifier}.vcf.gz"
+        output_file = f"result_of_job_{identifier}.vcf.gz"
         logger.info("Trying to determine if sample names overlap between temp files...")
 
         sample_names_per_VCF = self._get_all_sample_names_from_vcf_files(output_temp_files)

packages/divbase-api/src/divbase_api/worker/tasks.py

@@ -189,12 +189,13 @@ def bcftools_pipe_task(
     project_id: int,
     project_name: str,
     user_id: int,
+    job_id: int,
 ):
     """
     Run a full bcftools query command as a Celery task, with sample metadata filtering run first.
     """
     task_id = bcftools_pipe_task.request.id
-    logger.info(f"Starting bcftools_pipe_task with Celery, task ID: {task_id}")
+    logger.info(f"Starting bcftools_pipe_task with Celery, task ID: {task_id}, job ID: {job_id}")
 
     s3_file_manager = create_s3_file_manager(url=S3_ENDPOINT_URL)
 
@@ -258,7 +259,7 @@ def bcftools_pipe_task(
 
     # Let validation exceptions (BcftoolsPipeEmptyCommandError, BcftoolsPipeUnsupportedCommandError,
     # SidecarInvalidFilterError) propagate to mark task as FAILURE. Otherwise the tasks will incorrectly be marked as SUCCESS.
-    output_file = BcftoolsQueryManager().execute_pipe(command, bcftools_inputs, task_id)
+    output_file = BcftoolsQueryManager().execute_pipe(command, bcftools_inputs, job_id)
 
     _upload_results_file(output_file=Path(output_file), bucket_name=bucket_name, s3_file_manager=s3_file_manager)
     _delete_job_files_from_worker(vcf_paths=files_to_download, metadata_path=metadata_path, output_file=output_file)

packages/divbase-lib/src/divbase_lib/api_schemas/queries.py

@@ -45,6 +45,7 @@ class BcftoolsQueryKwargs(BaseModel):
     project_id: int
     project_name: str
     user_id: int
+    job_id: int
 
 
 class SampleMetadataQueryTaskResult(BaseModel):