Coverage for seedboxsync/core/taskmanager/track_taskstatus.py: 93%

1# 

2# Copyright (C) 2015-2026 Guillaume Kulakowski <guillaume@kulakowski.fr> 

3# 

4# For the full copyright and license information, please view the LICENSE 

5# file that was distributed with this source code. 

6# 

7"""Track task status by decorator.""" 

8 

9from collections.abc import Callable 

10from datetime import datetime 

11from functools import wraps 

12from typing import ParamSpec, TypeVar 

13from seedboxsync.core.dao import TaskStatus 

14 

15P = ParamSpec("P") 

16R = TypeVar("R") 

17 

18 

19def heartbeat_startup() -> None: 

20 """ 

21 Task manager heartbeat. 

22 

23 Update the TaskStatus with key "heartbeat". 

24 """ 

25 started = datetime.now() 

26 TaskStatus.insert( 

27 key="heartbeat", 

28 running=True, 

29 started=started, 

30 finished=started, 

31 ).on_conflict( 

32 conflict_target=[TaskStatus.key], 

33 update={ 

34 TaskStatus.running: True, 

35 TaskStatus.started: started, 

36 TaskStatus.finished: started, 

37 }, 

38 ).execute() 

39 

40 

41def heartbeat_shutdown() -> None: 

42 """ 

43 Task manager heartbeat. 

44 

45 Update the TaskStatus with key "heartbeat". 

46 """ 

47 finished = datetime.now() 

48 TaskStatus.insert( 

49 key="heartbeat", 

50 running=False, 

51 finished=finished, 

52 ).on_conflict( 

53 conflict_target=[TaskStatus.key], 

54 update={ 

55 TaskStatus.running: False, 

56 TaskStatus.finished: finished, 

57 }, 

58 ).execute() 

59 

60 

61def heartbeat() -> None: 

62 """ 

63 Task manager heartbeat. 

64 

65 Update the TaskStatus with key "heartbeat". 

66 """ 

67 finished = datetime.now() 

68 TaskStatus.insert( 

69 key="heartbeat", 

70 finished=finished, 

71 ).on_conflict( 

72 conflict_target=[TaskStatus.key], 

73 update={ 

74 TaskStatus.finished: finished, 

75 }, 

76 ).execute() 

77 

78 

79def track_taskstatus(key: str) -> Callable[[Callable[P, R]], Callable[P, R]]: 

80 """ 

81 Track the execution status of a task in the database. 

82 

83 Args: 

84 key: Unique identifier used to store the task status. 

85 

86 Returns: 

87 A decorator that updates the task status before and after execution. 

88 """ 

89 

90 def decorator(function: Callable[P, R]) -> Callable[P, R]: 

91 @wraps(function) 

92 def wrapper(*args: P.args, **kwargs: P.kwargs) -> R: 

93 started = datetime.now() 

94 

95 TaskStatus.insert( 

96 key=key, 

97 running=True, 

98 started=started, 

99 finished=None, 

100 ).on_conflict( 

101 conflict_target=[TaskStatus.key], 

102 update={ 

103 TaskStatus.running: True, 

104 TaskStatus.started: started, 

105 TaskStatus.finished: None, 

106 }, 

107 ).execute() 

108 # Call heartbeat() 

109 heartbeat() 

110 

111 try: 

112 return function(*args, **kwargs) 

113 finally: 

114 TaskStatus.update( 

115 running=False, 

116 finished=datetime.now(), 

117 ).where( 

118 TaskStatus.key == key, 

119 ).execute() 

120 # Call heartbeat() 

121 heartbeat() 

122 

123 return wrapper 

124 

125 return decorator