Coverage for /pythoncovmergedfiles/medio/medio/src/airflow/airflow/metrics/protocols.py: 88%

1# Licensed to the Apache Software Foundation (ASF) under one 

2# or more contributor license agreements. See the NOTICE file 

3# distributed with this work for additional information 

4# regarding copyright ownership. The ASF licenses this file 

5# to you under the Apache License, Version 2.0 (the 

6# "License"); you may not use this file except in compliance 

7# with the License. You may obtain a copy of the License at 

8# 

9# http://www.apache.org/licenses/LICENSE-2.0 

10# 

11# Unless required by applicable law or agreed to in writing, 

12# software distributed under the License is distributed on an 

13# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY 

14# KIND, either express or implied. See the License for the 

15# specific language governing permissions and limitations 

16# under the License. 

17 

18from __future__ import annotations 

19 

20import datetime 

21import time 

22from typing import Union 

23 

24from airflow.typing_compat import Protocol 

25 

26DeltaType = Union[int, float, datetime.timedelta] 

27 

28 

29class TimerProtocol(Protocol): 

30 """Type protocol for StatsLogger.timer.""" 

31 

32 def __enter__(self) -> Timer: ... 

33 

34 def __exit__(self, exc_type, exc_value, traceback) -> None: ... 

35 

36 def start(self) -> Timer: 

37 """Start the timer.""" 

38 ... 

39 

40 def stop(self, send: bool = True) -> None: 

41 """Stop, and (by default) submit the timer to StatsD.""" 

42 ... 

43 

44 

45class Timer(TimerProtocol): 

46 """ 

47 Timer that records duration, and optional sends to StatsD backend. 

48 

49 This class lets us have an accurate timer with the logic in one place (so 

50 that we don't use datetime math for duration -- it is error prone). 

51 

52 Example usage: 

53 

54 .. code-block:: python 

55 

56 with Stats.timer() as t: 

57 # Something to time 

58 frob_the_foos() 

59 

60 log.info("Frobbing the foos took %.2f", t.duration) 

61 

62 Or without a context manager: 

63 

64 .. code-block:: python 

65 

66 timer = Stats.timer().start() 

67 

68 # Something to time 

69 frob_the_foos() 

70 

71 timer.end() 

72 

73 log.info("Frobbing the foos took %.2f", timer.duration) 

74 

75 To send a metric: 

76 

77 .. code-block:: python 

78 

79 with Stats.timer("foos.frob"): 

80 # Something to time 

81 frob_the_foos() 

82 

83 Or both: 

84 

85 .. code-block:: python 

86 

87 with Stats.timer("foos.frob") as t: 

88 # Something to time 

89 frob_the_foos() 

90 

91 log.info("Frobbing the foos took %.2f", t.duration) 

92 """ 

93 

94 # pystatsd and dogstatsd both have a timer class, but present different API 

95 # so we can't use this as a mixin on those, instead this class contains the "real" timer 

96 

97 _start_time: float | None 

98 duration: float | None 

99 

100 def __init__(self, real_timer: Timer | None = None) -> None: 

101 self.real_timer = real_timer 

102 

103 def __enter__(self) -> Timer: 

104 return self.start() 

105 

106 def __exit__(self, exc_type, exc_value, traceback) -> None: 

107 self.stop() 

108 

109 def start(self) -> Timer: 

110 """Start the timer.""" 

111 if self.real_timer: 

112 self.real_timer.start() 

113 self._start_time = time.perf_counter() 

114 return self 

115 

116 def stop(self, send: bool = True) -> None: 

117 """Stop the timer, and optionally send it to stats backend.""" 

118 if self._start_time is not None: 

119 self.duration = time.perf_counter() - self._start_time 

120 if send and self.real_timer: 

121 self.real_timer.stop()