Coverage for app/services/task_service.py: 70.73%

1"""Task lifecycle helpers. 

2 

3Two-phase dispatch: 

4 

51. `prepare_task_in_tx` — INSERT a PENDING task row in the caller's 

6 transaction (no commit, no Celery I/O). The caller commits the 

7 surrounding business state and the new task row atomically. 

8 

92. `dispatch_to_celery` — outside the original TX, push the task to 

10 RabbitMQ. On success the row's `celeryTaskId` is stamped; on 

11 failure the row is marked FAILED. Either way the user sees a row 

12 in the deployment list reflecting reality — no splitbrain. 

13""" 

14from __future__ import annotations 

15 

16import logging 

17import uuid 

18 

19from sqlalchemy.exc import IntegrityError 

20from sqlalchemy.orm import Session 

21 

22from app.celery_app import celery_app 

23from app.crud import tasks as crud_tasks 

24from app.models import Task, TaskStatus, TaskType 

25 

26logger = logging.getLogger(__name__) 

27 

28 

29# Name of the Postgres partial unique index backing the "one active 

30# task per deployment" rule. String-matched when translating 

31# IntegrityError → ActiveTaskExistsError so other unique-constraint 

32# violations aren't swallowed. 

33_ACTIVE_TASK_UNIQUE_INDEX = "uq_tasks_active_per_deployment" 

34 

35 

36class ActiveTaskExistsError(Exception): 

37 """A PENDING/RUNNING task already exists for this deployment.""" 

38 

39 

40def prepare_task_in_tx( 

41 db: Session, 

42 deployment_id: uuid.UUID, 

43 task_type: TaskType, 

44) -> Task: 

45 """Insert a PENDING task row in the caller's transaction. 

46 

47 Does NOT call `db.commit()` — the caller is responsible for 

48 committing the surrounding state alongside this row, so that 

49 deployment + teams + task are all visible (or all rolled back) 

50 atomically. 

51 

52 Raises `ActiveTaskExistsError` if the deployment already has a 

53 PENDING/RUNNING task. A Postgres partial unique index enforces this 

54 at the DB level too: the pre-check catches the common case, and the 

55 ``except IntegrityError`` around ``db.flush()`` catches the racy one 

56 (two concurrent requests both pass the pre-check). 

57 """ 

58 existing = crud_tasks.get_tasks(db, deployment_id=deployment_id) 

59 for task in existing: 

60 if task.status in (TaskStatus.PENDING, TaskStatus.RUNNING): 

61 raise ActiveTaskExistsError( 

62 f"Deployment {deployment_id} already has active task {task.taskId}" 

63 ) 

64 

65 db_task = Task( 

66 deploymentId=deployment_id, 

67 type=task_type, 

68 status=TaskStatus.PENDING, 

69 celeryTaskId=None, 

70 ) 

71 db.add(db_task) 

72 try: 

73 db.flush() 

74 except IntegrityError as e: 

75 # Race: another transaction inserted a PENDING/RUNNING task for 

76 # the same deployment between our pre-check and flush. Translate 

77 # the constraint violation into the domain exception so the 

78 # caller's 409 branch fires. 

79 if _ACTIVE_TASK_UNIQUE_INDEX in str(e.orig): 

80 raise ActiveTaskExistsError( 

81 f"Deployment {deployment_id} already has an active task " 

82 "(detected via DB unique constraint)" 

83 ) from e 

84 raise 

85 db.refresh(db_task) 

86 return db_task 

87 

88 

89def dispatch_to_celery( 

90 db: Session, 

91 task: Task, 

92 celery_task_name: str, 

93 celery_args: list, 

94) -> tuple[Task, str]: 

95 """Push a prepared task to Celery. 

96 

97 MUST be called after the surrounding TX committed. Runs in fresh 

98 transactions so the task row is updated independently of any later 

99 request-handling commit. 

100 

101 On `send_task` failure the task is marked FAILED and the original 

102 exception is re-raised — the caller turns that into a 503. 

103 """ 

104 try: 

105 result = celery_app.send_task(celery_task_name, args=celery_args) 

106 except Exception: 

107 logger.exception( 

108 "Celery dispatch failed for task %s (deployment %s)", 

109 task.taskId, 

110 task.deploymentId, 

111 ) 

112 task.status = TaskStatus.FAILED 

113 task.logs = "Failed to dispatch to Celery" 

114 db.commit() 

115 db.refresh(task) 

116 raise 

117 

118 task.celeryTaskId = result.id 

119 db.commit() 

120 db.refresh(task) 

121 logger.info("Task %s dispatched as celery task %s", task.taskId, result.id) 

122 return task, result.id