Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/airflow/sdk/_shared/observability/metrics/protocols.py: 56%

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 TYPE_CHECKING, Protocol 

23 

24if TYPE_CHECKING: 

25 from .. import Self 

26 

27DeltaType = int | float | datetime.timedelta 

28 

29 

30class TimerProtocol(Protocol): 

31 """Type protocol for StatsLogger.timer.""" 

32 

33 def __enter__(self) -> Self: ... 

34 

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

36 

37 def start(self) -> Self: 

38 """Start the timer.""" 

39 ... 

40 

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

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

43 ... 

44 

45 

46class Timer(TimerProtocol): 

47 """ 

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

49 

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

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

52 

53 Example usage: 

54 

55 .. code-block:: python 

56 

57 with Stats.timer() as t: 

58 # Something to time 

59 frob_the_foos() 

60 

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

62 

63 Or without a context manager: 

64 

65 .. code-block:: python 

66 

67 timer = Stats.timer().start() 

68 

69 # Something to time 

70 frob_the_foos() 

71 

72 timer.end() 

73 

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

75 

76 To send a metric: 

77 

78 .. code-block:: python 

79 

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

81 # Something to time 

82 frob_the_foos() 

83 

84 Or both: 

85 

86 .. code-block:: python 

87 

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

89 # Something to time 

90 frob_the_foos() 

91 

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

93 """ 

94 

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

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

97 

98 _start_time: float | None 

99 duration: float | None 

100 

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

102 self.real_timer = real_timer 

103 

104 def __enter__(self) -> Self: 

105 return self.start() 

106 

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

108 self.stop() 

109 

110 def start(self) -> Self: 

111 """Start the timer.""" 

112 if self.real_timer: 

113 self.real_timer.start() 

114 self._start_time = time.perf_counter() 

115 return self 

116 

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

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

119 if self._start_time is not None: 

120 self.duration = 1000.0 * (time.perf_counter() - self._start_time) # Convert to milliseconds. 

121 if send and self.real_timer: 

122 self.real_timer.stop()