Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/events/models/composite_trigger_child_firing.py: 43%

1from typing import TYPE_CHECKING, Sequence 

2from uuid import UUID 

3 

4import sqlalchemy as sa 

5from sqlalchemy.dialects import postgresql 

6from sqlalchemy.ext.asyncio import AsyncSession 

7 

8from prefect.server.database import PrefectDBInterface, db_injector 

9from prefect.server.events.schemas.automations import CompositeTrigger, Firing 

10from prefect.server.utilities.database import get_dialect 

11from prefect.types._datetime import DateTime, now 

12 

13if TYPE_CHECKING: 13 ↛ 14line 13 didn't jump to line 14 because the condition on line 13 was never true

14 from prefect.server.database.orm_models import ORMCompositeTriggerChildFiring 

15 

16 

17async def acquire_composite_trigger_lock( 

18 session: AsyncSession, 

19 trigger: CompositeTrigger, 

20) -> None: 

21 """ 

22 Acquire a transaction-scoped advisory lock for the given composite trigger. 

23 

24 This serializes concurrent child trigger evaluations for the same compound 

25 trigger, preventing a race condition where multiple transactions each see 

26 only their own child firing and neither fires the parent. 

27 

28 The lock is automatically released when the transaction commits or rolls back. 

29 """ 

30 bind = session.get_bind() 

31 if bind is None: 

32 return 

33 

34 # Get the engine from either an Engine or Connection 

35 engine: sa.Engine = bind if isinstance(bind, sa.Engine) else bind.engine # type: ignore[union-attr] 

36 dialect = get_dialect(engine) 

37 

38 if dialect.name == "postgresql": 

39 # Use the trigger's UUID as the lock key 

40 # pg_advisory_xact_lock takes a bigint, so we use the UUID's int representation 

41 # truncated to fit (collision is extremely unlikely and benign) 

42 lock_key = int(trigger.id) % (2**63) 

43 await session.execute( 

44 sa.text("SELECT pg_advisory_xact_lock(:key)"), {"key": lock_key} 

45 ) 

46 # SQLite doesn't support advisory locks, but SQLite also serializes writes 

47 # at the database level, so the race condition is less likely to occur 

48 

49 

50@db_injector 

51async def upsert_child_firing( 

52 db: PrefectDBInterface, 

53 session: AsyncSession, 

54 firing: Firing, 

55): 

56 automation_id = firing.trigger.automation.id 

57 parent_trigger_id = firing.trigger.parent.id 

58 child_trigger_id = firing.trigger.id 

59 

60 upsert = ( 

61 postgresql.insert(db.CompositeTriggerChildFiring) 

62 .values( 

63 automation_id=automation_id, 

64 parent_trigger_id=parent_trigger_id, 

65 child_trigger_id=child_trigger_id, 

66 child_firing_id=firing.id, 

67 child_fired_at=firing.triggered, 

68 child_firing=firing.model_dump(), 

69 ) 

70 .on_conflict_do_update( 

71 index_elements=[ 

72 db.CompositeTriggerChildFiring.automation_id, 

73 db.CompositeTriggerChildFiring.parent_trigger_id, 

74 db.CompositeTriggerChildFiring.child_trigger_id, 

75 ], 

76 set_=dict( 

77 child_firing_id=firing.id, 

78 child_fired_at=firing.triggered, 

79 child_firing=firing.model_dump(), 

80 updated=now("UTC"), 

81 ), 

82 ) 

83 ) 

84 

85 await session.execute(upsert) 

86 

87 result = await session.execute( 

88 sa.select(db.CompositeTriggerChildFiring).filter( 

89 db.CompositeTriggerChildFiring.automation_id == automation_id, 

90 db.CompositeTriggerChildFiring.parent_trigger_id == parent_trigger_id, 

91 db.CompositeTriggerChildFiring.child_trigger_id == child_trigger_id, 

92 ) 

93 ) 

94 

95 return result.scalars().one() 

96 

97 

98@db_injector 

99async def get_child_firings( 

100 db: PrefectDBInterface, 

101 session: AsyncSession, 

102 trigger: CompositeTrigger, 

103) -> Sequence["ORMCompositeTriggerChildFiring"]: 

104 result = await session.execute( 

105 sa.select(db.CompositeTriggerChildFiring).filter( 

106 db.CompositeTriggerChildFiring.automation_id == trigger.automation.id, 

107 db.CompositeTriggerChildFiring.parent_trigger_id == trigger.id, 

108 db.CompositeTriggerChildFiring.child_trigger_id.in_( 

109 trigger.child_trigger_ids 

110 ), 

111 ) 

112 ) 

113 

114 return result.scalars().unique().all() 

115 

116 

117@db_injector 

118async def clear_old_child_firings( 

119 db: PrefectDBInterface, 

120 session: AsyncSession, 

121 trigger: CompositeTrigger, 

122 fired_before: DateTime, 

123) -> None: 

124 await session.execute( 

125 sa.delete(db.CompositeTriggerChildFiring).filter( 

126 db.CompositeTriggerChildFiring.automation_id == trigger.automation.id, 

127 db.CompositeTriggerChildFiring.parent_trigger_id == trigger.id, 

128 db.CompositeTriggerChildFiring.child_fired_at < fired_before, 

129 ) 

130 ) 

131 

132 

133@db_injector 

134async def clear_child_firings( 

135 db: PrefectDBInterface, 

136 session: AsyncSession, 

137 trigger: CompositeTrigger, 

138 firing_ids: Sequence[UUID], 

139) -> set[UUID]: 

140 """ 

141 Delete the specified child firings and return the IDs that were actually deleted. 

142 

143 Returns the set of child_firing_ids that were successfully deleted. Callers can 

144 compare this to the expected firing_ids to detect races and avoid double-firing 

145 composite triggers. 

146 """ 

147 result = await session.execute( 

148 sa.delete(db.CompositeTriggerChildFiring) 

149 .filter( 

150 db.CompositeTriggerChildFiring.automation_id == trigger.automation.id, 

151 db.CompositeTriggerChildFiring.parent_trigger_id == trigger.id, 

152 db.CompositeTriggerChildFiring.child_firing_id.in_(firing_ids), 

153 ) 

154 .returning(db.CompositeTriggerChildFiring.child_firing_id) 

155 ) 

156 

157 return set(result.scalars().all())