Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/airflow/metrics/protocols.py: 82%

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 

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

36 ... 

37 

38 def start(self) -> Timer: 

39 """Start the timer.""" 

40 ... 

41 

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

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

44 ... 

45 

46 

47class Timer(TimerProtocol): 

48 """ 

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

50 

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

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

53 

54 Example usage: 

55 

56 .. code-block:: python 

57 

58 with Stats.timer() as t: 

59 # Something to time 

60 frob_the_foos() 

61 

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

63 

64 Or without a context manager: 

65 

66 .. code-block:: python 

67 

68 timer = Stats.timer().start() 

69 

70 # Something to time 

71 frob_the_foos() 

72 

73 timer.end() 

74 

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

76 

77 To send a metric: 

78 

79 .. code-block:: python 

80 

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

82 # Something to time 

83 frob_the_foos() 

84 

85 Or both: 

86 

87 .. code-block:: python 

88 

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

90 # Something to time 

91 frob_the_foos() 

92 

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

94 """ 

95 

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

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

98 

99 _start_time: float | None 

100 duration: float | None 

101 

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

103 self.real_timer = real_timer 

104 

105 def __enter__(self) -> Timer: 

106 return self.start() 

107 

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

109 self.stop() 

110 

111 def start(self) -> Timer: 

112 """Start the timer.""" 

113 if self.real_timer: 

114 self.real_timer.start() 

115 self._start_time = time.perf_counter() 

116 return self 

117 

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

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

120 if self._start_time is not None: 

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

122 if send and self.real_timer: 

123 self.real_timer.stop()